[linalg] Vectorization failure - masked scalar read + broadcast

[banach-space]

REPRODUCER

func.func @vectorization_test(%extracted_slice : tensor<1x1x3xi32>, %arg0: index, %arg2: index, %3: tensor<2x4xi32>, %4: tensor<1x3x2x4xi32>) -> tensor<1x1x3xi32>{
%c0 = arith.constant 0 :index

%8 = linalg.generic {
  indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1, d2)>],
  iterator_types = ["parallel", "parallel", "parallel"]}
  outs(%extracted_slice : tensor<1x1x3xi32>) {
  ^bb0(%out: i32):
    %9 = linalg.index 0 : index
    %extracted = tensor.extract %3[%9, %c0] : tensor<2x4xi32>
    %14 = arith.index_cast %extracted : i32 to index
    %extracted_2 = tensor.extract %4[%c0, %14, %14, %14] : tensor<1x3x2x4xi32>
    linalg.yield %extracted_2 : i32
  } -> tensor<1x1x3xi32>

  return %8 : tensor<1x1x3xi32>
}
module attributes {transform.with_named_sequence} {
  transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
    %0 = transform.structured.match ops{["linalg.generic"]} in %arg1 : (!transform.any_op) -> !transform.any_op
    %1 = transform.get_parent_op %0 {isolated_from_above} : (!transform.any_op) -> !transform.any_op
    // %2 = transform.structured.vectorize_children_and_apply_patterns %1  { vectorize_nd_extract } : (!transform.any_op) -> !transform.any_op
    transform.structured.vectorize %0 vector_sizes [1, 1, 4] {vectorize_nd_extract} : !transform.any_op
    transform.yield
  }
}

ERROR LOG

../file.mlir:10:18: error: 'vector.mask' op expects a 'vector<i1>' mask for the maskable operation
    %extracted = tensor.extract %3[%9, %c0] : tensor<2x4xi32>
                 ^
../file.mlir:10:18: note: see current operation:
%17 = "vector.mask"(%6) ({
  %34 = "vector.transfer_read"(%arg3, %15, %0, %16) <{in_bounds = [true, true, true], operandSegmentSizes = array<i32: 1, 2, 1, 0>, permutation_map = affine_map<(d0, d1) -> (0, 0, 0)>}> : (tensor<2x4xi32>, index, index, i32) -> vector<1x1x4xi32>
  "vector.yield"(%34) : (vector<1x1x4xi32>) -> ()
}) : (vector<1x1x4xi1>) -> vector<1x1x4xi32>

ANALYSIS

1. Type of vectorization: masked
2. Op: tensor.extract --> %extracted = tensor.extract %3[%9, %c0] : tensor<2x4xi32>(effectively a scalar read + broadcast).
3. The Vectorizer output (generated by vectorizeAsTensorExtract):

 %34 = "vector.transfer_read"(%arg3, %15, %0, %16) <{in_bounds = [true, true, true], operandSegmentSizes = array<i32: 1, 2, 1, 0>, permutation_map = affine_map<(d0, d1) -> (0, 0, 0)>}> : (tensor<2x4xi32>, index, index, i32) -> vector<1x1x4xi32>

Looks like the mask generated by the vectorizer doesn't match the "expected mask" computed by MaskOp::verify:

• Linalg vectorizer, when creating vector.mask, generates a mask based on the static loop sizes and input vector sizes. This gives: vector<1x1x4xi1>.

• The Op verifier uses inferTransferOpMaskType, which has no access to the LinalgOp information and instead looks at the permutation map of the masked op, vector.transfer_read. And that yields vector<i1> (based on permutation_map = affine_map<(d0, d1) -> (0, 0, 0)).

To me, the output from the Vectorizer is correct.

RELEVANT DATA POINT

Looking at this example from "invalid.mlir":

llvm-project/mlir/test/Dialect/Vector/invalid.mlir

Lines 475 to 483 in 8ff2da7

	func.func @test_vector.transfer_read(%arg0: memref<?x?x?xf32>) {
	%c1 = arith.constant 1 : i1
	%c3 = arith.constant 3 : index
	%cst = arith.constant 3.0 : f32
	// expected-note@+1 {{prior use here}}
	%mask = vector.splat %c1 : vector<3x8x7xi1>
	// expected-error@+1 {{expects different type than prior uses: 'vector<3x7xi1>' vs 'vector<3x8x7xi1>'}}
	%0 = vector.transfer_read %arg0[%c3, %c3, %c3], %cst, %mask {permutation_map = affine_map<(d0, d1, d2)->(d0, 0, d2)>} : memref<?x?x?xf32>, vector<3x8x7xf32>
	}

To me, the error is wrong and the example is in fact correct (as in, vector<3x8x7xi1> as a mask makes sense to me). Specifically, as per the docs (emphasis mine):

