[mlir][vector] Implement subset op interface for xfer ops

Author: matthias-springer

mlir/include/mlir/Dialect/Vector/Transforms/SubsetOpInterfaceImpl.h

@@ -0,0 +1,20 @@
+//===- SubsetOpInterfaceImpl.h - Tensor subsets -----------------*- C++ -*-===//
+//
+// Part of the LLVM Project, 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 MLIR_DIALECT_VECTOR_SUBSETOPINTERFACEIMPL_H
+#define MLIR_DIALECT_VECTOR_SUBSETOPINTERFACEIMPL_H
+
+namespace mlir {
+class DialectRegistry;
+
+namespace vector {
+void registerSubsetOpInterfaceExternalModels(DialectRegistry &registry);
+} // namespace vector
+} // namespace mlir
+
+#endif // MLIR_DIALECT_VECTOR_SUBSETOPINTERFACEIMPL_H

mlir/include/mlir/InitAllDialects.h

@@ -85,6 +85,7 @@
 #include "mlir/Dialect/UB/IR/UBOps.h"
 #include "mlir/Dialect/Vector/IR/VectorOps.h"
 #include "mlir/Dialect/Vector/Transforms/BufferizableOpInterfaceImpl.h"
+#include "mlir/Dialect/Vector/Transforms/SubsetOpInterfaceImpl.h"
 #include "mlir/Dialect/X86Vector/X86VectorDialect.h"
 #include "mlir/IR/Dialect.h"
 #include "mlir/Interfaces/CastInterfaces.h"
@@ -171,6 +172,7 @@ inline void registerAllDialects(DialectRegistry &registry) {
  tensor::registerTilingInterfaceExternalModels(registry);
  tensor::registerValueBoundsOpInterfaceExternalModels(registry);
  vector::registerBufferizableOpInterfaceExternalModels(registry);
+ vector::registerSubsetOpInterfaceExternalModels(registry);
  NVVM::registerNVVMTargetInterfaceExternalModels(registry);
  ROCDL::registerROCDLTargetInterfaceExternalModels(registry);
 }

mlir/lib/Dialect/Vector/Transforms/CMakeLists.txt

@@ -10,6 +10,7 @@ add_mlir_dialect_library(MLIRVectorTransforms
  LowerVectorShapeCast.cpp
  LowerVectorTransfer.cpp
  LowerVectorTranspose.cpp
+ SubsetOpInterfaceImpl.cpp
  VectorDistribute.cpp
  VectorDropLeadUnitDim.cpp
  VectorEmulateNarrowType.cpp
@@ -40,6 +41,7 @@ add_mlir_dialect_library(MLIRVectorTransforms
  MLIRMemRefUtils
  MLIRSCFDialect
  MLIRSideEffectInterfaces
+ MLIRSubsetOpInterface
  MLIRTensorDialect
  MLIRTransforms
  MLIRVectorDialect

mlir/lib/Dialect/Vector/Transforms/SubsetOpInterfaceImpl.cpp

@@ -0,0 +1,82 @@
+//===- SubsetOpInterfaceImpl.cpp - Tensor subsets -------------------------===//
+//
+// Part of the LLVM Project, 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 "mlir/Dialect/Vector/Transforms/SubsetOpInterfaceImpl.h"
+
+#include "mlir/Dialect/Vector/IR/VectorOps.h"
+#include "mlir/Interfaces/SubsetOpInterface.h"
+
+using namespace mlir;
+using namespace mlir::vector;
+
+namespace {
+
+template <typename OpTy>
+struct XferOpSubsetOpInterface
+ : public SubsetOpInterface::ExternalModel<XferOpSubsetOpInterface<OpTy>,
+ OpTy> {
+ FailureOr<HyperrectangularSlice>
+ getAccessedHyperrectangularSlice(Operation *op) const {
+ auto xferOp = cast<OpTy>(op);
+ Builder b(xferOp->getContext());
+ SmallVector<OpFoldResult> offsets = llvm::map_to_vector(
+ xferOp.getIndices(), [](Value v) -> OpFoldResult { return v; });
+ SmallVector<OpFoldResult> sizes = llvm::map_to_vector(
+ xferOp.getTransferChunkAccessed(),
+ [&](int64_t sz) -> OpFoldResult { return b.getIndexAttr(sz); });
+ return HyperrectangularSlice(offsets, sizes);
+ }
+};
+
+struct TransferReadOpSubsetExtractionOpInterface
+ : public SubsetExtractionOpInterface::ExternalModel<
+ TransferReadOpSubsetExtractionOpInterface, vector::TransferReadOp> {
+ OpOperand &getSourceOperand(Operation *op) const {
+ return cast<vector::TransferReadOp>(op).getSourceMutable();
+ }
+};
+
+struct TransferWriteOpSubsetInsertionOpInterface
+ : public SubsetInsertionOpInterface::ExternalModel<
+ TransferWriteOpSubsetInsertionOpInterface, vector::TransferWriteOp> {
+ OpOperand &getSourceOperand(Operation *op) const {
+ return cast<vector::TransferWriteOp>(op).getVectorMutable();
+ }
+
+ OpOperand &getDestinationOperand(Operation *op) const {
+ return cast<vector::TransferWriteOp>(op).getSourceMutable();
+ }
+
+ Value buildSubsetExtraction(Operation *op, OpBuilder &builder,
+ Location loc) const {
+ // TODO: Implement when needed.
+ return Value();
+ }
+
+ SmallVector<Value>
+ getValuesNeededToBuildSubsetExtraction(Operation *op) const {
+ // TODO: Implement when needed.
+ return {};
+ }
+};
+
+} // namespace
+
+void mlir::vector::registerSubsetOpInterfaceExternalModels(
+ DialectRegistry &registry) {
+ registry.addExtension(+[](MLIRContext *ctx, vector::VectorDialect *dialect) {
+ TransferReadOp::attachInterface<XferOpSubsetOpInterface<TransferReadOp>>(
+ *ctx);
+ TransferReadOp::attachInterface<TransferReadOpSubsetExtractionOpInterface>(
+ *ctx);
+ TransferWriteOp::attachInterface<XferOpSubsetOpInterface<TransferWriteOp>>(
+ *ctx);
+ TransferWriteOp::attachInterface<TransferWriteOpSubsetInsertionOpInterface>(
+ *ctx);
+ });
+}