[Mlir-commits] [mlir] [MLIR][Vector] Fix transferOps optimization inside maskOp (PR #90835)
Hugo Trachino
llvmlistbot at llvm.org
Mon May 13 06:35:40 PDT 2024
================
@@ -111,3 +111,18 @@ func.func @insert_slice_of_transfer_write_rank_extending(%t1 : tensor<?x?x12xf32
%1 = tensor.insert_slice %0 into %t1[4, 3, %s] [1, 5, 6] [1, 1, 1] : tensor<5x6xf32> into tensor<?x?x12xf32>
return %1 : tensor<?x?x12xf32>
}
+
+// CHECK-LABEL: func @masked_transfer_read_of_extract_slice
+// CHECK-SAME: %[[t:.*]]: tensor<?x?xf32>, %[[s1:.*]]: index, %[[s2:.*]]: index
+// CHECK-DAG: %[[m:.*]] = vector.create_mask{{.*}} : vector<5x6xi1>
+// CHECK-DAG: %[[a:.*]] = affine.apply {{.*}}[[s1]]
+// CHECK: vector.mask %[[m]] { vector.transfer_read %[[t]]{{.*}}: tensor<?x?xf32>, vector<5x6xf32> } : vector<5x6xi1> -> vector<5x6xf32>
+func.func @masked_transfer_read_of_extract_slice(%t : tensor<?x?xf32>, %s1 : index, %s2 : index) -> vector<5x6xf32> {
+ %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>
+ %mask = vector.create_mask %c3, %c4 : vector<5x6xi1>
+ %1 = vector.mask %mask {vector.transfer_read %0[%c3, %c4], %cst {in_bounds = [true, true]} : tensor<10x?xf32>, vector<5x6xf32>} : vector<5x6xi1> -> vector<5x6xf32>
+ return %1 : vector<5x6xf32>
+}
----------------
nujaa wrote:
MR split to https://github.com/llvm/llvm-project/pull/91943
https://github.com/llvm/llvm-project/pull/90835
More information about the Mlir-commits
mailing list