[Mlir-commits] [mlir] [mlir][linalg] Fix mask rank for masked contiguous `tensor.extract` (PR #206207)

[Andrzej Warzyński]

================
@@ -477,3 +477,52 @@ module attributes {transform.with_named_sequence} {
     transform.yield
   }
 }
+
+// -----
+
+// A contiguous `tensor.extract` from a source whose rank is *smaller* than the
+// iteration space (here: a 1-D source inside a 2-D loop nest). The contiguous
+// `vector.transfer_read` only reads the trailing (contiguous) dim - the leading
+// dim is broadcast via the permutation map - so its mask must be rank-reduced
+// (`vector<4xi1>`) rather than the full iteration-space mask (`vector<1x4xi1>`).
+// Regression test: previously the full iteration mask was attached, producing
+// an invalid `vector.transfer_read` ("inferred mask type ... don't match").
+
+func.func @masked_contiguous_extract_rank_reducing_mask(
+    %src: tensor<16xf32>,
+    %output : tensor<1x3xf32>,
+    %idx: index) -> tensor<1x3xf32> {
+  %1 = linalg.generic {
+    indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>],
+    iterator_types = ["parallel", "parallel"]
+  } outs(%output : tensor<1x3xf32>) {
+  ^bb0(%out: f32):
+    %2 = linalg.index 1 : index
+    %3 = affine.apply affine_map<(d0, d1) -> (d0 + d1)>(%2, %idx)
+    %extracted = tensor.extract %src[%3] : tensor<16xf32>
+    linalg.yield %extracted : f32
+  } -> tensor<1x3xf32>
+  return %1 : tensor<1x3xf32>
----------------
banach-space wrote:

>  I checked. It looks like having a leading dimension > 1 doesn't actually use this code branch, it winds up being a gather instead. 

That's all I needed, thank you!

https://github.com/llvm/llvm-project/pull/206207