[VectorExt] Implement BufferizationInterface for transfer_gather (iree-org#21219)

Author: Groverkss

compiler/src/iree/compiler/Codegen/Common/test/iree_comprehensive_bufferize.mlir

@@ -3106,3 +3106,20 @@ func.func @cache_swizzle_resource_cast(%stride: index) {
 //       CHECK:   %[[CAST:.+]] = amdgpu.fat_raw_buffer_cast %[[INPUT]]
 //  CHECK-SAME:     cacheSwizzleStride(%[[TRUNC]]) resetOffset
 //  CHECK-SAME:     memref<2xf32, #hal.descriptor_type<storage_buffer>> to memref<2xf32, #amdgpu.address_space<fat_raw_buffer>>
+
+// -----
+
+func.func @transfer_gather(%source : tensor<?x64xf16>, %indices: vector<8xindex>) -> vector<8x64xf16> {
+  %c0 = arith.constant 0 : index
+  %cst = arith.constant 0.0 : f16
+  %out = iree_vector_ext.transfer_gather %source[%c0, %c0][%indices: vector<8xindex>, None], %cst {
+    indexed_maps = [affine_map<(d0, d1) -> (d0)>]
+  } : tensor<?x64xf16>, vector<8x64xf16>
+  return %out : vector<8x64xf16>
+}
+
+// CHECK-LABEL: func.func @transfer_gather
+// CHECK-SAME: %[[SOURCE:.+]]: tensor<?x64xf16>, %[[INDICES:.+]]: vector<8xindex>
+// CHECK: %[[C0:.+]] = arith.constant 0 : index
+// CHECK: %[[BUFFER:.+]] = bufferization.to_buffer %[[SOURCE]]
+// CHECK: iree_vector_ext.transfer_gather %[[BUFFER]][%[[C0]], %[[C0]]][%[[INDICES]]: vector<8xindex>, None]

compiler/src/iree/compiler/Codegen/Dialect/VectorExt/Transforms/BUILD.bazel

@@ -35,11 +35,13 @@ iree_gentbl_cc_library(
 iree_compiler_cc_library(
     name = "VectorExtTransforms",
     srcs = [
+        "BufferizationInterfaces.cpp",
         "Passes.cpp",
         "VectorExtFoldUnitExtentDims.cpp",
         "VectorizeIREEVectorExtOps.cpp",
     ],
     hdrs = [
+        "BufferizationInterfaces.h",
         "Passes.h",
         "Passes.h.inc",
         "Transforms.h",
@@ -51,6 +53,9 @@ iree_compiler_cc_library(
         "//compiler/src/iree/compiler/Dialect/LinalgExt/Utils",
         "@llvm-project//llvm:Support",
         "@llvm-project//mlir:ArithDialect",
+        "@llvm-project//mlir:BufferizationDialect",
+        "@llvm-project//mlir:BufferizationInterfaces",
+        "@llvm-project//mlir:BufferizationTransforms",
         "@llvm-project//mlir:FunctionInterfaces",
         "@llvm-project//mlir:IR",
         "@llvm-project//mlir:LinalgDialect",

compiler/src/iree/compiler/Codegen/Dialect/VectorExt/Transforms/BufferizationInterfaces.cpp

@@ -0,0 +1,78 @@
+// 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/Codegen/Dialect/VectorExt/Transforms/BufferizationInterfaces.h"
+#include "iree/compiler/Codegen/Dialect/VectorExt/IR/VectorExtDialect.h"
+
+#include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h"
+#include "mlir/Dialect/Bufferization/IR/Bufferization.h"
+#include "mlir/Dialect/Bufferization/Transforms/Transforms.h"
+#include "mlir/IR/BuiltinTypes.h"
+#include "mlir/IR/Value.h"
+
+namespace mlir::iree_compiler::IREE::VectorExt {
+
+using mlir::bufferization::AnalysisState;
+using mlir::bufferization::BufferizableOpInterface;
+using mlir::bufferization::BufferizationOptions;
+using mlir::bufferization::BufferizationState;
+using mlir::bufferization::BufferRelation;
+using mlir::bufferization::replaceOpWithNewBufferizedOp;
+
+namespace {
+
+struct TransferGatherOpInterface
+    : public BufferizableOpInterface::ExternalModel<
+          TransferGatherOpInterface, IREE::VectorExt::TransferGatherOp> {
+  bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
+                              const AnalysisState &state) const {
+    assert(isa<RankedTensorType>(opOperand.get().getType()) &&
+           "only tensor types expected");
+    return true;
+  }
+
+  bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,
+                               const AnalysisState &state) const {
+    assert(isa<RankedTensorType>(opOperand.get().getType()) &&
+           "only tensor types expected");
+    return false;
+  }
+
+  bufferization::AliasingValueList
+  getAliasingValues(Operation *op, OpOperand &opOperand,
+                    const AnalysisState &state) const {
+    return {};
+  }
+
+  LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
+                          const BufferizationOptions &options,
+                          BufferizationState &state) const {
+    auto gatherOp = cast<IREE::VectorExt::TransferGatherOp>(op);
+    assert(isa<TensorType>(gatherOp.getShapedType()) &&
+           "only tensor types expected");
+    FailureOr<Value> buffer =
+        getBuffer(rewriter, gatherOp.getBase(), options, state);
+    if (failed(buffer))
+      return failure();
+    replaceOpWithNewBufferizedOp<IREE::VectorExt::TransferGatherOp>(
+        rewriter, gatherOp, gatherOp.getVectorType(), *buffer,
+        gatherOp.getIndices(), gatherOp.getIndexVecs(), gatherOp.getIndexed(),
+        gatherOp.getIndexedMaps(), gatherOp.getPermutationMap(),
+        gatherOp.getPadding(), gatherOp.getMask(), gatherOp.getInBoundsAttr());
+    return success();
+  }
+};
+
+} // namespace
+
+void registerIREEVectorExtBufferizationInterfaces(DialectRegistry &registry) {
+  registry.addExtension(
+      +[](MLIRContext *context, IREEVectorExtDialect *dialect) {
+        TransferGatherOp::attachInterface<TransferGatherOpInterface>(*context);
+      });
+}
+
+} // namespace mlir::iree_compiler::IREE::VectorExt

compiler/src/iree/compiler/Codegen/Dialect/VectorExt/Transforms/BufferizationInterfaces.h

@@ -0,0 +1,19 @@
+// 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
+
+#ifndef IREE_COMPILER_CODEGEN_DIALECT_VECTOR_EXT_TRANSFORMS_BUFFERIZATIONINTERFACES_H_
+#define IREE_COMPILER_CODEGEN_DIALECT_VECTOR_EXT_TRANSFORMS_BUFFERIZATIONINTERFACES_H_
+
+#include "mlir/IR/Dialect.h"
+
+namespace mlir::iree_compiler::IREE::VectorExt {
+
+// Register all interfaces needed for bufferization.
+void registerIREEVectorExtBufferizationInterfaces(DialectRegistry &registry);
+
+} // namespace mlir::iree_compiler::IREE::VectorExt
+
+#endif // IREE_COMPILER_CODEGEN_DIALECT_VECTOR_EXT_TRANSFORMS_BUFFERIZATIONINTERFACES_H_

compiler/src/iree/compiler/Codegen/Dialect/VectorExt/Transforms/CMakeLists.txt

@@ -23,17 +23,21 @@ iree_cc_library(
   NAME
     VectorExtTransforms
   HDRS
+    "BufferizationInterfaces.h"
     "Passes.h"
     "Passes.h.inc"
     "Transforms.h"
   SRCS
+    "BufferizationInterfaces.cpp"
     "Passes.cpp"
     "VectorExtFoldUnitExtentDims.cpp"
     "VectorizeIREEVectorExtOps.cpp"
   DEPS
     ::PassesIncGen
     LLVMSupport
     MLIRArithDialect
+    MLIRBufferizationDialect
+    MLIRBufferizationTransforms
     MLIRFunctionInterfaces
     MLIRIR
     MLIRLinalgDialect

compiler/src/iree/compiler/Codegen/Interfaces/BUILD.bazel

@@ -91,6 +91,8 @@ iree_compiler_cc_library(
         "//compiler/src/iree/compiler/Codegen/Dialect/Codegen/IR:IREECodegenDialect",
         "//compiler/src/iree/compiler/Codegen/Dialect/GPU/IR:IREEGPUDialect",
         "//compiler/src/iree/compiler/Codegen/Dialect/GPU/Transforms:BufferizationInterfaces",
+        "//compiler/src/iree/compiler/Codegen/Dialect/VectorExt/IR:IREEVectorExtDialect",
+        "//compiler/src/iree/compiler/Codegen/Dialect/VectorExt/Transforms:VectorExtTransforms",
         "//compiler/src/iree/compiler/Codegen/Utils",
         "//compiler/src/iree/compiler/Dialect/HAL/IR",
         "//compiler/src/iree/compiler/Dialect/LinalgExt/IR",

compiler/src/iree/compiler/Codegen/Interfaces/BufferizationInterfaces.cpp

@@ -10,6 +10,8 @@
 #include "iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenOps.h"
 #include "iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUDialect.h"
 #include "iree/compiler/Codegen/Dialect/GPU/Transforms/BufferizationInterfaces.h"
+#include "iree/compiler/Codegen/Dialect/VectorExt/IR/VectorExtOps.h"
+#include "iree/compiler/Codegen/Dialect/VectorExt/Transforms/BufferizationInterfaces.h"
 #include "iree/compiler/Codegen/Utils/Utils.h"
 #include "iree/compiler/Dialect/HAL/IR/HALOps.h"
 #include "iree/compiler/Dialect/LinalgExt/IR/LinalgExtDialect.h"
@@ -675,6 +677,7 @@ void registerBufferizationInterfaces(DialectRegistry &registry) {
 
   // Register IREE operations.
   registerIREEGPUBufferizationInterfaces(registry);
+  IREE::VectorExt::registerIREEVectorExtBufferizationInterfaces(registry);
   registry.addExtension(
       +[](MLIRContext *ctx, IREE::TensorExt::IREETensorExtDialect *dialect) {
         // DispatchTensorLoadOp

compiler/src/iree/compiler/Codegen/Interfaces/CMakeLists.txt

@@ -79,6 +79,8 @@ iree_cc_library(
     iree::compiler::Codegen::Dialect::Codegen::IR::IREECodegenDialect
     iree::compiler::Codegen::Dialect::GPU::IR::IREEGPUDialect
     iree::compiler::Codegen::Dialect::GPU::Transforms::BufferizationInterfaces
+    iree::compiler::Codegen::Dialect::VectorExt::IR::IREEVectorExtDialect
+    iree::compiler::Codegen::Dialect::VectorExt::Transforms::VectorExtTransforms
     iree::compiler::Codegen::Utils
     iree::compiler::Dialect::HAL::IR
     iree::compiler::Dialect::LinalgExt::IR