[mlir][linalg] Fix vectorization of tensor.extract (llvm#118105)

The example below demonstrates a "scalar read followed by a broadcast" pattern for `tensor.extract`: ```mlir #map = affine_map<(d0, d1, d2) -> (d0, d1, d2)> func.func @scalar_broadcast( %init : tensor<1x1x3xi32>, %src: tensor<1x3x2x4xi32>, %idx :index) -> tensor<1x1x3xi32> { %c0 = arith.constant 0 :index %res = linalg.generic { indexing_maps = [#map], iterator_types = ["parallel", "parallel", "parallel"]} outs(%init : tensor<1x1x3xi32>) { ^bb0(%out: i32): %val = tensor.extract %src[%idx, %idx, %idx, %idx] : tensor<1x3x2x4xi32> linalg.yield %val : i32 } -> tensor<1x1x3xi32> return %res : tensor<1x1x3xi32> } ``` The default masking path within the Linalg vectorizer, which assumes an identity masking map, is not suitable here. Indeed: * identity != broadcast. This patch ensures masking is handled in the `vectorizeTensorExtract` hook, which has the necessary context for proper handling. Fixes llvm#116197

Author: banach-space

mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp

@@ -1165,8 +1165,18 @@ vectorizeTensorExtract(RewriterBase &rewriter, VectorizationState &state,
  loc, resultType, extractOp.getTensor(), transferReadIdxs,
  permutationMap, inBounds);
 
+ // Mask this broadcasting xfer_read here rather than relying on the generic
+ // path (the generic path assumes identity masking map, which wouldn't be
+ // valid here).
+ SmallVector<int64_t> readMaskShape = {1};
+ auto readMaskType = VectorType::get(readMaskShape, rewriter.getI1Type());
+ auto allTrue = rewriter.create<vector::ConstantMaskOp>(
+ loc, readMaskType, vector::ConstantMaskKind::AllTrue);
+ auto *maskedReadOp =
+ mlir::vector::maskOperation(rewriter, transferReadOp, allTrue);
+
  LDBG("Vectorised as scalar broadcast load: " << extractOp << "\n");
- return VectorizationResult{VectorizationStatus::NewOp, transferReadOp};
+ return VectorizationResult{VectorizationStatus::NewOp, maskedReadOp};
  }
 
  // 2b. Handle contiguous access.

mlir/test/Dialect/Linalg/vectorize-tensor-extract-masked.mlir

@@ -425,3 +425,55 @@ module attributes {transform.with_named_sequence} {
  transform.yield
  }
 }
+
+// -----
+
+#map = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
+func.func @scalar_broadcast(%init : tensor<1x1x3xi32>, %src: tensor<1x3x2x4xi32>, %idx :index) -> tensor<1x1x3xi32> {
+
+ %c0 = arith.constant 0 :index
+
+ %res = linalg.generic {
+ indexing_maps = [#map],
+ iterator_types = ["parallel", "parallel", "parallel"]}
+ outs(%init : tensor<1x1x3xi32>) {
+ ^bb0(%out: i32):
+ %val = tensor.extract %src[%idx, %idx, %idx, %idx] : tensor<1x3x2x4xi32>
+ linalg.yield %val : i32
+ } -> tensor<1x1x3xi32>
+
+ return %res : tensor<1x1x3xi32>
+}
+
+// CHECK: #[[$MAP:.+]] = affine_map<(d0, d1, d2, d3) -> (0, 0, 0)>
+// CHECK-LABEL: func.func @scalar_broadcast(
+// CHECK-SAME: %[[INIT:.*]]: tensor<1x1x3xi32>,
+// CHECK-SAME: %[[SRC:.*]]: tensor<1x3x2x4xi32>,
+// CHECK-SAME: %[[IDX:.*]]: index) -> tensor<1x1x3xi32> {
+
+/// Compute the mask for saving the final result
+// CHECK: %[[C1:.*]] = arith.constant 1 : index
+// CHECK: %[[C1_2:.*]] = arith.constant 1 : index
+// CHECK: %[[C3:.*]] = arith.constant 3 : index
+// CHECK: %[[MASK_RES:.*]] = vector.create_mask %[[C1]], %[[C1_2]], %[[C3]] : vector<1x1x4xi1>
+
+/// Read and broadcast the scalar
+// CHECK: %[[PAD:.*]] = arith.constant 0 : i32
+// CHECK: %[[MASK_READ:.*]] = vector.constant_mask [1] : vector<1xi1>
+// CHECK: %[[READ:.*]] = vector.mask %[[MASK_READ]] {
+// CHECK-SAME: vector.transfer_read %[[SRC]]{{\[}}%[[IDX]], %[[IDX]], %[[IDX]], %[[IDX]]], %[[PAD]]
+// CHECK-SAME: {in_bounds = [true, true, true], permutation_map = #[[$MAP]]} : tensor<1x3x2x4xi32>, vector<1x1x4xi32>
+// CHECK-SAME: } : vector<1xi1> -> vector<1x1x4xi32>
+
+/// Save the result in the output tensor
+// CHECK: vector.mask %[[MASK_RES]] {
+// CHECK-SAME: vector.transfer_write %[[READ]], %[[INIT]]{{.*}} {in_bounds = [true, true, true]} : vector<1x1x4xi32>, tensor<1x1x3xi32>
+// CHECK-SAME: } : vector<1x1x4xi1> -> tensor<1x1x3xi32>
+
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%module: !transform.any_op {transform.readonly}) {
+ %0 = transform.structured.match ops{["linalg.generic"]} in %module : (!transform.any_op) -> !transform.any_op
+ transform.structured.vectorize %0 vector_sizes [1, 1, 4] {vectorize_nd_extract} : !transform.any_op
+ transform.yield
+ }
+}

