// RUN: mlir-opt %s -for-loop-canonicalization -split-input-file | FileCheck %s

// CHECK-LABEL: func @scf_for_canonicalize_min
//       CHECK:   %[[C2:.*]] = arith.constant 2 : i64
//       CHECK:   scf.for
//       CHECK:     memref.store %[[C2]], %{{.*}}[] : memref<i64>
func @scf_for_canonicalize_min(%A : memref<i64>) {
  %c0 = arith.constant 0 : index
  %c2 = arith.constant 2 : index
  %c4 = arith.constant 4 : index

  scf.for %i = %c0 to %c4 step %c2 {
    %1 = affine.min affine_map<(d0, d1)[] -> (2, d1 - d0)> (%i, %c4)
    %2 = arith.index_cast %1: index to i64
    memref.store %2, %A[]: memref<i64>
  }
  return
}

// -----

// CHECK-LABEL: func @scf_for_canonicalize_max
//       CHECK:   %[[Cneg2:.*]] = arith.constant -2 : i64
//       CHECK:   scf.for
//       CHECK:     memref.store %[[Cneg2]], %{{.*}}[] : memref<i64>
func @scf_for_canonicalize_max(%A : memref<i64>) {
  %c0 = arith.constant 0 : index
  %c2 = arith.constant 2 : index
  %c4 = arith.constant 4 : index

  scf.for %i = %c0 to %c4 step %c2 {
    %1 = affine.max affine_map<(d0, d1)[] -> (-2, -(d1 - d0))> (%i, %c4)
    %2 = arith.index_cast %1: index to i64
    memref.store %2, %A[]: memref<i64>
  }
  return
}

// -----

// CHECK-LABEL: func @scf_for_max_not_canonicalizable
//       CHECK:   scf.for
//       CHECK:     affine.max
//       CHECK:     arith.index_cast
func @scf_for_max_not_canonicalizable(%A : memref<i64>) {
  %c0 = arith.constant 0 : index
  %c2 = arith.constant 2 : index
  %c3 = arith.constant 3 : index
  %c4 = arith.constant 4 : index

  scf.for %i = %c0 to %c4 step %c2 {
    %1 = affine.max affine_map<(d0, d1)[] -> (-2, -(d1 - d0))> (%i, %c3)
    %2 = arith.index_cast %1: index to i64
    memref.store %2, %A[]: memref<i64>
  }
  return
}

// -----

// CHECK-LABEL: func @scf_for_loop_nest_canonicalize_min
//       CHECK:   %[[C5:.*]] = arith.constant 5 : i64
//       CHECK:   scf.for
//       CHECK:     scf.for
//       CHECK:       memref.store %[[C5]], %{{.*}}[] : memref<i64>
func @scf_for_loop_nest_canonicalize_min(%A : memref<i64>) {
  %c0 = arith.constant 0 : index
  %c2 = arith.constant 2 : index
  %c3 = arith.constant 3 : index
  %c4 = arith.constant 4 : index
  %c6 = arith.constant 6 : index

  scf.for %i = %c0 to %c4 step %c2 {
    scf.for %j = %c0 to %c6 step %c3 {
      %1 = affine.min affine_map<(d0, d1, d2, d3)[] -> (5, d1 + d3 - d0 - d2)> (%i, %c4, %j, %c6)
      %2 = arith.index_cast %1: index to i64
      memref.store %2, %A[]: memref<i64>
    }
  }
  return
}

// -----

// CHECK-LABEL: func @scf_for_not_canonicalizable_1
//       CHECK:   scf.for
//       CHECK:     affine.min
//       CHECK:     arith.index_cast
func @scf_for_not_canonicalizable_1(%A : memref<i64>) {
  // This should not canonicalize because: 4 - %i may take the value 1 < 2.
  %c1 = arith.constant 1 : index
  %c2 = arith.constant 2 : index
  %c4 = arith.constant 4 : index

  scf.for %i = %c1 to %c4 step %c2 {
    %1 = affine.min affine_map<(d0)[s0] -> (2, s0 - d0)> (%i)[%c4]
    %2 = arith.index_cast %1: index to i64
    memref.store %2, %A[]: memref<i64>
  }
  return
}