The masked-off lanes in the result vector are taken from the corresponding lanes of the pass-thru argument, if provided, or left unmodified, otherwise.

Doesn't this mean that the mask shape should always match the result vector shape?

PROPOSED SOLUTION

• Either fix how the vectorizer handles broadcast dims, or
• Update inferTransferOpMaskType and, in general, the semantics of broadcast dims when masking is used.

CC @dcaballe

[llvmbot]

@llvm/issue-subscribers-mlir

Author: Andrzej Warzyński (banach-space)

**REPRODUCER** ```mlir func.func @vectorization_test(%extracted_slice : tensor<1x1x3xi32>, %arg0: index, %arg2: index, %3: tensor<2x4xi32>, %4: tensor<1x3x2x4xi32>) -> tensor<1x1x3xi32>{ %c0 = arith.constant 0 :index

%8 = linalg.generic {
indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1, d2)>],
iterator_types = ["parallel", "parallel", "parallel"]}
outs(%extracted_slice : tensor<1x1x3xi32>) {
^bb0(%out: i32):
%9 = linalg.index 0 : index
%extracted = tensor.extract %3[%9, %c0] : tensor<2x4xi32>
%14 = arith.index_cast %extracted : i32 to index
%extracted_2 = tensor.extract %4[%c0, %14, %14, %14] : tensor<1x3x2x4xi32>
linalg.yield %extracted_2 : i32
} -> tensor<1x1x3xi32>

return %8 : tensor<1x1x3xi32>
}
module attributes {transform.with_named_sequence} {
transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
%0 = transform.structured.match ops{["linalg.generic"]} in %arg1 : (!transform.any_op) -> !transform.any_op
%1 = transform.get_parent_op %0 {isolated_from_above} : (!transform.any_op) -> !transform.any_op
// %2 = transform.structured.vectorize_children_and_apply_patterns %1 { vectorize_nd_extract } : (!transform.any_op) -> !transform.any_op
transform.structured.vectorize %0 vector_sizes [1, 1, 4] {vectorize_nd_extract} : !transform.any_op
transform.yield
}
}

**ERROR LOG**
```bash
../file.mlir:10:18: error: 'vector.mask' op expects a 'vector<i1>' mask for the maskable operation
    %extracted = tensor.extract %3[%9, %c0] : tensor<2x4xi32>
                 ^
../file.mlir:10:18: note: see current operation:
%17 = "vector.mask"(%6) ({
  %34 = "vector.transfer_read"(%arg3, %15, %0, %16) <{in_bounds = [true, true, true], operandSegmentSizes = array<i32: 1, 2, 1, 0>, permutation_map = affine_map<(d0, d1) -> (0, 0, 0)>}> : (tensor<2x4xi32>, index, index, i32) -> vector<1x1x4xi32>
  "vector.yield"(%34) : (vector<1x1x4xi32>) -> ()
}) : (vector<1x1x4xi1>) -> vector<1x1x4xi32>

ANALYSIS

1. Type of vectorization: masked
2. Op: tensor.extract --> %extracted = tensor.extract %3[%9, %c0] : tensor<2x4xi32>(effectively a scalar read + broadcast).
3. The Vectorizer output (generated by vectorizeAsTensorExtract):

 %34 = "vector.transfer_read"(%arg3, %15, %0, %16) <{in_bounds = [true, true, true], operandSegmentSizes = array<i32: 1, 2, 1, 0>, permutation_map = affine_map<(d0, d1) -> (0, 0, 0)>}> : (tensor<2x4xi32>, index, index, i32) -> vector<1x1x4xi32>

Looks like the mask generated by the vectorizer doesn't match the "expected mask" computed by MaskOp::verify:

• Linalg vectorizer, when creating vector.mask, generates a mask based on the static loop sizes and input vector sizes. This gives: vector<1x1x4xi1>.

• The Op verifier uses inferTransferOpMaskType, which has no access to the LinalgOp information and instead looks at the permutation map of the masked op, vector.transfer_read. And that yields vector<i1> (based on permutation_map = affine_map<(d0, d1) -> (0, 0, 0)).

To me, the output from the Vectorizer is correct.

RELEVANT DATA POINT

Looking at this example from "invalid.mlir":

llvm-project/mlir/test/Dialect/Vector/invalid.mlir

Lines 475 to 483 in 8ff2da7

	func.func @test_vector.transfer_read(%arg0: memref<?x?x?xf32>) {
	%c1 = arith.constant 1 : i1
	%c3 = arith.constant 3 : index
	%cst = arith.constant 3.0 : f32
	// expected-note@+1 {{prior use here}}
	%mask = vector.splat %c1 : vector<3x8x7xi1>
	// expected-error@+1 {{expects different type than prior uses: 'vector<3x7xi1>' vs 'vector<3x8x7xi1>'}}
	%0 = vector.transfer_read %arg0[%c3, %c3, %c3], %cst, %mask {permutation_map = affine_map<(d0, d1, d2)->(d0, 0, d2)>} : memref<?x?x?xf32>, vector<3x8x7xf32>
	}