mlir/test/Dialect/Linalg/vectorize-tensor-extract.mlir

@@ -66,7 +66,7 @@ module attributes {transform.with_named_sequence} {
 // -----
 
 #map = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
-func.func @vectorize_nd_tensor_extract_constant_idx(%arg0: tensor<3x3xf32>, %arg2: tensor<1x1x3xf32>) -> tensor<1x1x3xf32> {
+func.func @vectorize_nd_tensor_extract_scalar_broadcast(%arg0: tensor<3x3xf32>, %arg2: tensor<1x1x3xf32>) -> tensor<1x1x3xf32> {
  %c0 = arith.constant 1 : index
  %c1 = arith.constant 2 : index
  %2 = linalg.generic {
@@ -80,17 +80,17 @@ func.func @vectorize_nd_tensor_extract_constant_idx(%arg0: tensor<3x3xf32>, %arg
  return %2 : tensor<1x1x3xf32>
 }
 
-// CHECK: #[[$MAP:.*]] = affine_map<(d0, d1) -> (0, 0, 0)>
-// CHECK-LABEL: func.func @vectorize_nd_tensor_extract_constant_idx(
+// CHECK: #[[$MAP:.+]] = affine_map<(d0, d1) -> (0, 0, 0)>
+// CHECK-LABEL: func.func @vectorize_nd_tensor_extract_scalar_broadcast(
 // CHECK-SAME: %[[ARG_0:.*]]: tensor<3x3xf32>,
 // CHECK-SAME: %[[ARG_1:.*]]: tensor<1x1x3xf32>) -> tensor<1x1x3xf32> {
 // CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index
 // CHECK-DAG: %[[C2:.*]] = arith.constant 2 : index
-// CHECK-DAG: %[[C0_f32_2:.*]] = arith.constant 0.000000e+00 : f32
-// CHECK-DAG: %[[C0_f32:.*]] = arith.constant 0.000000e+00 : f32
-// CHECK: %[[READ:.*]] = vector.transfer_read %[[ARG_0]][%[[C1]], %[[C2]]], %[[C0_f32]] {in_bounds = [true, true, true], permutation_map = #[[$MAP]]} : tensor<3x3xf32>, vector<1x1x3xf32>
-// CHECK: %[[C0_4:.*]] = arith.constant 0 : index
-// CHECK: vector.transfer_write %[[READ]], %[[ARG_1]][%[[C0_4]], %[[C0_4]], %[[C0_4]]]  : vector<1x1x3xf32>, tensor<1x1x3xf32>
+// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
+// CHECK:  %[[MASK:.*]] = vector.constant_mask [1] : vector<1xi1>
+// CHECK: %[[READ:.*]] = vector.mask %[[MASK]] { vector.transfer_read %[[ARG_0]][%[[C1]], %[[C2]]], {{.*}} {in_bounds = [true, true, true], permutation_map = #[[$MAP]]} : tensor<3x3xf32>, vector<1x1x3xf32> } : vector<1xi1> -> vector<1x1x3xf32>
+// CHECK: %[[C0_2:.*]] = arith.constant 0 : index
+// CHECK: vector.transfer_write %[[READ]], %[[ARG_1]]{{\[}}%[[C0_2]], %[[C0_2]], %[[C0_2]]] : vector<1x1x3xf32>, tensor<1x1x3xf32>
 
 module attributes {transform.with_named_sequence} {
  transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
@@ -823,7 +823,7 @@ func.func @vectorize_scalar_broadcast_column_tensor(%in: tensor<1x1x4xi32>) -> t
  return %out:tensor<1x1x4xi32>
 }
 
-// CHECK: #[[$ATTR_1:.+]] = affine_map<(d0, d1) -> (0, 0, 0)>
+// CHECK: #[[$MAP:.+]] = affine_map<(d0, d1) -> (0, 0, 0)>
 // CHECK-LABEL: func.func @vectorize_scalar_broadcast_column_tensor(
 // CHECK-SAME: %[[VAL_0:.*]]: tensor<1x1x4xi32>) -> tensor<1x1x4xi32> {
 // CHECK: %[[VAL_1:.*]] = arith.constant 4 : index
@@ -844,12 +844,14 @@ func.func @vectorize_scalar_broadcast_column_tensor(%in: tensor<1x1x4xi32>) -> t
 // CHECK: %[[VAL_16:.*]] = arith.constant dense<true> : vector<1x1x4xi1>
 // CHECK: %[[VAL_17:.*]] = arith.constant dense<0> : vector<1x1x4xi32>
 // CHECK: %[[VAL_18:.*]] = arith.constant 0 : index
-// CHECK: %[[VAL_20:.*]] = vector.shape_cast %[[VAL_15]] : vector<1x1x4xindex> to vector<4xindex>
-// CHECK: %[[VAL_21:.*]] = vector.extract %[[VAL_20]][0] : index from vector<4xindex>
-// CHECK: %[[VAL_22:.*]] = arith.constant 0 : i32
-// CHECK: %[[VAL_23:.*]] = vector.transfer_read %[[VAL_3]]{{\[}}%[[VAL_21]], %[[VAL_2]]], %[[VAL_22]] {in_bounds = [true, true, true], permutation_map = #[[$ATTR_1]]} : tensor<15x1xi32>, vector<1x1x4xi32>
+// CHECK: %[[VAL_19:.*]] = vector.shape_cast %[[VAL_15]] : vector<1x1x4xindex> to vector<4xindex>
+// CHECK: %[[VAL_20:.*]] = vector.extract %[[VAL_19]][0] : index from vector<4xindex>
+// CHECK: %[[VAL_21:.*]] = arith.constant 0 : i32
+// CHECK: %[[VAL_22:.*]] = vector.constant_mask [1] : vector<1xi1>
+// CHECK: %[[VAL_23:.*]] = vector.mask %[[VAL_22]] { vector.transfer_read %[[VAL_3]]{{\[}}%[[VAL_20]], %[[VAL_2]]], %[[VAL_21]] {in_bounds = [true, true, true], permutation_map = #[[$MAP]]} : tensor<15x1xi32>, vector<1x1x4xi32> } : vector<1xi1> -> vector<1x1x4xi32>
 // CHECK: %[[VAL_24:.*]] = arith.constant 0 : index
 // CHECK: %[[VAL_25:.*]] = vector.transfer_write %[[VAL_23]], %[[VAL_0]]{{\[}}%[[VAL_24]], %[[VAL_24]], %[[VAL_24]]] : vector<1x1x4xi32>, tensor<1x1x4xi32>
+// CHECK: return %[[VAL_25]] : tensor<1x1x4xi32>
 
 module attributes {transform.with_named_sequence} {
  transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {