[TensorExt][NFC] Remove dependency on Encoding through ExternalModel (#20612)

There is no strong dependency of `DispatchTensorType` on the Encoding
dialect as the encoding can be just any attribute and is optional, so I
don't think there's a need to have the TensorExt dialect explicitly
depend on the Encoding just for the interface. Decoupling them would
avoid potential circular dependency issues (I ran into that while
working on
#20160 (comment)),
remove the need for other (potential) users of the TensorExt dialect to
use the encoding and allow users to bring their own encoding interfaces
while avoiding having to bring in the Encoding dialect.

Signed-off-by: Jorn Tuyls <[email protected]>

Author: jtuyls

compiler/src/iree/compiler/Dialect/TensorExt/IR/BUILD.bazel

@@ -49,7 +49,6 @@ iree_compiler_cc_library(
     ],
     deps = [
         ":TensorExtOpsGen",
-        "//compiler/src/iree/compiler/Dialect/Encoding/IR",
         "//compiler/src/iree/compiler/Dialect/Util/IR",
         "@llvm-project//llvm:Support",
         "@llvm-project//mlir:ArithDialect",

compiler/src/iree/compiler/Dialect/TensorExt/IR/CMakeLists.txt

@@ -34,7 +34,6 @@ iree_cc_library(
     MLIRIR
     MLIRInferTypeOpInterface
     MLIRTensorDialect
-    iree::compiler::Dialect::Encoding::IR
     iree::compiler::Dialect::Util::IR
   PUBLIC
 )

compiler/src/iree/compiler/Dialect/TensorExt/IR/TensorExtTypes.cpp

@@ -108,13 +108,6 @@ bool DispatchTensorType::hasStaticShape(ArrayRef<int64_t> shape) const {
   return hasStaticShape() && getShape() == shape;
 }
 
-Type DispatchTensorType::getEncodingType() const { return getBoundType(); }
-
-Type DispatchTensorType::updateEncoding(Attribute encoding) const {
-  return DispatchTensorType::get(getAccess(), getShape(), getBoundElementType(),
-                                 encoding);
-}
-
 LogicalResult
 DispatchTensorType::verify(function_ref<InFlightDiagnostic()> emitError,
                            uint32_t access, Type boundType) {

compiler/src/iree/compiler/Dialect/TensorExt/IR/TensorExtTypes.h

@@ -7,7 +7,6 @@
 #ifndef IREE_COMPILER_DIALECT_TENSOREXT_IR_TENSOREXTTYPES_H_
 #define IREE_COMPILER_DIALECT_TENSOREXT_IR_TENSOREXTTYPES_H_
 
-#include "iree/compiler/Dialect/Encoding/IR/EncodingTypes.h"
 #include "iree/compiler/Dialect/TensorExt/IR/TensorExtDialect.h"
 #include "llvm/ADT/DenseMapInfo.h"
 #include "llvm/ADT/SmallVector.h"
@@ -40,8 +39,7 @@ enum class TensorAccess : uint32_t {
 // we can't extend it and reuse all of this.
 class DispatchTensorType
     : public Type::TypeBase<DispatchTensorType, Type,
-                            detail::DispatchTensorTypeStorage,
-                            IREE::Encoding::EncodingTypeInterface::Trait> {
+                            detail::DispatchTensorTypeStorage> {
 public:
   using ImplType = detail::DispatchTensorTypeStorage;
 
@@ -126,9 +124,6 @@ class DispatchTensorType
     }
     return llvm::cast<RankedTensorType>(boundType);
   }
-
-  Type getEncodingType() const;
-  Type updateEncoding(Attribute encoding) const;
 };
 
 void printType(DispatchTensorType &type, DialectAsmPrinter &p);

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

@@ -6,12 +6,17 @@
 
 #include "iree/compiler/ExternalInterfaces/TensorExtExternalModels.h"
 
+#include "iree/compiler/Dialect/Encoding/IR/EncodingTypes.h"
 #include "iree/compiler/Dialect/TensorExt/IR/TensorExtOps.h"
 #include "mlir/Interfaces/ValueBoundsOpInterface.h"
 
 namespace mlir::iree_compiler {
 namespace {
 
+//===----------------------------------------------------------------------===//
+// Op Interfaces
+//===----------------------------------------------------------------------===//
+
 struct DispatchTensorLoadOpInterface
     : public ValueBoundsOpInterface::ExternalModel<
           DispatchTensorLoadOpInterface,
@@ -36,6 +41,27 @@ struct WorkloadOrdinalOpInterface
   }
 };
 
+//===----------------------------------------------------------------------===//
+// Type Interfaces
+//===----------------------------------------------------------------------===//
+
+struct EncodingTypeExternalModel
+    : public IREE::Encoding::EncodingTypeInterface::ExternalModel<
+          EncodingTypeExternalModel, IREE::TensorExt::DispatchTensorType> {
+
+  Type getEncodingType(Type type) const {
+    auto dispatchTensorType = cast<IREE::TensorExt::DispatchTensorType>(type);
+    return dispatchTensorType.getBoundType();
+  }
+
+  Type updateEncoding(Type type, Attribute encoding) const {
+    auto dispatchTensorType = cast<IREE::TensorExt::DispatchTensorType>(type);
+    return IREE::TensorExt::DispatchTensorType::get(
+        dispatchTensorType.getAccess(), dispatchTensorType.getShape(),
+        dispatchTensorType.getBoundElementType(), encoding);
+  }
+};
+
 } // namespace
 
 void registerTensorExtExternalModels(DialectRegistry &registry) {
@@ -45,6 +71,8 @@ void registerTensorExtExternalModels(DialectRegistry &registry) {
             DispatchTensorLoadOpInterface>(*ctx);
         IREE::TensorExt::DispatchWorkloadOrdinalOp::attachInterface<
             WorkloadOrdinalOpInterface>(*ctx);
+        IREE::TensorExt::DispatchTensorType::attachInterface<
+            EncodingTypeExternalModel>(*ctx);
       });
 }