[LinalgExt] Add simple vectorization for map_scatter (#21090)

Allows vector types for the map_scatter input operand, and adds a
trivial vectorization pass, which converts the input of map_scatter ops
to vector operands with a vector.transfer_read.

This vectorization lowering is not complete, because the map_scatter op
in this form will not be able to lower into other operations. The final
lowering will happen as a decomposition after the map_scatter op has
been both vectorized and bufferized, and the map_scatter op will be
decomposed into a vector.scatter op. The reason this decomposition will
need to happen after bufferization, is that the vector.scatter op
requires a memref base. The decomposition pattern of the vectorized
map_scatter will come as a follow-up PR.

Eventually, the vectorized map_scatter op could instead be decomposed
into an `iree_vector_ext.transfer_scatter` op, but since the op does not
exist yet, the current plan is to use `vector.scatter` first.

---------

Signed-off-by: Max Dawkins <[email protected]>

Author: Max191

compiler/src/iree/compiler/Dialect/LinalgExt/IR/LinalgExtOps.cpp

@@ -468,6 +468,9 @@ LogicalResult MapScatterOp::verify() {
   if (getInputType().getElementType() != getOutputType().getElementType()) {
     return emitOpError("expected input and output element types to match");
   }
+  if (getInputType().getRank() == 0) {
+    return emitOpError("expected input type to have non-zero rank");
+  }
   Region &transformRegion = getTransformationRegion();
   Block &transformBody = transformRegion.getBlocks().front();
   if (transformBody.getNumArguments() != getInputRank()) {

compiler/src/iree/compiler/Dialect/LinalgExt/IR/LinalgExtOps.td

@@ -348,7 +348,7 @@ def IREELinalgExt_MapScatterOp : IREELinalgExt_PureOp<"map_scatter",
     the mask value is true, the input value will be written.
   }];
   let arguments = (ins
-      AnyRankedTensorOrMemRefType:$input,
+      AnyShaped:$input,
       AnyRankedTensorOrMemRefType:$output
   );
   let results = (outs Variadic<AnyRankedTensor>:$results);
@@ -374,6 +374,11 @@ def IREELinalgExt_MapScatterOp : IREELinalgExt_PureOp<"map_scatter",
       return getOutputType().getRank();
     }
 
+    // Return true if the map_scatter op has vector semantics.
+    bool isVectorized() {
+      return isa<VectorType>(getInputType());
+    }
+
     // Helper to apply transformations to the source index block arguments of
     // the transformation body, and replace the uses of the previous source
     // indices with the values returned by `transformationBuilder`. The argument

compiler/src/iree/compiler/Dialect/LinalgExt/IR/test/invalid.mlir