// -----

// CHECK-LABEL: func @scf_for_canonicalize_partly
//       CHECK:   scf.for
//       CHECK:     affine.apply
//       CHECK:     arith.index_cast
func @scf_for_canonicalize_partly(%A : memref<i64>) {
  // This should canonicalize only partly: 256 - %i <= 256.
  %c1 = arith.constant 1 : index
  %c16 = arith.constant 16 : index
  %c256 = arith.constant 256 : index

  scf.for %i = %c1 to %c256 step %c16 {
    %1 = affine.min affine_map<(d0) -> (256, 256 - d0)> (%i)
    %2 = arith.index_cast %1: index to i64
    memref.store %2, %A[]: memref<i64>
  }
  return
}

// -----

// CHECK-LABEL: func @scf_for_not_canonicalizable_2
//       CHECK: scf.for
//       CHECK:   affine.min
//       CHECK:   arith.index_cast
func @scf_for_not_canonicalizable_2(%A : memref<i64>, %step : index) {
  // This example should simplify but affine_map is currently missing
  // semi-affine canonicalizations: `((s0 * 42 - 1) floordiv s0) * s0`
  // should evaluate to 41 * s0.
  // Note that this may require positivity assumptions on `s0`.
  // Revisit when support is added.
  %c0 = arith.constant 0 : index

  %ub = affine.apply affine_map<(d0) -> (42 * d0)> (%step)
  scf.for %i = %c0 to %ub step %step {
    %1 = affine.min affine_map<(d0, d1, d2) -> (d0, d1 - d2)> (%step, %ub, %i)
    %2 = arith.index_cast %1: index to i64
    memref.store %2, %A[]: memref<i64>
  }
  return
}

// -----

// CHECK-LABEL: func @scf_for_not_canonicalizable_3
//       CHECK: scf.for
//       CHECK:   affine.min
//       CHECK:   arith.index_cast
func @scf_for_not_canonicalizable_3(%A : memref<i64>, %step : index) {
  // This example should simplify but affine_map is currently missing
  // semi-affine canonicalizations: `-(((s0 * s0 - 1) floordiv s0) * s0)`
  // should evaluate to (s0 - 1) * s0.
  // Note that this may require positivity assumptions on `s0`.
  // Revisit when support is added.
  %c0 = arith.constant 0 : index

  %ub2 = affine.apply affine_map<(d0)[s0] -> (s0 * d0)> (%step)[%step]
  scf.for %i = %c0 to %ub2 step %step {
    %1 = affine.min affine_map<(d0, d1, d2) -> (d0, d2 - d1)> (%step, %i, %ub2)
    %2 = arith.index_cast %1: index to i64
    memref.store %2, %A[]: memref<i64>
  }
  return
}

// -----

// CHECK-LABEL: func @scf_for_invalid_loop
//       CHECK: scf.for
//       CHECK:   affine.min
//       CHECK:   arith.index_cast
func @scf_for_invalid_loop(%A : memref<i64>, %step : index) {
  // This is an invalid loop. It should not be touched by the canonicalization
  // pattern.
  %c1 = arith.constant 1 : index
  %c7 = arith.constant 7 : index
  %c256 = arith.constant 256 : index

  scf.for %i = %c256 to %c1 step %c1 {
    %1 = affine.min affine_map<(d0)[s0] -> (s0 + d0, 0)> (%i)[%c7]
    %2 = arith.index_cast %1: index to i64
    memref.store %2, %A[]: memref<i64>
  }
  return
}

// -----

