132 lines · plain
1// RUN: mlir-opt %s -pass-pipeline='builtin.module(func.func(test-affine-reify-value-bounds))' -verify-diagnostics \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{use-arith-ops}))' \5// RUN: -verify-diagnostics -split-input-file | \6// RUN: FileCheck %s --check-prefix=CHECK-ARITH7 8// CHECK: #[[$map:.*]] = affine_map<()[s0] -> (s0 + 5)>9// CHECK-LABEL: func @arith_addi(10// CHECK-SAME: %[[a:.*]]: index11// CHECK: %[[apply:.*]] = affine.apply #[[$map]]()[%[[a]]]12// CHECK: return %[[apply]]13 14// CHECK-ARITH-LABEL: func @arith_addi(15// CHECK-ARITH-SAME: %[[a:.*]]: index16// CHECK-ARITH: %[[c5:.*]] = arith.constant 5 : index17// CHECK-ARITH: %[[add:.*]] = arith.addi %[[c5]], %[[a]]18// CHECK-ARITH: %[[c5:.*]] = arith.constant 5 : index19// CHECK-ARITH: %[[add:.*]] = arith.addi %[[a]], %[[c5]]20// CHECK-ARITH: return %[[add]]21func.func @arith_addi(%a: index) -> index {22 %0 = arith.constant 5 : index23 %1 = arith.addi %0, %a : index24 %2 = "test.reify_bound"(%1) : (index) -> (index)25 return %2 : index26}27 28// -----29 30// CHECK: #[[$map:.*]] = affine_map<()[s0] -> (-s0 + 5)>31// CHECK-LABEL: func @arith_subi(32// CHECK-SAME: %[[a:.*]]: index33// CHECK: %[[apply:.*]] = affine.apply #[[$map]]()[%[[a]]]34// CHECK: return %[[apply]]35func.func @arith_subi(%a: index) -> index {36 %0 = arith.constant 5 : index37 %1 = arith.subi %0, %a : index38 %2 = "test.reify_bound"(%1) : (index) -> (index)39 return %2 : index40}41 42// -----43 44// CHECK: #[[$map:.*]] = affine_map<()[s0] -> (s0 * 5)>45// CHECK-LABEL: func @arith_muli(46// CHECK-SAME: %[[a:.*]]: index47// CHECK: %[[apply:.*]] = affine.apply #[[$map]]()[%[[a]]]48// CHECK: return %[[apply]]49func.func @arith_muli(%a: index) -> index {50 %0 = arith.constant 5 : index51 %1 = arith.muli %0, %a : index52 %2 = "test.reify_bound"(%1) : (index) -> (index)53 return %2 : index54}55 56// -----57 58func.func @arith_muli_non_pure(%a: index, %b: index) -> index {59 %0 = arith.muli %a, %b : index60 // Semi-affine expressions (such as "symbol * symbol") are not supported.61 // expected-error @below{{could not reify bound}}62 %1 = "test.reify_bound"(%0) : (index) -> (index)63 return %1 : index64}65 66// -----67 68// CHECK: #[[$map:.*]] = affine_map<()[s0] -> (s0 floordiv 5)>69// CHECK-LABEL: func @arith_floordivsi(70// CHECK-SAME: %[[a:.*]]: index71// CHECK: %[[apply:.*]] = affine.apply #[[$map]]()[%[[a]]]72// CHECK: return %[[apply]]73func.func @arith_floordivsi(%a: index) -> index {74 %0 = arith.constant 5 : index75 %1 = arith.floordivsi %a, %0 : index76 %2 = "test.reify_bound"(%1) : (index) -> (index)77 return %2 : index78}79 80// -----81 82func.func @arith_floordivsi_non_pure(%a: index, %b: index) -> index {83 %0 = arith.floordivsi %a, %b : index84 // Semi-affine expressions (such as "symbol * symbol") are not supported.85 // expected-error @below{{could not reify bound}}86 %1 = "test.reify_bound"(%0) : (index) -> (index)87 return %1 : index88}89 90// -----91 92// CHECK-LABEL: func @arith_const()93// CHECK: %[[c5:.*]] = arith.constant 5 : index94// CHECK: %[[c5:.*]] = arith.constant 5 : index95// CHECK: return %[[c5]]96func.func @arith_const() -> index {97 %c5 = arith.constant 5 : index98 %0 = "test.reify_bound"(%c5) : (index) -> (index)99 return %0 : index100}101 102// -----103 104// CHECK-LABEL: func @arith_select(105func.func @arith_select(%c: i1) -> (index, index) {106 // CHECK: arith.constant 5 : index107 %c5 = arith.constant 5 : index108 // CHECK: arith.constant 9 : index109 %c9 = arith.constant 9 : index110 %r = arith.select %c, %c5, %c9 : index111 // CHECK: %[[c5:.*]] = arith.constant 5 : index112 // CHECK: %[[c10:.*]] = arith.constant 10 : index113 %0 = "test.reify_bound"(%r) {type = "LB"} : (index) -> (index)114 %1 = "test.reify_bound"(%r) {type = "UB"} : (index) -> (index)115 // CHECK: return %[[c5]], %[[c10]]116 return %0, %1 : index, index117}118 119// -----120 121// CHECK-LABEL: func @arith_select_elementwise(122// CHECK-SAME: %[[a:.*]]: tensor<?xf32>, %[[b:.*]]: tensor<?xf32>, %[[c:.*]]: tensor<?xi1>)123func.func @arith_select_elementwise(%a: tensor<?xf32>, %b: tensor<?xf32>, %c: tensor<?xi1>) -> index {124 %r = arith.select %c, %a, %b : tensor<?xi1>, tensor<?xf32>125 // CHECK: %[[c0:.*]] = arith.constant 0 : index126 // CHECK: %[[dim:.*]] = tensor.dim %[[a]], %[[c0]]127 %0 = "test.reify_bound"(%r) {type = "EQ", dim = 0}128 : (tensor<?xf32>) -> (index)129 // CHECK: return %[[dim]]130 return %0 : index131}132