[Mlir-commits] [mlir] 55c45c2 - [mlir][vector] Do not unroll ops masked by vector.mask (#230461)

[llvmlistbot at llvm.org]

Author: Vaisman
Date: 2026-10-14T06:56:28+02:00
New Revision: 55c45c20269f89e6d70eecc94dd00fecd41cff64

URL: https://github.com/llvm/llvm-project/commit/55c45c20269f89e6d70eecc94dd00fecd41cff64
DIFF: https://github.com/llvm/llvm-project/commit/55c45c20269f89e6d70eecc94dd00fecd41cff64.diff

LOG: [mlir][vector] Do not unroll ops masked by vector.mask (#230461)

Added: 
    

Modified: 
    mlir/lib/Dialect/Vector/Transforms/VectorUnroll.cpp
    mlir/test/Dialect/Vector/vector-transfer-unroll.mlir
    mlir/test/Dialect/Vector/vector-unroll-options.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/Vector/Transforms/VectorUnroll.cpp b/mlir/lib/Dialect/Vector/Transforms/VectorUnroll.cpp
index dfb92c78eabc2..0ea6bc78d3b0c 100644
--- a/mlir/lib/Dialect/Vector/Transforms/VectorUnroll.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/VectorUnroll.cpp
@@ -94,6 +94,13 @@ getTargetShape(const vector::UnrollVectorOptions &options, Operation *op) {
     LDBG() << "--no filter constraint -> BAIL";
     return std::nullopt;
   }
+  // Unrolling a masked op would put several ops inside its vector.mask region,
+  // and the enclosing mask would not be sliced to match them.
+  auto maskableOp = dyn_cast<vector::MaskableOpInterface>(op);
+  if (maskableOp && maskableOp.isMasked()) {
+    LDBG() << "--op is masked -> BAIL";
+    return std::nullopt;
+  }
   assert(options.nativeShape &&
          "vector unrolling expects the native shape or native"
          "shape call back function to be set");

diff  --git a/mlir/test/Dialect/Vector/vector-transfer-unroll.mlir b/mlir/test/Dialect/Vector/vector-transfer-unroll.mlir
index bfa76ecaf77f7..828c0b3043ec6 100644
--- a/mlir/test/Dialect/Vector/vector-transfer-unroll.mlir
+++ b/mlir/test/Dialect/Vector/vector-transfer-unroll.mlir
@@ -416,3 +416,27 @@ func.func @transfer_write_unroll_dynamic_index(%mem : memref<4x4xf32>, %vec : ve
   vector.transfer_write %vec, %mem[%idx, %idx] : vector<4x2xf32>, memref<4x4xf32>
   return
 }
+
+// -----
+
+// Region-masked transfers are not unrolled: the enclosing mask would not be
+// sliced.
+
+// ALL-LABEL: func @masked_transfer_read_unroll(
+//       ALL:   %[[R:.*]] = vector.mask %{{.*}} { vector.transfer_read {{.*}} : memref<4x4xf32>, vector<4x4xf32> }
+//  ALL-NEXT:   return %[[R]]
+func.func @masked_transfer_read_unroll(%mem : memref<4x4xf32>, %idx : index, %mask : vector<4x4xi1>) -> vector<4x4xf32> {
+  %cf0 = arith.constant 0.0 : f32
+  %res = vector.mask %mask { vector.transfer_read %mem[%idx, %idx], %cf0 : memref<4x4xf32>, vector<4x4xf32> } : vector<4x4xi1> -> vector<4x4xf32>
+  return %res : vector<4x4xf32>
+}
+
+// -----
+
+// ALL-LABEL: func @masked_transfer_write_unroll(
+//  ALL-NEXT:   vector.mask %{{.*}} { vector.transfer_write {{.*}} : vector<4x4xf32>, memref<4x4xf32> }
+//  ALL-NEXT:   return
+func.func @masked_transfer_write_unroll(%mem : memref<4x4xf32>, %vec : vector<4x4xf32>, %idx : index, %mask : vector<4x4xi1>) {
+  vector.mask %mask { vector.transfer_write %vec, %mem[%idx, %idx] : vector<4x4xf32>, memref<4x4xf32> } : vector<4x4xi1>
+  return
+}

diff  --git a/mlir/test/Dialect/Vector/vector-unroll-options.mlir b/mlir/test/Dialect/Vector/vector-unroll-options.mlir
index 94f81a9ec6f30..61a74744e7a72 100644
--- a/mlir/test/Dialect/Vector/vector-unroll-options.mlir
+++ b/mlir/test/Dialect/Vector/vector-unroll-options.mlir
@@ -907,3 +907,31 @@ func.func @vector_multi_reduction_rank_mismatch(%v : vector<2x2x4xf32>, %acc: ve
 //       CHECK:   %[[V1:.*]] = vector.insert_strided_slice %[[R1]], %[[V0]] offsets = [0, 0], strides = [1, 1] : vector<1x2xf32> into vector<2x2xf32>
 //       CHECK:   %[[V2:.*]] = vector.insert_strided_slice %[[R3]], %[[V1]] offsets = [1, 0], strides = [1, 1] : vector<1x2xf32> into vector<2x2xf32>
 //       CHECK:   return %[[V2]] : vector<2x2xf32>
+
+// -----
+
+// Region-masked ops are not unrolled: the enclosing mask would not be sliced.
+
+func.func @masked_vector_contract(%lhs : vector<8x4xf32>, %rhs : vector<8x4xf32>,
+                                  %init : vector<8x8xf32>, %mask : vector<8x8x4xi1>) -> vector<8x8xf32> {
+  %0 = vector.mask %mask { vector.contract
+         {indexing_maps = [affine_map<(i, j, k) -> (i, k)>,
+                           affine_map<(i, j, k) -> (j, k)>,
+                           affine_map<(i, j, k) -> (i, j)>],
+          iterator_types = ["parallel", "parallel", "reduction"]}
+       %lhs, %rhs, %init : vector<8x4xf32>, vector<8x4xf32> into vector<8x8xf32> } : vector<8x8x4xi1> -> vector<8x8xf32>
+  return %0 : vector<8x8xf32>
+}
+// CHECK-LABEL: func @masked_vector_contract
+//  CHECK-NEXT:   %[[R:.*]] = vector.mask %{{.*}} { vector.contract {{.*}} : vector<8x4xf32>, vector<8x4xf32> into vector<8x8xf32> }
+//  CHECK-NEXT:   return %[[R]]
+
+// -----
+
+func.func @masked_vector_multi_reduction(%v : vector<4x6xf32>, %acc: vector<4xf32>, %mask : vector<4x6xi1>) -> vector<4xf32> {
+  %0 = vector.mask %mask { vector.multi_reduction #vector.kind<add>, %v, %acc [1] : vector<4x6xf32> to vector<4xf32> } : vector<4x6xi1> -> vector<4xf32>
+  return %0 : vector<4xf32>
+}
+// CHECK-LABEL: func @masked_vector_multi_reduction
+//  CHECK-NEXT:   %[[R:.*]] = vector.mask %{{.*}} { vector.multi_reduction <add>, {{.*}} : vector<4x6xf32> to vector<4xf32> }
+//  CHECK-NEXT:   return %[[R]]