To me, the error is wrong and the example is in fact correct (as in, vector<3x8x7xi1> as a mask makes sense to me). Specifically, as per the docs (emphasis mine):

> The masked-off lanes in the result vector are taken from the corresponding lanes of the pass-thru argument, if provided, or left unmodified, otherwise.

Doesn't this mean that the mask shape should always match the result vector shape?

PROPOSED SOLUTION

• Either fix how the vectorizer handles broadcast dims, or
• Update inferTransferOpMaskType and, in general, the semantics of broadcast dims when masking is used.

CC @dcaballe

[dcaballe]

Thanks for reporting this! Indeed, the result of inferTransferOpMaskType is wrong. The mask for transfer ops is inferred by inversing its permutation map as masking applies to the memory operation side of it, which is a single element in this case. The code inversing the mask is probably not handling this case properly. I think compressUnusedDims might be turning the broadcasting map (e.g., (d0, d1, ..., dn) -> (0, 0, ... 0) into a map with no inputs (e.g., () -> (0, 0, ... 0)) and things get off rail after that.

I can look into this later this week, if needed. Just let me know.

[banach-space]

Thanks for replying!

I can look into this later this week, if needed. Just let me know.

Can you take look at the test that I linked? I think that the current semantics are not 100% correct.

Also, let me share what I've tried already (see diff below). I've updated inferTransferOpMaskType with a pre-processing step so that:

1. (new!) Replace broadcast dims with identities (d0, d1, ..., dn) -> (0, 0, ... 0) --> (d0, d1, ..., dn) -> (d0, d1, ... dn)
2. (old) Inverse the map: (d0, d1, ..., dn) -> (d0, d1, ... dn) --> (d0, d1, ..., dn) -> (d0, d1, ... dn).
3. (old) Finally, proceed with the rest.

This makes sense under the assumption that:

for broadcast dims, the corresponding mask size should be identical to the output vector size.

This would mean that this example, extracted from invalid.mlir, is actually valid:

%mask = vector.splat %c1 : vector<3x8x7xi1> 
%0 = vector.transfer_read %arg0[%c3, %c3, %c3], %cst, %mask {permutation_map = affine_map<(d0, d1, d2)->(d0, 0, d2)>} : memref<?x?x?xf32>, vector<3x8x7xf32> 

My update to inferTransferOpMaskType fixes the repro, but ... brakes "invalid.mlir" and other tests. I'm trying to figure out whether it's my assumptions or the tests that are wrong 😅

---

diff --git a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
index db199a46e163..5760c8243261 100644
--- a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
+++ b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
@@ -4118,7 +4118,21 @@ void TransferReadOp::print(OpAsmPrinter &p) {
 VectorType mlir::vector::inferTransferOpMaskType(VectorType vecType,
                                                  AffineMap permMap) {
   auto i1Type = IntegerType::get(permMap.getContext(), 1);
-  AffineMap invPermMap = inversePermutation(compressUnusedDims(permMap));
+  SmallVector<AffineExpr> newResults;
+  for (auto [idx, expr] : llvm::enumerate(permMap.getResults())) {
+    auto constExpr = dyn_cast<AffineConstantExpr>(expr);
+    if (constExpr && constExpr.getValue() == 0) {
+      auto dim = getAffineDimExpr(idx, vecType.getContext());
+      newResults.push_back(dim);
+    }
+    else
+      newResults.push_back(expr);
+  }
+  AffineMap newMap =
+      AffineMap::get(std::max<int64_t>(permMap.getNumDims(), newResults.size()),
+                     permMap.getNumSymbols(), newResults, vecType.getContext());
+  AffineMap invPermMap = inversePermutation(compressUnusedDims(newMap));
   assert(invPermMap && "Inversed permutation map couldn't be computed");
   SmallVector<int64_t, 8> maskShape = invPermMap.compose(vecType.getShape());

[dcaballe]

Actually, the mask inference logic is (almost) correct. There is a type mismatch due to usual 0-D vs vector<1xi1> problem but I fixed that here: #116526

The actual problem is in the tensor::ExtractOp vectorization logic. It looks like it's using the default masking path, which assumes an identity masking map (i.e., it will mask the operation based on its result type). However, masking a transfer op is more complicated than that. As described in the vector.transfer_read doc, the mask applies to the loaded elements before any broadcasting or permutation. That is why the transfer op's mask inference logic is walking the result type through the inversed permutation map to compute the mask. We already have some utilities in the vectorizer that are building vector transfer ops and masking them properly. You should be able to follow a similar approach for the tensor::ExtractOp case.

Hopefully that helps!