@@ -568,6 +568,21 @@ func.func @map_scatter_wrong_num_yielded_values(
 
 // -----
 
+func.func @map_scatter_0D(
+    %input: vector<f32>, %output: memref<4xf32>
+) {
+  // expected-error@+1 {{expected input type to have non-zero rank}}
+  iree_linalg_ext.map_scatter %input into %output {
+    ^bb0():
+      %mask = arith.constant true
+      %zero = arith.constant 0 : index
+      iree_linalg_ext.yield %zero, %mask : index, i1
+  } : vector<f32> into memref<4xf32>
+  return
+}
+
+// -----
+
 func.func @arg_compare_invalid_too_many_inputs(
     %input_val: tensor<2x10xf32>,
     %input_extra: tensor<2x10xf32>,

compiler/src/iree/compiler/Dialect/LinalgExt/IR/test/roundtrip.mlir

@@ -833,6 +833,28 @@ func.func @map_scatter_memref_static(
 
 // -----
 
+func.func @map_scatter_vector(
+    %input: vector<4x16x64xf32>, %output: tensor<4x16x64xf32>
+) -> tensor<4x16x64xf32> {
+  %0 = iree_linalg_ext.map_scatter %input into %output {
+    ^bb0(%idx0: index, %idx1: index, %idx2: index):
+      %mask = arith.constant true
+      iree_linalg_ext.yield %idx0, %idx1, %idx2, %mask : index, index, index, i1
+  } : vector<4x16x64xf32> into tensor<4x16x64xf32> -> tensor<4x16x64xf32>
+  return %0 : tensor<4x16x64xf32>
+}
+// CHECK-LABEL: func.func @map_scatter_vector(
+//  CHECK-SAME:   %[[INPUT:[a-zA-Z0-9_]+]]
+//  CHECK-SAME:   %[[OUTPUT:[a-zA-Z0-9_]+]]
+//       CHECK:   %[[RES:.+]] = iree_linalg_ext.map_scatter %[[INPUT]] into %[[OUTPUT]] {
+//       CHECK:     ^bb0(%[[IDX0:.+]]: index, %[[IDX1:.+]]: index, %[[IDX2:.+]]: index):
+//       CHECK:       %[[MASK:.+]] = arith.constant true
+//       CHECK:       iree_linalg_ext.yield %[[IDX0]], %[[IDX1]], %[[IDX2]], %[[MASK]]
+//       CHECK:   } : vector<4x16x64xf32> into tensor<4x16x64xf32> -> tensor<4x16x64xf32>
+//       CHECK:   return %[[RES]] : tensor<4x16x64xf32>
+
+// -----
+
 func.func @arg_compare_static(
     %input : tensor<2x6xf32>,
     %outv : tensor<2xf32>, %outi : tensor<2xindex>

compiler/src/iree/compiler/Dialect/LinalgExt/Transforms/BUILD.bazel

@@ -45,6 +45,7 @@ iree_compiler_cc_library(
         "TestReshapeFusion.cpp",
         "TileAttention.cpp",
         "TransposeFusion.cpp",
+        "VectorizeIREELinalgExtOps.cpp",
     ],
     hdrs = [
         "Passes.h",
@@ -85,5 +86,6 @@ iree_compiler_cc_library(
         "@llvm-project//mlir:TransformDialect",
         "@llvm-project//mlir:TransformUtils",
         "@llvm-project//mlir:Transforms",
+        "@llvm-project//mlir:VectorDialect",
     ],
 )

compiler/src/iree/compiler/Dialect/LinalgExt/Transforms/CMakeLists.txt

@@ -42,6 +42,7 @@ iree_cc_library(
     "TestReshapeFusion.cpp"
     "TileAttention.cpp"
     "TransposeFusion.cpp"
+    "VectorizeIREELinalgExtOps.cpp"
   DEPS
     ::PassesIncGen
     LLVMSupport
@@ -72,6 +73,7 @@ iree_cc_library(
     MLIRTransformDialect
     MLIRTransformUtils
     MLIRTransforms
+    MLIRVectorDialect
     iree::compiler::Dialect::LinalgExt::IR
     iree::compiler::Dialect::LinalgExt::Utils
     iree::compiler::Dialect::Util::IR

compiler/src/iree/compiler/Dialect/LinalgExt/Transforms/Passes.td

@@ -108,4 +108,13 @@ def TestReshapeFusionPass :
   let summary = "Test reshape fusion patterns";
 }
 
+def VectorizeIREELinalgExtOpsPass :
+    InterfacePass<"iree-linalg-ext-vectorize-ops", "mlir::FunctionOpInterface"> {
+  let summary = "Convert linalg_ext ops into their vector form.";
+  let dependentDialects = [
+    "::mlir::vector::VectorDialect",
+    "::mlir::arith::ArithDialect"
+  ];
+}
+
 #endif  // IREE_DIALECT_LINALGEXT_PASSES

compiler/src/iree/compiler/Dialect/LinalgExt/Transforms/VectorizeIREELinalgExtOps.cpp

@@ -0,0 +1,64 @@
+// Copyright 2025 The IREE Authors
+//
+// Licensed under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+#include "iree/compiler/Dialect/LinalgExt/IR/LinalgExtOps.h"
+#include "iree/compiler/Dialect/LinalgExt/Transforms/Passes.h"
+#include "iree/compiler/Dialect/LinalgExt/Transforms/Transforms.h"
+#include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/Vector/IR/VectorOps.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+
+namespace mlir::iree_compiler::IREE::LinalgExt {
+
+#define GEN_PASS_DEF_VECTORIZEIREELINALGEXTOPSPASS
+#include "iree/compiler/Dialect/LinalgExt/Transforms/Passes.h.inc"
+
+namespace {
+
+struct VectorizeStaticMapScatterOpPattern final
+    : OpRewritePattern<IREE::LinalgExt::MapScatterOp> {
+  using OpRewritePattern<IREE::LinalgExt::MapScatterOp>::OpRewritePattern;
+  LogicalResult matchAndRewrite(IREE::LinalgExt::MapScatterOp mapScatterOp,
+                                PatternRewriter &rewriter) const override {
+    if (mapScatterOp.isVectorized()) {
+      return rewriter.notifyMatchFailure(mapScatterOp,
+                                         "map_scatter is already vectorized");
+    }
+    ShapedType inputType = mapScatterOp.getInputType();
+    if (!inputType.hasStaticShape()) {
+      return rewriter.notifyMatchFailure(mapScatterOp,
+                                         "map_scatter has non-static shape");
+    }
+    Location loc = mapScatterOp.getLoc();
+    rewriter.setInsertionPoint(mapScatterOp);
+    Value zero = rewriter.create<arith::ConstantIndexOp>(loc, 0);
+    SmallVector<Value> zeros(inputType.getRank(), zero);
+    auto inputVectorType =
+        VectorType::get(inputType.getShape(), inputType.getElementType());
+    Value inputVector = rewriter.create<vector::TransferReadOp>(
+        loc, inputVectorType, mapScatterOp.getInput(), /*indices=*/zeros);
+    auto vectorizedMapScatterOp =
+        clone(rewriter, mapScatterOp, mapScatterOp.getResultTypes(),
+              {inputVector, mapScatterOp.getOutput()});
+    rewriter.replaceOp(mapScatterOp, vectorizedMapScatterOp);
+    return success();
+  }
+};
+
+struct VectorizeIREELinalgExtOpsPass final
+    : impl::VectorizeIREELinalgExtOpsPassBase<VectorizeIREELinalgExtOpsPass> {
+  void runOnOperation() {
+    MLIRContext *context = &getContext();
+    RewritePatternSet patterns(context);
+    patterns.add<VectorizeStaticMapScatterOpPattern>(context);
+    if (failed(applyPatternsGreedily(getOperation(), std::move(patterns)))) {
+      return signalPassFailure();
+    }
+  }
+};
+} // namespace
+
+} // namespace mlir::iree_compiler::IREE::LinalgExt

compiler/src/iree/compiler/Dialect/LinalgExt/Transforms/test/BUILD.bazel

@@ -28,6 +28,7 @@ iree_lit_test_suite(
             "reshape_fusion.mlir",
             "split_reduction.mlir",
             "tiling.mlir",
+            "vectorize_iree_linalg_ext_ops.mlir",
         ],
         include = ["*.mlir"],
     ),

compiler/src/iree/compiler/Dialect/LinalgExt/Transforms/test/CMakeLists.txt

@@ -26,6 +26,7 @@ iree_lit_test_suite(
     "reshape_fusion.mlir"
     "split_reduction.mlir"
     "tiling.mlir"
+    "vectorize_iree_linalg_ext_ops.mlir"
   TOOLS
     FileCheck
     iree-opt