Skip to content

Commit

Permalink
[mlir][vector] Check minor identity map in FoldExtractSliceIntoTransf…
Browse files Browse the repository at this point in the history
…erRead

vecotr.transfer_read ops with minor identity indexing map is rank
reducing, with implicit leading unit dimensions. This should be
a natural extension to support in addition to full identity indexing
maps.

Reviewed By: ThomasRaoux

Differential Revision: https://reviews.llvm.org/D133883
  • Loading branch information
antiagainst committed Sep 14, 2022
1 parent 6e675fb commit 8a5cb93
Show file tree
Hide file tree
Showing 2 changed files with 19 additions and 1 deletion.
2 changes: 1 addition & 1 deletion mlir/lib/Dialect/Vector/IR/VectorOps.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -3294,7 +3294,7 @@ struct FoldExtractSliceIntoTransferRead
return failure();
if (xferOp.hasOutOfBoundsDim())
return failure();
if (!xferOp.getPermutationMap().isIdentity())
if (!xferOp.getPermutationMap().isMinorIdentity())
return failure();
if (xferOp.getMask())
return failure();
Expand Down
18 changes: 18 additions & 0 deletions mlir/test/Dialect/Vector/canonicalize.mlir
Original file line number Diff line number Diff line change
Expand Up @@ -1169,6 +1169,24 @@ func.func @transfer_read_of_extract_slice(%t : tensor<?x?xf32>, %s1 : index, %s2

// -----

// CHECK-LABEL: func @transfer_read_of_extract_slice(
// CHECK-SAME: %[[t:.*]]: tensor<?x?xf32>, %[[s1:.*]]: index, %[[s2:.*]]: index
// CHECK-DAG: %[[c4:.*]] = arith.constant 4 : index
// CHECK-DAG: %[[c8:.*]] = arith.constant 8 : index
// CHECK: %[[add:.*]] = arith.addi %[[s1]], %[[c4]]
// CHECK: %[[r:.*]] = vector.transfer_read %[[t]][%[[c8]], %[[add]]], %{{.*}} {in_bounds = [true]} : tensor<?x?xf32>, vector<6xf32>
// CHECK: return %[[r]]
func.func @transfer_read_of_extract_slice(%t : tensor<?x?xf32>, %s1 : index, %s2 : index) -> vector<6xf32> {
%c3 = arith.constant 3 : index
%c4 = arith.constant 4 : index
%cst = arith.constant 0.0 : f32
%0 = tensor.extract_slice %t[5, %s1] [10, %s2] [1, 1] : tensor<?x?xf32> to tensor<10x?xf32>
%1 = vector.transfer_read %0[%c3, %c4], %cst {in_bounds = [true]} : tensor<10x?xf32>, vector<6xf32>
return %1 : vector<6xf32>
}

// -----

// CHECK-LABEL: func @transfer_read_of_extract_slice_rank_reducing(
// CHECK-SAME: %[[t:.*]]: tensor<?x?x?xf32>, %[[s1:.*]]: index, %[[s2:.*]]: index
// CHECK-DAG: %[[c3:.*]] = arith.constant 3 : index
Expand Down

0 comments on commit 8a5cb93

Please sign in to comment.