Skip to content

Commit 6479a5a

Browse files
banach-spacebanach-space
authored
[mlir][vector] Restrict DropInnerMostUnitDimsTransfer{Read|Write} (#96218)
Restrict `DropInnerMostUnitDimsTransfer{Read|Write}` so that it fails when one of the indices to be dropped could be != 0 and "out of bounds": ```mlir func.func @negative_example(%arg0: memref<16x1xf32>, %arg1: vector<8x1xf32>, %idx_1: index, %idx_2: index) { vector.transfer_write %arg1, %arg0[%idx_1, %idx_2] {in_bounds = [true, false]} : vector<8x1xf32>, memref<16x1xf32> return } ``` This is an edge case that could represent an out-of-bounds access, though that will depend on the actual value of %i. Importantly, without this change it would be transformed as follows: ```mlir func.func @negative_example(%arg0: memref<16x1xf32>, %arg1: vector<8x1xf32>, %arg2: index, %arg3: index) { %subview = memref.subview %arg0[0, 0] [16, 1] [1, 1] : memref<16x1xf32> to memref<16xf32, strided<[1]>> %0 = vector.shape_cast %arg1 : vector<8x1xf32> to vector<8xf32> vector.transfer_write %0, %subview[%arg2] {in_bounds = [true]} : vector<8xf32>, memref<16xf32, strided<[1]>> return } ``` This is incorrect - `%idx_2` is ignored and the "out of bounds" flags is not propagated. Hence the extra restriction to avoid such cases. NOTE: This is a follow-up for: #94904
1 parent fe97671 commit 6479a5a

File tree

2 files changed

+124
-35
lines changed

2 files changed

+124
-35
lines changed

‎mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp

Lines changed: 8 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1300,9 +1300,9 @@ class DropInnerMostUnitDimsTransferRead
13001300
if (dimsToDrop == 0)
13011301
return failure();
13021302

1303-
// Make sure that the indices to be dropped are equal 0.
1304-
// TODO: Deal with cases when the indices are not 0.
1305-
if (!llvm::all_of(readOp.getIndices().take_back(dimsToDrop), isZeroIndex))
1303+
auto inBounds = readOp.getInBoundsValues();
1304+
auto droppedInBounds = ArrayRef<bool>(inBounds).take_back(dimsToDrop);
1305+
if (llvm::is_contained(droppedInBounds, false))
13061306
return failure();
13071307

13081308
auto resultTargetVecType =
@@ -1394,6 +1394,11 @@ class DropInnerMostUnitDimsTransferWrite
13941394
if (dimsToDrop == 0)
13951395
return failure();
13961396

1397+
auto inBounds = writeOp.getInBoundsValues();
1398+
auto droppedInBounds = ArrayRef<bool>(inBounds).take_back(dimsToDrop);
1399+
if (llvm::is_contained(droppedInBounds, false))
1400+
return failure();
1401+
13971402
auto resultTargetVecType =
13981403
VectorType::get(targetType.getShape().drop_back(dimsToDrop),
13991404
targetType.getElementType(),

‎mlir/test/Dialect/Vector/vector-transfer-collapse-inner-most-dims.mlir

Lines changed: 116 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -113,46 +113,69 @@ func.func @contiguous_inner_most_outer_dim_dyn_scalable_inner_dim(%a: index, %b:
113113

114114
// -----
115115

116-
func.func @contiguous_inner_most_dim_non_zero_idx(%A: memref<16x1xf32>, %i:index) -> (vector<8x1xf32>) {
116+
// Test the impact of changing the in_bounds attribute. The behaviour will
117+
// depend on whether the index is == 0 or != 0.
118+
119+
// The index to be dropped is == 0, so it's safe to collapse. The other index
120+
// should be preserved correctly.
121+
func.func @contiguous_inner_most_zero_idx_in_bounds(%A: memref<16x1xf32>, %i:index) -> (vector<8x1xf32>) {
122+
%pad = arith.constant 0.0 : f32
117123
%c0 = arith.constant 0 : index
118-
%f0 = arith.constant 0.0 : f32
119-
%1 = vector.transfer_read %A[%i, %c0], %f0 : memref<16x1xf32>, vector<8x1xf32>
124+
%1 = vector.transfer_read %A[%i, %c0], %pad {in_bounds = [true, true]} : memref<16x1xf32>, vector<8x1xf32>
120125
return %1 : vector<8x1xf32>
121126
}
122-
// CHECK: func @contiguous_inner_most_dim_non_zero_idx(%[[SRC:.+]]: memref<16x1xf32>, %[[I:.+]]: index) -> vector<8x1xf32>
123-
// CHECK: %[[SRC_0:.+]] = memref.subview %[[SRC]]
124-
// CHECK-SAME: memref<16x1xf32> to memref<16xf32, strided<[1]>>
125-
// CHECK: %[[V:.+]] = vector.transfer_read %[[SRC_0]]
126-
// CHECK: %[[RESULT:.+]] = vector.shape_cast %[[V]] : vector<8xf32> to vector<8x1xf32>
127-
// CHECK: return %[[RESULT]]
128-
129-
// The index to be dropped is != 0 - this is currently not supported.
130-
func.func @negative_contiguous_inner_most_dim_non_zero_idxs(%A: memref<16x1xf32>, %i:index) -> (vector<8x1xf32>) {
131-
%f0 = arith.constant 0.0 : f32
132-
%1 = vector.transfer_read %A[%i, %i], %f0 : memref<16x1xf32>, vector<8x1xf32>
127+
// CHECK-LABEL: func.func @contiguous_inner_most_zero_idx_in_bounds(
128+
// CHECK-SAME: %[[MEM:.*]]: memref<16x1xf32>,
129+
// CHECK-SAME: %[[IDX:.*]]: index) -> vector<8x1xf32> {
130+
// CHECK: %[[PAD:.*]] = arith.constant 0.000000e+00 : f32
131+
// CHECK: %[[SV:.*]] = memref.subview %[[MEM]][0, 0] [16, 1] [1, 1] : memref<16x1xf32> to memref<16xf32, strided<[1]>>
132+
// CHECK: %[[READ:.*]] = vector.transfer_read %[[SV]]{{\[}}%[[IDX]]], %[[PAD]] {in_bounds = [true]} : memref<16xf32, strided<[1]>>, vector<8xf32>
133+
// CHECK: vector.shape_cast %[[READ]] : vector<8xf32> to vector<8x1xf32>
134+
135+
// The index to be dropped is == 0, so it's safe to collapse. The "out of
136+
// bounds" attribute is too conservative and will be folded to "in bounds"
137+
// before the pattern runs. The other index should be preserved correctly.
138+
func.func @contiguous_inner_most_zero_idx_out_of_bounds(%A: memref<16x1xf32>, %i:index) -> (vector<8x1xf32>) {
139+
%pad = arith.constant 0.0 : f32
140+
%c0 = arith.constant 0 : index
141+
%1 = vector.transfer_read %A[%i, %c0], %pad {in_bounds = [true, false]} : memref<16x1xf32>, vector<8x1xf32>
133142
return %1 : vector<8x1xf32>
134143
}
135-
// CHECK-LABEL: func @negative_contiguous_inner_most_dim_non_zero_idxs
144+
// CHECK-LABEL: func.func @contiguous_inner_most_zero_idx_out_of_bounds(
145+
// CHECK-SAME: %[[MEM:.*]]: memref<16x1xf32>,
146+
// CHECK-SAME: %[[IDX:.*]]: index) -> vector<8x1xf32> {
147+
// CHECK: %[[PAD:.*]] = arith.constant 0.000000e+00 : f32
148+
// CHECK: %[[SV:.*]] = memref.subview %[[MEM]][0, 0] [16, 1] [1, 1] : memref<16x1xf32> to memref<16xf32, strided<[1]>>
149+
// CHECK: %[[READ:.*]] = vector.transfer_read %[[SV]]{{\[}}%[[IDX]]], %[[PAD]] {in_bounds = [true]} : memref<16xf32, strided<[1]>>, vector<8xf32>
150+
// CHECK: vector.shape_cast %[[READ]] : vector<8xf32> to vector<8x1xf32>
151+
152+
// The index to be dropped is unknown, but since it's "in bounds", it has to be
153+
// == 0. It's safe to collapse the corresponding dim.
154+
func.func @contiguous_inner_most_non_zero_idx_in_bounds(%A: memref<16x1xf32>, %i:index) -> (vector<8x1xf32>) {
155+
%pad = arith.constant 0.0 : f32
156+
%1 = vector.transfer_read %A[%i, %i], %pad {in_bounds = [true, true]} : memref<16x1xf32>, vector<8x1xf32>
157+
return %1 : vector<8x1xf32>
158+
}
159+
// CHECK-LABEL: func.func @contiguous_inner_most_non_zero_idx_in_bounds(
160+
// CHECK-SAME: %[[MEM:.*]]: memref<16x1xf32>,
161+
// CHECK-SAME: %[[IDX:.*]]: index) -> vector<8x1xf32> {
162+
// CHECK: %[[PAD:.*]] = arith.constant 0.000000e+00 : f32
163+
// CHECK: %[[SV:.*]] = memref.subview %[[MEM]][0, 0] [16, 1] [1, 1] : memref<16x1xf32> to memref<16xf32, strided<[1]>>
164+
// CHECK: %[[READ:.*]] = vector.transfer_read %[[SV]]{{\[}}%[[IDX]]], %[[PAD]] {in_bounds = [true]} : memref<16xf32, strided<[1]>>, vector<8xf32>
165+
// CHECK: vector.shape_cast %[[READ]] : vector<8xf32> to vector<8x1xf32>
166+
167+
// The index to be dropped is unknown and "out of bounds" - not safe to
168+
// collapse.
169+
func.func @negative_contiguous_inner_most_non_zero_idx_out_of_bounds(%A: memref<16x1xf32>, %i:index) -> (vector<8x1xf32>) {
170+
%pad = arith.constant 0.0 : f32
171+
%1 = vector.transfer_read %A[%i, %i], %pad {in_bounds = [true, false]} : memref<16x1xf32>, vector<8x1xf32>
172+
return %1 : vector<8x1xf32>
173+
}
174+
// CHECK-LABEL: func.func @negative_contiguous_inner_most_non_zero_idx_out_of_bounds(
136175
// CHECK-NOT: memref.subview
176+
// CHECK-NOT: memref.shape_cast
137177
// CHECK: vector.transfer_read
138178

139-
// Same as the top example within this split, but with the outer vector
140-
// dim scalable. Note that this example only makes sense when "8 = [8]" (i.e.
141-
// vscale = 1). This is assumed (implicitly) via the `in_bounds` attribute.
142-
143-
func.func @contiguous_inner_most_dim_non_zero_idx_scalable_inner_dim(%A: memref<16x1xf32>, %i:index) -> (vector<[8]x1xf32>) {
144-
%c0 = arith.constant 0 : index
145-
%f0 = arith.constant 0.0 : f32
146-
%1 = vector.transfer_read %A[%i, %c0], %f0 : memref<16x1xf32>, vector<[8]x1xf32>
147-
return %1 : vector<[8]x1xf32>
148-
}
149-
// CHECK-LABEL: func @contiguous_inner_most_dim_non_zero_idx_scalable_inner_dim(
150-
// CHECK-SAME: %[[SRC:.+]]: memref<16x1xf32>, %[[I:.+]]: index) -> vector<[8]x1xf32>
151-
// CHECK: %[[SRC_0:.+]] = memref.subview %[[SRC]]
152-
// CHECK-SAME: memref<16x1xf32> to memref<16xf32, strided<[1]>>
153-
// CHECK: %[[V:.+]] = vector.transfer_read %[[SRC_0]]
154-
// CHECK: %[[RESULT:.+]] = vector.shape_cast %[[V]] : vector<[8]xf32> to vector<[8]x1xf32>
155-
// CHECK: return %[[RESULT]]
156179

157180
// -----
158181

@@ -367,6 +390,67 @@ func.func @contiguous_inner_most_dynamic_outer_scalable_inner_dim(%a: index, %b:
367390

368391
// -----
369392

393+
// Test the impact of changing the in_bounds attribute. The behaviour will
394+
// depend on whether the index is == 0 or != 0.
395+
396+
// The index to be dropped is == 0, so it's safe to collapse. The other index
397+
// should be preserved correctly.
398+
func.func @contiguous_inner_most_zero_idx_in_bounds(%arg0: memref<16x1xf32>, %arg1: vector<8x1xf32>, %i: index) {
399+
%c0 = arith.constant 0 : index
400+
vector.transfer_write %arg1, %arg0[%i, %c0] {in_bounds = [true, true]} : vector<8x1xf32>, memref<16x1xf32>
401+
return
402+
}
403+
// CHECK-LABEL: func.func @contiguous_inner_most_zero_idx_in_bounds(
404+
// CHECK-SAME: %[[MEM:.*]]: memref<16x1xf32>,
405+
// CHECK-SAME: %[[VEC:.*]]: vector<8x1xf32>,
406+
// CHECK-SAME: %[[IDX:.*]]: index) {
407+
// CHECK: %[[SV:.*]] = memref.subview %[[MEM]][0, 0] [16, 1] [1, 1] : memref<16x1xf32> to memref<16xf32, strided<[1]>>
408+
// CHECK: %[[SC:.*]] = vector.shape_cast %[[VEC]] : vector<8x1xf32> to vector<8xf32>
409+
// CHECK: vector.transfer_write %[[SC]], %[[SV]]{{\[}}%[[IDX]]] {in_bounds = [true]} : vector<8xf32>, memref<16xf32, strided<[1]>>
410+
411+
// The index to be dropped is == 0, so it's safe to collapse. The "out of
412+
// bounds" attribute is too conservative and will be folded to "in bounds"
413+
// before the pattern runs. The other index should be preserved correctly.
414+
func.func @contiguous_inner_most_zero_idx_out_of_bounds(%arg0: memref<16x1xf32>, %arg1: vector<8x1xf32>, %i: index) {
415+
%c0 = arith.constant 0 : index
416+
vector.transfer_write %arg1, %arg0[%i, %c0] {in_bounds = [true, false]} : vector<8x1xf32>, memref<16x1xf32>
417+
return
418+
}
419+
// CHECK-LABEL: func.func @contiguous_inner_most_zero_idx_out_of_bounds
420+
// CHECK-SAME: %[[MEM:.*]]: memref<16x1xf32>,
421+
// CHECK-SAME: %[[VEC:.*]]: vector<8x1xf32>,
422+
// CHECK-SAME: %[[IDX:.*]]: index) {
423+
// CHECK: %[[SV:.*]] = memref.subview %[[MEM]][0, 0] [16, 1] [1, 1] : memref<16x1xf32> to memref<16xf32, strided<[1]>>
424+
// CHECK: %[[SC:.*]] = vector.shape_cast %[[VEC]] : vector<8x1xf32> to vector<8xf32>
425+
// CHECK: vector.transfer_write %[[SC]], %[[SV]]{{\[}}%[[IDX]]] {in_bounds = [true]} : vector<8xf32>, memref<16xf32, strided<[1]>>
426+
427+
// The index to be dropped is unknown, but since it's "in bounds", it has to be
428+
// == 0. It's safe to collapse the corresponding dim.
429+
func.func @contiguous_inner_most_dim_non_zero_idx_in_bounds(%arg0: memref<16x1xf32>, %arg1: vector<8x1xf32>, %i: index) {
430+
vector.transfer_write %arg1, %arg0[%i, %i] {in_bounds = [true, true]} : vector<8x1xf32>, memref<16x1xf32>
431+
return
432+
}
433+
// CHECK-LABEL: func @contiguous_inner_most_dim_non_zero_idx_in_bounds
434+
// CHECK-SAME: %[[MEM:.*]]: memref<16x1xf32>,
435+
// CHECK-SAME: %[[VEC:.*]]: vector<8x1xf32>,
436+
// CHECK-SAME: %[[IDX:.*]]: index) {
437+
// CHECK: %[[SV:.*]] = memref.subview %[[MEM]][0, 0] [16, 1] [1, 1] : memref<16x1xf32> to memref<16xf32, strided<[1]>>
438+
// CHECK: %[[SC:.*]] = vector.shape_cast %[[VEC]] : vector<8x1xf32> to vector<8xf32>
439+
// CHECK: vector.transfer_write %[[SC]], %[[SV]]{{\[}}%[[IDX]]] {in_bounds = [true]} : vector<8xf32>, memref<16xf32, strided<[1]>>
440+
441+
// The index to be dropped is unknown and "out of bounds" - not safe to
442+
// collapse.
443+
func.func @negative_contiguous_inner_most_dim_non_zero_idx_out_of_bounds(%arg0: memref<16x1xf32>, %arg1: vector<8x1xf32>, %i: index) {
444+
vector.transfer_write %arg1, %arg0[%i, %i] {in_bounds = [true, false]} : vector<8x1xf32>, memref<16x1xf32>
445+
return
446+
}
447+
// CHECK-LABEL: func @negative_contiguous_inner_most_dim_non_zero_idx_out_of_bounds
448+
// CHECK-NOT: memref.subview
449+
// CHECK-NOT: memref.shape_cast
450+
// CHECK: vector.transfer_write
451+
452+
// -----
453+
370454
func.func @drop_inner_most_dim(%arg0: memref<1x512x16x1xf32, strided<[8192, 16, 1, 1], offset: ?>>, %arg1: vector<1x16x16x1xf32>, %arg2: index) {
371455
%c0 = arith.constant 0 : index
372456
vector.transfer_write %arg1, %arg0[%c0, %arg2, %c0, %c0]

0 commit comments

Comments
 (0)