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

[llvmlistbot at llvm.org]

https://github.com/Vaisman created https://github.com/llvm/llvm-project/pull/230461

Unrolling an operation nested in `vector.mask` creates the unrolled   slices inside the mask region, which then holds more than one operation and fails verification. The enclosing mask is not sliced to match the slices.
This affects the `contract`, `reduction`, `multi_reduction`,  `transfer_read`/`transfer_write`, and `gather` unroll patterns.

Skip region-masked operations in `getTargetShape`, which is shared by all affected unroll patterns.
Regression tests cover contractions, multi-reductions, and transfers, and fail without the fix.

This addresses the `VectorUnroll` item in #227821.

>From 04579f87349fafe4f930b1a44409ca0532d81bd4 Mon Sep 17 00:00:00 2001
From: Vasili Svirski <vasili.svirski at gmail.com>
Date: Fri, 25 Sep 2026 21:12:52 +0200
Subject: [PATCH] [mlir][vector] Do not unroll ops masked by vector.mask

Unrolling an op inside vector.mask creates the unrolled slices inside the
mask region, which then holds more than one op and fails verification.
The enclosing mask is not sliced to match the slices. This affects the
contract, reduction, multi_reduction, transfer_read/write and gather
unroll patterns.

Skip region-masked ops in getTargetShape, which is shared by all affected
unroll patterns.

Regression tests cover contract, multi_reduction and transfers, and fail
without the fix.
---
 .../Vector/Transforms/VectorUnroll.cpp        |  7 +++++
 .../Vector/vector-transfer-unroll.mlir        | 24 ++++++++++++++++
 .../Dialect/Vector/vector-unroll-options.mlir | 28 +++++++++++++++++++
 3 files changed, 59 insertions(+)

diff --git a/mlir/lib/Dialect/Vector/Transforms/VectorUnroll.cpp b/mlir/lib/Dialect/Vector/Transforms/VectorUnroll.cpp
index d58aedcc9d0234..54ed0d40f2f0df 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 bfa76ecaf77f77..828c0b3043ec62 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 d1f0961bbb735c..468b818dece275 100644
--- a/mlir/test/Dialect/Vector/vector-unroll-options.mlir
+++ b/mlir/test/Dialect/Vector/vector-unroll-options.mlir
@@ -888,3 +888,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]]