From 545c8cf77a3383f9ae9be89879c41b7f3a559286 Mon Sep 17 00:00:00 2001 From: Andrzej Warzynski Date: Sun, 9 Jun 2024 16:38:18 +0100 Subject: [PATCH 1/3] [mlir][vector] Restrict `DropInnerMostUnitDimsTransferRead` Restrict `DropInnerMostUnitDimsTransferRead` so that it fails when one of the indices to be dropped could be != 0, e.g. ``` func.func @negative_example(%A: memref<16x1xf32>, %i:index, %j:index) -> (vector<8x1xf32>) { %f0 = arith.constant 0.0 : f32 %1 = vector.transfer_read %A[%i, %j], %f0 : memref<16x1xf32>, vector<8x1xf32> return %1 : vector<8x1xf32> } ``` This is an edge case that could represent an out-of-bounds access, though that will depend on the actual value of `%j`. NOTE: This PR is limited to tests for `vector.transfer_read`. Depends on: #94490, #94604 --- .../Vector/Transforms/VectorTransforms.cpp | 15 ++++++++++++ ...tor-transfer-collapse-inner-most-dims.mlir | 24 +++++++++++++------ 2 files changed, 32 insertions(+), 7 deletions(-) diff --git a/mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp b/mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp index f29eba90c3ceb..caf1506e0db93 100644 --- a/mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp +++ b/mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp @@ -1293,6 +1293,21 @@ class DropInnerMostUnitDimsTransferRead if (dimsToDrop == 0) return failure(); + // Make sure that the indixes to be dropped are equal 0. + // TODO: Deal with cases when the indices are not 0. + auto isZeroIdx = [](Value idx) { + Attribute attr; + APInt value; + if (!matchPattern(idx, m_Constant(&attr))) + return false; + if (matchPattern(attr, m_ConstantInt(&value))) + if (!value.isZero()) + return false; + return true; + }; + if (!llvm::all_of(readOp.getIndices().take_back(dimsToDrop), isZeroIdx)) + return failure(); + auto resultTargetVecType = VectorType::get(targetType.getShape().drop_back(dimsToDrop), targetType.getElementType(), diff --git a/mlir/test/Dialect/Vector/vector-transfer-collapse-inner-most-dims.mlir b/mlir/test/Dialect/Vector/vector-transfer-collapse-inner-most-dims.mlir index a50c01898c62e..bb37d5b45520c 100644 --- a/mlir/test/Dialect/Vector/vector-transfer-collapse-inner-most-dims.mlir +++ b/mlir/test/Dialect/Vector/vector-transfer-collapse-inner-most-dims.mlir @@ -111,31 +111,41 @@ func.func @contiguous_inner_most_outer_dim_dyn_scalable_inner_dim(%a: index, %b: // ----- -func.func @contiguous_inner_most_dim_non_zero_idxs(%A: memref<16x1xf32>, %i:index, %j:index) -> (vector<8x1xf32>) { +func.func @contiguous_inner_most_dim_non_zero_idx(%A: memref<16x1xf32>, %i:index) -> (vector<8x1xf32>) { %c0 = arith.constant 0 : index %f0 = arith.constant 0.0 : f32 - %1 = vector.transfer_read %A[%i, %j], %f0 : memref<16x1xf32>, vector<8x1xf32> + %1 = vector.transfer_read %A[%i, %c0], %f0 : memref<16x1xf32>, vector<8x1xf32> return %1 : vector<8x1xf32> } -// CHECK: func @contiguous_inner_most_dim_non_zero_idxs(%[[SRC:.+]]: memref<16x1xf32>, %[[I:.+]]: index, %[[J:.+]]: index) -> vector<8x1xf32> +// CHECK: func @contiguous_inner_most_dim_non_zero_idx(%[[SRC:.+]]: memref<16x1xf32>, %[[I:.+]]: index) -> vector<8x1xf32> // CHECK: %[[SRC_0:.+]] = memref.subview %[[SRC]] // CHECK-SAME: memref<16x1xf32> to memref<16xf32, strided<[1]>> // CHECK: %[[V:.+]] = vector.transfer_read %[[SRC_0]] // CHECK: %[[RESULT:.+]] = vector.shape_cast %[[V]] : vector<8xf32> to vector<8x1xf32> // CHECK: return %[[RESULT]] +// The index to be dropped is != 0 - this is currently not supported. +func.func @negative_contiguous_inner_most_dim_non_zero_idxs(%A: memref<16x1xf32>, %i:index) -> (vector<8x1xf32>) { + %f0 = arith.constant 0.0 : f32 + %1 = vector.transfer_read %A[%i, %i], %f0 : memref<16x1xf32>, vector<8x1xf32> + return %1 : vector<8x1xf32> +} +// CHECK-LABEL: func @negative_contiguous_inner_most_dim_non_zero_idxs +// CHECK-NOT: memref.subview +// CHECK: vector.transfer_read + // Same as the top example within this split, but with the outer vector // dim scalable. Note that this example only makes sense when "8 = [8]" (i.e. // vscale = 1). This is assumed (implicitly) via the `in_bounds` attribute. -func.func @contiguous_inner_most_dim_non_zero_idxs_scalable_inner_dim(%A: memref<16x1xf32>, %i:index, %j:index) -> (vector<[8]x1xf32>) { +func.func @contiguous_inner_most_dim_non_zero_idx_scalable_inner_dim(%A: memref<16x1xf32>, %i:index) -> (vector<[8]x1xf32>) { %c0 = arith.constant 0 : index %f0 = arith.constant 0.0 : f32 - %1 = vector.transfer_read %A[%i, %j], %f0 : memref<16x1xf32>, vector<[8]x1xf32> + %1 = vector.transfer_read %A[%i, %c0], %f0 : memref<16x1xf32>, vector<[8]x1xf32> return %1 : vector<[8]x1xf32> } -// CHECK-LABEL: func @contiguous_inner_most_dim_non_zero_idxs_scalable_inner_dim( -// CHECK-SAME: %[[SRC:.+]]: memref<16x1xf32>, %[[I:.+]]: index, %[[J:.+]]: index) -> vector<[8]x1xf32> +// CHECK-LABEL: func @contiguous_inner_most_dim_non_zero_idx_scalable_inner_dim( +// CHECK-SAME: %[[SRC:.+]]: memref<16x1xf32>, %[[I:.+]]: index) -> vector<[8]x1xf32> // CHECK: %[[SRC_0:.+]] = memref.subview %[[SRC]] // CHECK-SAME: memref<16x1xf32> to memref<16xf32, strided<[1]>> // CHECK: %[[V:.+]] = vector.transfer_read %[[SRC_0]] From a9651e1c427357a3abf8220c86201ef03f8b6029 Mon Sep 17 00:00:00 2001 From: Andrzej Warzynski Date: Tue, 11 Jun 2024 15:32:46 +0100 Subject: [PATCH 2/3] fixup! [mlir][vector] Restrict `DropInnerMostUnitDimsTransferRead` Switch to using isZeroIndex from StaticValueUtils.h --- .../Dialect/Vector/Transforms/VectorTransforms.cpp | 12 +----------- 1 file changed, 1 insertion(+), 11 deletions(-) diff --git a/mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp b/mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp index caf1506e0db93..ecb8e8060aed0 100644 --- a/mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp +++ b/mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp @@ -1295,17 +1295,7 @@ class DropInnerMostUnitDimsTransferRead // Make sure that the indixes to be dropped are equal 0. // TODO: Deal with cases when the indices are not 0. - auto isZeroIdx = [](Value idx) { - Attribute attr; - APInt value; - if (!matchPattern(idx, m_Constant(&attr))) - return false; - if (matchPattern(attr, m_ConstantInt(&value))) - if (!value.isZero()) - return false; - return true; - }; - if (!llvm::all_of(readOp.getIndices().take_back(dimsToDrop), isZeroIdx)) + if (!llvm::all_of(readOp.getIndices().take_back(dimsToDrop), isZeroIndex)) return failure(); auto resultTargetVecType = From c55adde64eec9087611bcf12da793584d69bce34 Mon Sep 17 00:00:00 2001 From: Andrzej Warzynski Date: Tue, 11 Jun 2024 19:10:04 +0100 Subject: [PATCH 3/3] fixup! fixup! [mlir][vector] Restrict `DropInnerMostUnitDimsTransferRead` Fix typo --- mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp b/mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp index ecb8e8060aed0..ea4a02f2f2e77 100644 --- a/mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp +++ b/mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp @@ -1293,7 +1293,7 @@ class DropInnerMostUnitDimsTransferRead if (dimsToDrop == 0) return failure(); - // Make sure that the indixes to be dropped are equal 0. + // Make sure that the indices to be dropped are equal 0. // TODO: Deal with cases when the indices are not 0. if (!llvm::all_of(readOp.getIndices().take_back(dimsToDrop), isZeroIndex)) return failure();