[Integrate] Update bufferization related codes for upstream custom types support. (iree-org#21250)

Author: hanhanW

compiler/src/iree/compiler/Codegen/Common/IREEComprehensiveBufferizePass.cpp

@@ -78,7 +78,12 @@ class EliminateEmptyTensorsPass final
     : public impl::EliminateEmptyTensorsPassBase<EliminateEmptyTensorsPass> {
 public:
   void getDependentDialects(DialectRegistry &registry) const override {
-    registry.insert<tensor::TensorDialect>();
+    // BufferizationDialect is needed for using type interfaces, like
+    // TensorLikeType. Because the builtin types, e.g., RankedTensorType, etc.,
+    // implement the type interface in
+    // bufferization::BufferizationDialect::initialize().
+    registry
+        .insert<bufferization::BufferizationDialect, tensor::TensorDialect>();
   }
 
   void runOnOperation() override;

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

@@ -14,6 +14,7 @@
 #include "mlir/Dialect/Bufferization/IR/Bufferization.h"
 #include "mlir/Dialect/Bufferization/Transforms/Transforms.h"
 #include "mlir/Dialect/GPU/IR/GPUDialect.h"
+#include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/Value.h"
 
@@ -108,7 +109,7 @@ struct BarrierRegionOpBufferizationInterface
                 SmallVector<Value> &invocationStack) const {
     auto barrierOp = cast<IREE::GPU::BarrierRegionOp>(op);
 
-    FailureOr<BaseMemRefType> memrefType = failure();
+    FailureOr<mlir::bufferization::BufferLikeType> memrefType = failure();
     if (auto opResult = dyn_cast<OpResult>(value)) {
       int64_t resultNum = opResult.getResultNumber();
       memrefType = bufferization::getBufferType(
@@ -121,7 +122,7 @@ struct BarrierRegionOpBufferizationInterface
     }
     if (failed(memrefType))
       return failure();
-    return memrefType;
+    return cast<BaseMemRefType>(*memrefType);
   }
 
   LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
@@ -146,10 +147,13 @@ struct BarrierRegionOpBufferizationInterface
         tensorizedOperands.push_back(replacement);
         continue;
       }
-      tensorizedOperands.push_back(rewriter
-                                       .create<bufferization::ToTensorOp>(
-                                           replacement.getLoc(), replacement)
-                                       .getResult());
+      tensorizedOperands.push_back(
+          rewriter
+              .create<bufferization::ToTensorOp>(
+                  replacement.getLoc(),
+                  memref::getTensorTypeFromMemRefType(replacement.getType()),
+                  replacement)
+              .getResult());
     }
 
     rewriter.setInsertionPoint(barrierOp);
@@ -205,7 +209,7 @@ struct ValueBarrierOpBufferizationInterface
         state, invocationStack);
     if (failed(srcMemrefType))
       return failure();
-    return srcMemrefType;
+    return cast<BaseMemRefType>(*srcMemrefType);
   }
 
   LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
@@ -341,8 +345,9 @@ struct BufferResourceCastOpBufferizationInterface
     if (failed(srcMemrefType))
       return failure();
 
-    if (!hasStorageBufferMemSpace(srcMemrefType.value())) {
-      return srcMemrefType;
+    auto baseMemrefType = cast<BaseMemRefType>(srcMemrefType.value());
+    if (!hasStorageBufferMemSpace(baseMemrefType)) {
+      return baseMemrefType;
     }
 
     auto rankedSrcType = cast<MemRefType>(srcMemrefType.value());

compiler/src/iree/compiler/ExternalInterfaces/BUILD.bazel

@@ -44,6 +44,7 @@ iree_compiler_cc_library(
         "//compiler/src/iree/compiler/Dialect/Util/IR",
         "@llvm-project//llvm:Support",
         "@llvm-project//mlir:ArithDialect",
+        "@llvm-project//mlir:BufferizationInterfaces",
         "@llvm-project//mlir:GPUDialect",
         "@llvm-project//mlir:IR",
         "@llvm-project//mlir:LinalgDialect",

compiler/src/iree/compiler/ExternalInterfaces/TensorExtExternalModels.cpp

@@ -8,6 +8,9 @@
 
 #include "iree/compiler/Dialect/Encoding/IR/EncodingTypes.h"
 #include "iree/compiler/Dialect/TensorExt/IR/TensorExtOps.h"
+#include "iree/compiler/Dialect/TensorExt/IR/TensorExtTypes.h"
+#include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h"
+#include "mlir/Dialect/Bufferization/IR/BufferizationTypeInterfaces.h"
 #include "mlir/Interfaces/ValueBoundsOpInterface.h"
 
 namespace mlir::iree_compiler {
@@ -62,6 +65,39 @@ struct EncodingTypeExternalModel
   }
 };
 
+struct TensorLikeTypeExternalModel
+    : bufferization::TensorLikeType::ExternalModel<
+          TensorLikeTypeExternalModel, IREE::TensorExt::DispatchTensorType> {
+  FailureOr<bufferization::BufferLikeType> getBufferType(
+      Type type, const bufferization::BufferizationOptions &options,
+      llvm::function_ref<mlir::InFlightDiagnostic()> emitError) const {
+    auto dispatchTensorType = cast<IREE::TensorExt::DispatchTensorType>(type);
+    auto tensorType = cast<TensorType>(dispatchTensorType.asRankedTensorType());
+    auto memSpace = options.defaultMemorySpaceFn(tensorType);
+    if (!memSpace.has_value()) {
+      return emitError() << "could not infer memory space";
+    }
+    return cast<bufferization::BufferLikeType>(
+        getMemRefType(tensorType, options, /*layout=*/{}, *memSpace));
+  }
+
+  LogicalResult verifyCompatibleBufferType(
+      Type type, bufferization::BufferLikeType bufferType,
+      llvm::function_ref<mlir::InFlightDiagnostic()> emitError) const {
+    auto dispatchTensorType = cast<IREE::TensorExt::DispatchTensorType>(type);
+    assert(isa<BaseMemRefType>(bufferType) && "expected memref type");
+    auto memrefType = cast<ShapedType>(bufferType);
+    if (dispatchTensorType.getShape() != memrefType.getShape()) {
+      return emitError() << "shapes do not match";
+    }
+    if (dispatchTensorType.getBoundElementType() !=
+        memrefType.getElementType()) {
+      return emitError() << "element types do not match";
+    }
+    return success();
+  }
+};
+
 } // namespace
 
 void registerTensorExtExternalModels(DialectRegistry &registry) {
@@ -72,7 +108,7 @@ void registerTensorExtExternalModels(DialectRegistry &registry) {
         IREE::TensorExt::DispatchWorkloadOrdinalOp::attachInterface<
             WorkloadOrdinalOpInterface>(*ctx);
         IREE::TensorExt::DispatchTensorType::attachInterface<
-            EncodingTypeExternalModel>(*ctx);
+            EncodingTypeExternalModel, TensorLikeTypeExternalModel>(*ctx);
       });
 }
 

compiler/src/iree/compiler/Preprocessing/Common/Passes.td

@@ -159,6 +159,9 @@ def TransposeMatmulPass : Pass<"iree-preprocessing-transpose-matmul-pass"> {
            )}]>
   ];
   let dependentDialects = [
+    // TODO(hanchung): Remove the dep after switching upstream patterns to not
+    // use bufferization::hasTensorSemantics method.
+    "mlir::bufferization::BufferizationDialect",
     "mlir::linalg::LinalgDialect",
   ];
 }