126 lines · plain
1// RUN: mlir-opt %s -pass-pipeline='builtin.module(func.func(test-affine-reify-value-bounds{reify-to-func-args}))' \2// RUN: -verify-diagnostics -split-input-file | FileCheck %s3 4// RUN: mlir-opt %s -pass-pipeline='builtin.module(func.func(test-affine-reify-value-bounds{reify-to-func-args use-arith-ops}))' \5// RUN: -verify-diagnostics -split-input-file | FileCheck %s --check-prefix=CHECK-ARITH6 7// CHECK-LABEL: func @reify_through_chain(8// CHECK-SAME: %[[sz0:.*]]: index, %[[sz2:.*]]: index9// CHECK: %[[c10:.*]] = arith.constant 10 : index10// CHECK: return %[[sz0]], %[[c10]], %[[sz2]]11 12// CHECK-ARITH-LABEL: func @reify_through_chain(13// CHECK-ARITH-SAME: %[[sz0:.*]]: index, %[[sz2:.*]]: index14// CHECK-ARITH: %[[c10:.*]] = arith.constant 10 : index15// CHECK-ARITH: return %[[sz0]], %[[c10]], %[[sz2]]16func.func @reify_through_chain(%sz0: index, %sz2: index) -> (index, index, index) {17 %c2 = arith.constant 2 : index18 %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 : index21 %f = arith.constant 0.0 : f3222 %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, index30}31 32// -----33 34// CHECK-LABEL: func @reify_slice_bound(35// CHECK: %[[c5:.*]] = arith.constant 5 : index36// CHECK: "test.some_use"(%[[c5]])37// CHECK: %[[c5:.*]] = arith.constant 5 : index38// CHECK: "test.some_use"(%[[c5]])39func.func @reify_slice_bound(%t: tensor<?x?xi32>, %idx: index, %ub: index, %f: i32) {40 %c0 = arith.constant 0 : index41 %c4 = arith.constant 4 : index42 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 : i32) 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 return54}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:.*]]: index61// CHECK: %[[bound:.*]] = affine.apply #[[$map]]()[%[[ub]], %[[lb]]]62// CHECK: "test.some_use"(%[[bound]])63func.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 return70}71 72// -----73 74// CHECK-LABEL: func @reify_slice_bound2(75func.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 : index79 %c1 = arith.constant 1 : index80 %c32 = arith.constant 32 : index81 scf.for %iv0 = %lb0 to %ub0 step %step0 {82 // CHECK: %[[c129:.*]] = arith.constant 129 : index83 // 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 : index89 // 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 : index95 // 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 : index101 // 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 : index113 // 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 : index118 // 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 return125}126