1// RUN: mlir-opt %s -pass-pipeline='builtin.module(func.func(test-affine-reify-value-bounds, cse))' -verify-diagnostics \ 2// RUN: -verify-diagnostics -split-input-file | FileCheck %s 3 4#map_dim_i = affine_map<(d0)[s0] -> (-d0 + 32400, s0)> 5#map_dim_j = affine_map<(d0)[s0] -> (-d0 + 16, s0)> 6 7// Here the upper bound for min_i is 4 x vscale, as we know 4 x vscale is 8// always less than 32400. The bound for min_j is 16, as 16 is always less 9// 4 x vscale_max (vscale_max is the UB for vscale). 10 11// CHECK: #[[$SCALABLE_BOUND_MAP_0:.*]] = affine_map<()[s0] -> (s0 * 4)> 12 13// CHECK-LABEL: @fixed_size_loop_nest 14// CHECK-DAG: %[[VSCALE:.*]] = vector.vscale 15// CHECK-DAG: %[[UB_i:.*]] = affine.apply #[[$SCALABLE_BOUND_MAP_0]]()[%[[VSCALE]]] 16// CHECK-DAG: %[[UB_j:.*]] = arith.constant 16 : index 17// CHECK: "test.some_use"(%[[UB_i]], %[[UB_j]]) : (index, index) -> () 18func.func @fixed_size_loop_nest() { 19 %c16 = arith.constant 16 : index 20 %c32400 = arith.constant 32400 : index 21 %c4 = arith.constant 4 : index 22 %c0 = arith.constant 0 : index 23 %vscale = vector.vscale 24 %c4_vscale = arith.muli %vscale, %c4 : index 25 scf.for %i = %c0 to %c32400 step %c4_vscale { 26 %min_i = affine.min #map_dim_i(%i)[%c4_vscale] 27 scf.for %j = %c0 to %c16 step %c4_vscale { 28 %min_j = affine.min #map_dim_j(%j)[%c4_vscale] 29 %bound_i = "test.reify_bound"(%min_i) {type = "UB", vscale_min = 1, vscale_max = 16, scalable} : (index) -> index 30 %bound_j = "test.reify_bound"(%min_j) {type = "UB", vscale_min = 1, vscale_max = 16, scalable} : (index) -> index 31 "test.some_use"(%bound_i, %bound_j) : (index, index) -> () 32 } 33 } 34 return 35} 36 37// ----- 38 39#map_dynamic_dim = affine_map<(d0)[s0, s1] -> (-d0 + s1, s0)> 40 41// Here upper bounds for both min_i and min_j are both (conservatively) 42// 4 x vscale, as we know that is always the largest value they could take. As 43// if `dim < 4 x vscale` then 4 x vscale is an overestimate, and if 44// `dim > 4 x vscale` then the min will be clamped to 4 x vscale. 45 46// CHECK: #[[$SCALABLE_BOUND_MAP_1:.*]] = affine_map<()[s0] -> (s0 * 4)> 47 48// CHECK-LABEL: @dynamic_size_loop_nest 49// CHECK: %[[VSCALE:.*]] = vector.vscale 50// CHECK: %[[UB_ij:.*]] = affine.apply #[[$SCALABLE_BOUND_MAP_1]]()[%[[VSCALE]]] 51// CHECK: "test.some_use"(%[[UB_ij]], %[[UB_ij]]) : (index, index) -> () 52func.func @dynamic_size_loop_nest(%dim0: index, %dim1: index) { 53 %c4 = arith.constant 4 : index 54 %c0 = arith.constant 0 : index 55 %vscale = vector.vscale 56 %c4_vscale = arith.muli %vscale, %c4 : index 57 scf.for %i = %c0 to %dim0 step %c4_vscale { 58 %min_i = affine.min #map_dynamic_dim(%i)[%c4_vscale, %dim0] 59 scf.for %j = %c0 to %dim1 step %c4_vscale { 60 %min_j = affine.min #map_dynamic_dim(%j)[%c4_vscale, %dim1] 61 %bound_i = "test.reify_bound"(%min_i) {type = "UB", vscale_min = 1, vscale_max = 16, scalable} : (index) -> index 62 %bound_j = "test.reify_bound"(%min_j) {type = "UB", vscale_min = 1, vscale_max = 16, scalable} : (index) -> index 63 "test.some_use"(%bound_i, %bound_j) : (index, index) -> () 64 } 65 } 66 return 67} 68 69// ----- 70 71// Here the bound is just a value + a constant. 72 73// CHECK: #[[$SCALABLE_BOUND_MAP_2:.*]] = affine_map<()[s0] -> (s0 + 8)> 74 75// CHECK-LABEL: @add_to_vscale 76// CHECK: %[[VSCALE:.*]] = vector.vscale 77// CHECK: %[[SCALABLE_BOUND:.*]] = affine.apply #[[$SCALABLE_BOUND_MAP_2]]()[%[[VSCALE]]] 78// CHECK: "test.some_use"(%[[SCALABLE_BOUND]]) : (index) -> () 79func.func @add_to_vscale() { 80 %vscale = vector.vscale 81 %c8 = arith.constant 8 : index 82 %vscale_plus_c8 = arith.addi %vscale, %c8 : index 83 %bound = "test.reify_bound"(%vscale_plus_c8) {type = "EQ", vscale_min = 1, vscale_max = 16, scalable} : (index) -> index 84 "test.some_use"(%bound) : (index) -> () 85 return 86} 87 88// ----- 89 90// Here we know vscale is always 2 so we get a constant bound. 91 92// CHECK-LABEL: @vscale_fixed_size 93// CHECK: %[[C2:.*]] = arith.constant 2 : index 94// CHECK: "test.some_use"(%[[C2]]) : (index) -> () 95func.func @vscale_fixed_size() { 96 %vscale = vector.vscale 97 %bound = "test.reify_bound"(%vscale) {type = "EQ", vscale_min = 2, vscale_max = 2, scalable} : (index) -> index 98 "test.some_use"(%bound) : (index) -> () 99 return 100} 101 102// ----- 103 104// Here we don't know the upper bound (%a is underspecified) 105 106func.func @unknown_bound(%a: index) { 107 %vscale = vector.vscale 108 %vscale_plus_a = arith.muli %vscale, %a : index 109 // expected-error @below{{could not reify bound}} 110 %bound = "test.reify_bound"(%vscale_plus_a) {type = "UB", vscale_min = 1, vscale_max = 16, scalable} : (index) -> index 111 "test.some_use"(%bound) : (index) -> () 112 return 113} 114 115// ----- 116 117// Here we have two vscale values (that have not been CSE'd), but they should 118// still be treated as equivalent. 119 120// CHECK: #[[$SCALABLE_BOUND_MAP_3:.*]] = affine_map<()[s0] -> (s0 * 6)> 121 122// CHECK-LABEL: @duplicate_vscale_values 123// CHECK: %[[VSCALE:.*]] = vector.vscale 124// CHECK: %[[SCALABLE_BOUND:.*]] = affine.apply #[[$SCALABLE_BOUND_MAP_3]]()[%[[VSCALE]]] 125// CHECK: "test.some_use"(%[[SCALABLE_BOUND]]) : (index) -> () 126func.func @duplicate_vscale_values() { 127 %c4 = arith.constant 4 : index 128 %vscale_0 = vector.vscale 129 130 %c2 = arith.constant 2 : index 131 %vscale_1 = vector.vscale 132 133 %c4_vscale = arith.muli %vscale_0, %c4 : index 134 %c2_vscale = arith.muli %vscale_1, %c2 : index 135 %add = arith.addi %c2_vscale, %c4_vscale : index 136 137 %bound = "test.reify_bound"(%add) {type = "EQ", vscale_min = 1, vscale_max = 16, scalable} : (index) -> index 138 "test.some_use"(%bound) : (index) -> () 139 return 140} 141 142// ----- 143 144// Test some non-scalable code to ensure that works too: 145 146#map_dim_i = affine_map<(d0)[s0] -> (-d0 + 1024, s0)> 147 148// CHECK-LABEL: @non_scalable_code 149// CHECK: %[[C4:.*]] = arith.constant 4 : index 150// CHECK: "test.some_use"(%[[C4]]) : (index) -> () 151func.func @non_scalable_code() { 152 %c1024 = arith.constant 1024 : index 153 %c4 = arith.constant 4 : index 154 %c0 = arith.constant 0 : index 155 scf.for %i = %c0 to %c1024 step %c4 { 156 %min_i = affine.min #map_dim_i(%i)[%c4] 157 %bound_i = "test.reify_bound"(%min_i) {type = "UB", vscale_min = 1, vscale_max = 16, scalable} : (index) -> index 158 "test.some_use"(%bound_i) : (index) -> () 159 } 160 return 161} 162 163// ----- 164 165#remainder_start_index = affine_map<()[s0] -> (-(1000 mod s0) + 1000)> 166#remaining_iterations = affine_map<(d0) -> (-d0 + 1000)> 167 168// CHECK: #[[$REMAINDER_START_MAP:.*]] = affine_map<()[s0] -> (-(1000 mod s0) + 1000)> 169// CHECK: #[[$SCALABLE_BOUND_MAP_4:.*]] = affine_map<()[s0] -> (s0 * 8 - 1)> 170 171// CHECK-LABEL: @test_scalable_remainder_loop 172// CHECK: %[[VSCALE:.*]] = vector.vscale 173// CHECK: %[[SCALABLE_BOUND:.*]] = affine.apply #[[$SCALABLE_BOUND_MAP_4]]()[%[[VSCALE]]] 174// CHECK: "test.some_use"(%[[SCALABLE_BOUND]]) : (index) -> () 175func.func @test_scalable_remainder_loop() { 176 %c8 = arith.constant 8 : index 177 %c1000 = arith.constant 1000 : index 178 %vscale = vector.vscale 179 %c8_vscale = arith.muli %vscale, %c8 : index 180 %0 = affine.apply #remainder_start_index()[%c8_vscale] 181 scf.for %arg1 = %0 to %c1000 step %c8_vscale { 182 %remaining_iterations = affine.apply #remaining_iterations(%arg1) 183 // The upper bound for the remainder loop iterations should be: %c8_vscale - 1 184 // (expressed as an affine map, affine_map<()[s0] -> (s0 * 8 - 1)>, where s0 is vscale) 185 %bound = "test.reify_bound"(%remaining_iterations) <{scalable, type = "UB", vscale_min = 1 : i64, vscale_max = 16 : i64}> : (index) -> index 186 "test.some_use"(%bound) : (index) -> () 187 } 188 return 189} 190 191// ----- 192 193#unsupported_semi_affine = affine_map<()[s0] -> (s0 * s0)> 194 195func.func @unsupported_semi_affine() { 196 %vscale = vector.vscale 197 %0 = affine.apply #unsupported_semi_affine()[%vscale] 198 // expected-error @below{{could not reify bound}} 199 %bound = "test.reify_bound"(%0) <{scalable, type = "UB", vscale_min = 1 : i64, vscale_max = 16 : i64}> : (index) -> index 200 "test.some_use"(%bound) : (index) -> () 201 return 202} 203 204// ----- 205 206#map_mod = affine_map<()[s0] -> (1000 mod s0)> 207 208func.func @unsupported_negative_mod() { 209 %c_minus_1 = arith.constant -1 : index 210 %vscale = vector.vscale 211 %negative_vscale = arith.muli %vscale, %c_minus_1 : index 212 %0 = affine.apply #map_mod()[%negative_vscale] 213 // expected-error @below{{could not reify bound}} 214 %bound = "test.reify_bound"(%0) <{scalable, type = "UB", vscale_min = 1 : i64, vscale_max = 16 : i64}> : (index) -> index 215 "test.some_use"(%bound) : (index) -> () 216 return 217} 218 219// ----- 220 221// CHECK: #[[$SCALABLE_BOUND_MAP_5:.*]] = affine_map<()[s0] -> (s0 * 4)> 222 223// CHECK-LABEL: @extract_slice_loop 224// CHECK: %[[VSCALE:.*]] = vector.vscale 225// CHECK: %[[SCALABLE_BOUND:.*]] = affine.apply #[[$SCALABLE_BOUND_MAP_5]]()[%[[VSCALE]]] 226// CHECK: "test.some_use"(%[[SCALABLE_BOUND]]) : (index) -> () 227 228func.func @extract_slice_loop(%tensor: tensor<1x1x3x?xf32>) { 229 %vscale = vector.vscale 230 %c0 = arith.constant 0 : index 231 %c1 = arith.constant 1 : index 232 %c2 = arith.constant 2 : index 233 %c3 = arith.constant 3 : index 234 %c4 = arith.constant 4 : index 235 %cst = arith.constant 0.0 : f32 236 %c4_vscale = arith.muli %c4, %vscale : index 237 %slice = tensor.extract_slice %tensor[0, 0, 0, 0] [1, 1, 3, %c4_vscale] [1, 1, 1, 1] : tensor<1x1x3x?xf32> to tensor<1x3x?xf32> 238 %15 = scf.for %arg6 = %c0 to %c3 step %c1 iter_args(%arg = %slice) -> (tensor<1x3x?xf32>) { 239 %dim = tensor.dim %arg, %c2 : tensor<1x3x?xf32> 240 %bound = "test.reify_bound"(%dim) {type = "LB", vscale_min = 1, vscale_max = 16, scalable} : (index) -> index 241 "test.some_use"(%bound) : (index) -> () 242 scf.yield %arg : tensor<1x3x?xf32> 243 } 244 return 245} 246