// CHECK-LABEL: func @scf_parallel_canonicalize_min_1
//       CHECK:   %[[C2:.*]] = arith.constant 2 : i64
//       CHECK:   scf.parallel
//  CHECK-NEXT:     memref.store %[[C2]], %{{.*}}[] : memref<i64>
func @scf_parallel_canonicalize_min_1(%A : memref<i64>) {
  %c0 = arith.constant 0 : index
  %c2 = arith.constant 2 : index
  %c4 = arith.constant 4 : index

  scf.parallel (%i) = (%c0) to (%c4) step (%c2) {
    %1 = affine.min affine_map<(d0, d1)[] -> (2, d1 - d0)> (%i, %c4)
    %2 = arith.index_cast %1: index to i64
    memref.store %2, %A[]: memref<i64>
  }
  return
}

// -----

// CHECK-LABEL: func @scf_parallel_canonicalize_min_2
//       CHECK:   %[[C2:.*]] = arith.constant 2 : i64
//       CHECK:   scf.parallel
//  CHECK-NEXT:     memref.store %[[C2]], %{{.*}}[] : memref<i64>
func @scf_parallel_canonicalize_min_2(%A : memref<i64>) {
  %c1 = arith.constant 1 : index
  %c2 = arith.constant 2 : index
  %c7 = arith.constant 7 : index

  scf.parallel (%i) = (%c1) to (%c7) step (%c2) {
    %1 = affine.min affine_map<(d0)[s0] -> (2, s0 - d0)> (%i)[%c7]
    %2 = arith.index_cast %1: index to i64
    memref.store %2, %A[]: memref<i64>
  }
  return
}

// -----

// CHECK-LABEL: func @tensor_dim_of_iter_arg(
//  CHECK-SAME:     %[[t:.*]]: tensor<?x?xf32>
//       CHECK:   scf.for
//       CHECK:     tensor.dim %[[t]]
func @tensor_dim_of_iter_arg(%t : tensor<?x?xf32>) -> index {
  %c0 = arith.constant 0 : index
  %c1 = arith.constant 1 : index
  %c10 = arith.constant 10 : index
  %0, %1 = scf.for %i = %c0 to %c10 step %c1 iter_args(%arg0 = %t, %arg1 = %c0)
      -> (tensor<?x?xf32>, index) {
    %dim = tensor.dim %arg0, %c0 : tensor<?x?xf32>
    scf.yield %arg0, %dim : tensor<?x?xf32>, index
  }
  return %1 : index
}

// -----

// CHECK-LABEL: func @tensor_dim_of_iter_arg_insertslice(
//  CHECK-SAME:     %[[t:.*]]: tensor<?x?xf32>,
//       CHECK:   scf.for
//       CHECK:     tensor.dim %[[t]]
func @tensor_dim_of_iter_arg_insertslice(%t : tensor<?x?xf32>,
                                         %t2 : tensor<10x10xf32>) -> index {
  %c0 = arith.constant 0 : index
  %c1 = arith.constant 1 : index
  %c10 = arith.constant 10 : index
  %0, %1 = scf.for %i = %c0 to %c10 step %c1 iter_args(%arg0 = %t, %arg1 = %c0)
      -> (tensor<?x?xf32>, index) {
    %dim = tensor.dim %arg0, %c0 : tensor<?x?xf32>
    %2 = tensor.insert_slice %t2 into %arg0[0, 0] [10, 10] [1, 1]
        : tensor<10x10xf32> into tensor<?x?xf32>
    %3 = tensor.insert_slice %t2 into %2[1, 1] [10, 10] [1, 1]
        : tensor<10x10xf32> into tensor<?x?xf32>
    scf.yield %3, %dim : tensor<?x?xf32>, index
  }
  return %1 : index
}

// -----

