mlir/test/Dialect/Affine/value-bounds-reification.mlir

   1 // RUN: mlir-opt %s -test-affine-reify-value-bounds="reify-to-func-args" \
   2 // RUN:     -verify-diagnostics -split-input-file | FileCheck %s
   3
   4 // RUN: mlir-opt %s -test-affine-reify-value-bounds="reify-to-func-args use-arith-ops" \
   5 // RUN:     -verify-diagnostics -split-input-file | FileCheck %s --check-prefix=CHECK-ARITH
   6
   7 // CHECK-LABEL: func @reify_through_chain(
   8 //  CHECK-SAME:     %[[sz0:.*]]: index, %[[sz2:.*]]: index
   9 //       CHECK:   %[[c10:.*]] = arith.constant 10 : index
  10 //       CHECK:   return %[[sz0]], %[[c10]], %[[sz2]]
  11
  12 // CHECK-ARITH-LABEL: func @reify_through_chain(
  13 //  CHECK-ARITH-SAME:     %[[sz0:.*]]: index, %[[sz2:.*]]: index
  14 //       CHECK-ARITH:   %[[c10:.*]] = arith.constant 10 : index
  15 //       CHECK-ARITH:   return %[[sz0]], %[[c10]], %[[sz2]]
  16 func.func @reify_through_chain(%sz0: index, %sz2: index) -> (index, index, index) {
  17   %c2 = arith.constant 2 : index
  18   %0 = tensor.empty(%sz0, %sz2) : tensor<?x10x?xf32>
  19   %1 = tensor.cast %0 : tensor<?x10x?xf32> to tensor<?x?x?xf32>
  20   %pos = arith.constant 0 : index
  21   %f = arith.constant 0.0 : f32
  22   %2 = tensor.insert %f into %1[%pos, %pos, %pos] : tensor<?x?x?xf32>
  23   %3 = tensor.dim %2, %c2 : tensor<?x?x?xf32>
  24
  25   %4 = "test.reify_bound"(%2) {dim = 0} : (tensor<?x?x?xf32>) -> (index)
  26   %5 = "test.reify_bound"(%2) {dim = 1} : (tensor<?x?x?xf32>) -> (index)
  27   %6 = "test.reify_bound"(%3) : (index) -> (index)
  28
  29   return %4, %5, %6 : index, index, index
  30 }
  31
  32 // -----
  33
  34 // CHECK-LABEL: func @reify_slice_bound(
  35 //       CHECK:   %[[c5:.*]] = arith.constant 5 : index
  36 //       CHECK:   "test.some_use"(%[[c5]])
  37 //       CHECK:   %[[c5:.*]] = arith.constant 5 : index
  38 //       CHECK:   "test.some_use"(%[[c5]])
  39 func.func @reify_slice_bound(%t: tensor<?x?xi32>, %idx: index, %ub: index, %f: f32) {
  40   %c0 = arith.constant 0 : index
  41   %c4 = arith.constant 4 : index
  42   scf.for %iv = %c0 to %ub step %c4 {
  43     %sz = affine.min affine_map<(d0)[s0] -> (-d0 + s0, 4)>(%iv)[%ub]
  44     %slice = tensor.extract_slice %t[%idx, %iv] [1, %sz] [1, 1] : tensor<?x?xi32> to tensor<1x?xi32>
  45     %filled = linalg.fill ins(%f : f32) outs(%slice : tensor<1x?xi32>) -> tensor<1x?xi32>
  46
  47     %bound = "test.reify_bound"(%filled) {dim = 1, type = "UB"} : (tensor<1x?xi32>) -> (index)
  48     "test.some_use"(%bound) : (index) -> ()
  49
  50     %bound_const = "test.reify_bound"(%filled) {dim = 1, type = "UB", constant} : (tensor<1x?xi32>) -> (index)
  51     "test.some_use"(%bound_const) : (index) -> ()
  52   }
  53   return
  54 }
  55
  56 // -----
  57
  58 // CHECK: #[[$map:.*]] = affine_map<()[s0, s1] -> (s0 - s1 + 1)>
  59 // CHECK-LABEL: func @scf_for(
  60 //  CHECK-SAME:     %[[lb:.*]]: index, %[[ub:.*]]: index, %[[step:.*]]: index
  61 //       CHECK:   %[[bound:.*]] = affine.apply #[[$map]]()[%[[ub]], %[[lb]]]
  62 //       CHECK:   "test.some_use"(%[[bound]])
  63 func.func @scf_for(%lb: index, %ub: index, %step: index) {
  64   scf.for %iv = %lb to %ub step %step {
  65     %0 = affine.apply affine_map<(d0)[s0] -> (-d0 + s0)>(%iv)[%ub]
  66     %bound = "test.reify_bound"(%0) {type = "UB"} : (index) -> (index)
  67     "test.some_use"(%bound) : (index) -> ()
  68   }
  69   return
  70 }
  71
  72 // -----
  73
  74 // CHECK-LABEL: func @reify_slice_bound2(
  75 func.func @reify_slice_bound2(%lb0: index, %ub0: index, %step0: index,
  76                               %ub2: index, %t1: tensor<1x?xi8>,
  77                               %t2: tensor<?x?xi8>, %t3: tensor<1x?xi32>) {
  78   %c0 = arith.constant 0 : index
  79   %c1 = arith.constant 1 : index
  80   %c32 = arith.constant 32 : index
  81   scf.for %iv0 = %lb0 to %ub0 step %step0 {
  82     // CHECK: %[[c129:.*]] = arith.constant 129 : index
  83     // CHECK: "test.some_use"(%[[c129]])
  84     %ub1 = affine.min affine_map<(d0)[s0] -> (-d0 + s0, 128)>(%iv0)[%ub0]
  85     %ub1_ub = "test.reify_bound"(%ub1) {type = "UB"} : (index) -> (index)
  86     "test.some_use"(%ub1_ub) : (index) -> ()
  87
  88     // CHECK: %[[c129:.*]] = arith.constant 129 : index
  89     // CHECK: "test.some_use"(%[[c129]])
  90     %lb1 = affine.apply affine_map<()[s0] -> ((s0 floordiv 32) * 32)>()[%ub1]
  91     %lb1_ub = "test.reify_bound"(%lb1) {type = "UB"} : (index) -> (index)
  92     "test.some_use"(%lb1_ub) : (index) -> ()
  93
  94     // CHECK: %[[c129:.*]] = arith.constant 129 : index
  95     // CHECK: "test.some_use"(%[[c129]])
  96     %lb1_ub_const = "test.reify_bound"(%lb1) {type = "UB", constant} : (index) -> (index)
  97     "test.some_use"(%lb1_ub_const) : (index) -> ()
  98
  99     scf.for %iv1 = %lb1 to %ub1 step %c32 {
 100       // CHECK: %[[c32:.*]] = arith.constant 32 : index
 101       // CHECK: "test.some_use"(%[[c32]])
 102       %sz = affine.apply affine_map<(d0)[s0] -> (-d0 + s0)>(%iv1)[%ub1]
 103       %sz_ub = "test.reify_bound"(%sz) {type = "UB"} : (index) -> (index)
 104       "test.some_use"(%sz_ub) : (index) -> ()
 105
 106       scf.for %iv2 = %c0 to %ub2 step %c1 {
 107         %slice1 = tensor.extract_slice %t1[0, %iv2] [1, 1] [1, 1] : tensor<1x?xi8> to tensor<1x1xi8>
 108         %slice2 = tensor.extract_slice %t2[%iv2, 0] [1, %sz] [1, 1] : tensor<?x?xi8> to tensor<1x?xi8>
 109         %slice3 = tensor.extract_slice %t3[0, 0] [1, %sz] [1, 1] : tensor<1x?xi32> to tensor<1x?xi32>
 110         %matmul = linalg.matmul ins(%slice1, %slice2 : tensor<1x1xi8>, tensor<1x?xi8>) outs(%slice3 : tensor<1x?xi32>) -> tensor<1x?xi32>
 111
 112         // CHECK: %[[c32:.*]] = arith.constant 32 : index
 113         // CHECK: "test.some_use"(%[[c32]])
 114         %matmul_ub = "test.reify_bound"(%matmul) {dim = 1, type = "UB"} : (tensor<1x?xi32>) -> (index)
 115         "test.some_use"(%matmul_ub) : (index) -> ()
 116
 117         // CHECK: %[[c32:.*]] = arith.constant 32 : index
 118         // CHECK: "test.some_use"(%[[c32]])
 119         %matmul_ub_const = "test.reify_bound"(%matmul) {dim = 1, type = "UB", constant} : (tensor<1x?xi32>) -> (index)
 120         "test.some_use"(%matmul_ub_const) : (index) -> ()
 121       }
 122     }
 123   }
 124   return
 125 }