// CHECK-LABEL: func @tensor_dim_of_iter_arg_nested_for(
//  CHECK-SAME:     %[[t:.*]]: tensor<?x?xf32>,
//       CHECK:   scf.for
//       CHECK:     scf.for
//       CHECK:       tensor.dim %[[t]]
func @tensor_dim_of_iter_arg_nested_for(%t : tensor<?x?xf32>,
                                        %t2 : tensor<10x10xf32>) -> index {
  %c0 = arith.constant 0 : index
  %c1 = arith.constant 1 : index
  %c10 = arith.constant 10 : index
  %0, %1 = scf.for %i = %c0 to %c10 step %c1 iter_args(%arg0 = %t, %arg1 = %c0)
      -> (tensor<?x?xf32>, index) {
    %2, %3 = scf.for %j = %c0 to %c10 step %c1 iter_args(%arg2 = %arg0, %arg3 = %arg1)
        -> (tensor<?x?xf32>, index) {
      %dim = tensor.dim %arg2, %c0 : tensor<?x?xf32>
      %4 = tensor.insert_slice %t2 into %arg2[0, 0] [10, 10] [1, 1]
          : tensor<10x10xf32> into tensor<?x?xf32>
      scf.yield %4, %dim : tensor<?x?xf32>, index
    }
    scf.yield %2, %3 : tensor<?x?xf32>, index
  }
  return %1 : index
}


// -----

// A test case that should not canonicalize because the loop is not shape
// conserving.

// CHECK-LABEL: func @tensor_dim_of_iter_arg_no_canonicalize(
//  CHECK-SAME:     %[[t:.*]]: tensor<?x?xf32>,
//       CHECK:   scf.for {{.*}} iter_args(%[[arg0:.*]] = %[[t]]
//       CHECK:     tensor.dim %[[arg0]]
func @tensor_dim_of_iter_arg_no_canonicalize(%t : tensor<?x?xf32>,
                                             %t2 : tensor<?x?xf32>) -> index {
  %c0 = arith.constant 0 : index
  %c1 = arith.constant 1 : index
  %c10 = arith.constant 10 : index
  %0, %1 = scf.for %i = %c0 to %c10 step %c1 iter_args(%arg0 = %t, %arg1 = %c0)
      -> (tensor<?x?xf32>, index) {
    %dim = tensor.dim %arg0, %c0 : tensor<?x?xf32>
    scf.yield %t2, %dim : tensor<?x?xf32>, index
  }
  return %1 : index
}

// -----

// CHECK-LABEL: func @tensor_dim_of_loop_result(
//  CHECK-SAME:     %[[t:.*]]: tensor<?x?xf32>
//       CHECK:   tensor.dim %[[t]]
func @tensor_dim_of_loop_result(%t : tensor<?x?xf32>) -> index {
  %c0 = arith.constant 0 : index
  %c1 = arith.constant 1 : index
  %c10 = arith.constant 10 : index
  %0 = scf.for %i = %c0 to %c10 step %c1 iter_args(%arg0 = %t)
      -> (tensor<?x?xf32>) {
    scf.yield %arg0 : tensor<?x?xf32>
  }
  %dim = tensor.dim %0, %c0 : tensor<?x?xf32>
  return %dim : index
}

// -----

// CHECK-LABEL: func @tensor_dim_of_loop_result_no_canonicalize(
//       CHECK:   %[[loop:.*]]:2 = scf.for
//       CHECK:   tensor.dim %[[loop]]#1
func @tensor_dim_of_loop_result_no_canonicalize(%t : tensor<?x?xf32>,
                                                %u : tensor<?x?xf32>) -> index {
  %c0 = arith.constant 0 : index
  %c1 = arith.constant 1 : index
  %c10 = arith.constant 10 : index
  %0, %1 = scf.for %i = %c0 to %c10 step %c1 iter_args(%arg0 = %t, %arg1 = %u)
      -> (tensor<?x?xf32>, tensor<?x?xf32>) {
    scf.yield %arg0, %u : tensor<?x?xf32>, tensor<?x?xf32>
  }
  %dim = tensor.dim %1, %c0 : tensor<?x?xf32>
  return %dim : index
}