[CPU][DT] Define boundary encodings and execution encodings.

The boundary operations (i.e., bindings, flow.tensor.load/store, etc)
has EncodingSolver attached. It defines the layout for inputs/outputs.

The encodings on the operations (e.g., compute ops) captures all the
original encoding fields and we do know that they will be executed on
the execution device, so we should be able to insert some ops which
bring it from layout_a to layout_b.

It adds the indexing maps back, we can revisit it later.

Signed-off-by: hanhanW <[email protected]>

compiler/src/iree/compiler/Codegen/Common/CPU/CPUMaterializeEncodings.cpp

@@ -433,7 +433,7 @@ materializeEncodingForTarget(RankedTensorType tensorType,
                              IREE::HAL::ExecutableTargetAttr targetAttr) {
   auto encoding =
       dyn_cast_or_null<IREE::Encoding::EncodingAttr>(tensorType.getEncoding());
-  if (!encoding) {
+  if (!encoding || !encoding.getUserIndexingMaps()) {
     return failure();
   }
 
@@ -617,7 +617,6 @@ struct CPUMaterializeDeviceEncodingPass
   void runOnOperation() override {
     auto funcOp = getOperation();
     bool unsupported = false;
-    SetVector<ArrayAttr> encodingSolvers;
     funcOp.walk([&](IREE::HAL::InterfaceBindingSubspanOp op) {
       auto resType = dyn_cast<IREE::Flow::DispatchTensorType>(op.getType());
       if (!resType) {
@@ -645,24 +644,8 @@ struct CPUMaterializeDeviceEncodingPass
         unsupported = true;
         return;
       }
-      encodingSolvers.insert(encoding.getTargets());
     });
 
-    if (!unsupported && encodingSolvers.size() != 1) {
-      LLVM_DEBUG(llvm::dbgs()
-                 << "unsupported: encodings are not from the same device\n");
-      unsupported = true;
-    }
-
-#if 1
-    if (!unsupported && !encodingSolvers.empty() &&
-        encodingSolvers[0].getValue().size() != 1) {
-      LLVM_DEBUG(llvm::dbgs()
-                 << "unsupported: only single target is supported atm\n");
-      unsupported = true;
-    }
-#endif
-
     if (unsupported) {
       LLVM_DEBUG(llvm::dbgs() << "drop encodings\n");
       OpPassManager pipeline(func::FuncOp::getOperationName());
@@ -674,22 +657,8 @@ struct CPUMaterializeDeviceEncodingPass
       return;
     }
 
-    auto solver =
-        *encodingSolvers[0]
-             .getAsRange<IREE::Encoding::EncodingSolverInterfaceAttr>()
-             .begin();
     auto executableTargetAttr = IREE::HAL::ExecutableTargetAttr::lookup(funcOp);
-    // TODO: The fake target attribute is not needed. This is only correct when
-    // the encoding target is as same as the execution device. These
-    // materialization implementations should be moved into interface methods,
-    // and we query all the needed information from the solver.
-    // Note: perhaps a new Codegen attribute interface, not
-    // IREE::Encoding::EncodingSolverInterfaceAttr.
-    auto fakeTargetAttr = IREE::HAL::ExecutableTargetAttr::get(
-        &getContext(), executableTargetAttr.getBackend(),
-        executableTargetAttr.getFormat(), solver.getConfig());
-
-    if (failed(materializeFuncOpEncodings(funcOp, fakeTargetAttr))) {
+    if (failed(materializeFuncOpEncodings(funcOp, executableTargetAttr))) {
       return signalPassFailure();
     }
   }

compiler/src/iree/compiler/Codegen/Common/EncodingUtils.h

@@ -29,6 +29,15 @@ struct MaterializeEncodingInfo {
 
   // The optional swizzle, see the comment on TileSwizzle. Only used on GPU.
   std::optional<TileSwizzle> swizzle;
+
+  bool operator==(const MaterializeEncodingInfo &other) const {
+    return innerDimsPos == other.innerDimsPos &&
+           innerTileSizes == other.innerTileSizes &&
+           outerDimsPerm == other.outerDimsPerm && swizzle == other.swizzle;
+  }
+  bool operator!=(const MaterializeEncodingInfo &other) const {
+    return !(*this == other);
+  }
 };
 
 DictionaryAttr serializeMaterializeEncodingInfo(MLIRContext *ctx,

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

@@ -12,6 +12,7 @@
 #include "iree/compiler/Codegen/Common/Passes.h"
 #include "iree/compiler/Codegen/Utils/Utils.h"
 #include "iree/compiler/Dialect/Encoding/IR/EncodingOps.h"
+#include "iree/compiler/Dialect/Encoding/IR/EncodingTypes.h"
 #include "iree/compiler/Dialect/Flow/IR/FlowOps.h"
 #include "iree/compiler/Dialect/HAL/IR/HALTypes.h"
 #include "iree/compiler/Dialect/Util/IR/UtilOps.h"
@@ -543,11 +544,34 @@ static FailureOr<SmallVector<OpFoldResult>> getPackedDimsForDispatchTensor(
     return failure();
   }
 
+  auto encoding =
+      dyn_cast<IREE::Encoding::EncodingAttr>(boundTensorType.getEncoding());
+  if (encoding.getTargets().empty()) {
+    // Unexpected in the prototype. All the targets should be reoslved in the
+    // boundary.
+    return failure();
+  }
+  auto target = dyn_cast<IREE::Encoding::EncodingSolverInterfaceAttr>(
+      encoding.getTargets().getValue()[0]);
+  if (!target) {
+    return failure();
+  }
+  std::optional<MaterializeEncodingInfo> srcLayout =
+      deserializeMaterializeEncodingInfo(target.getConfig());
+
   FailureOr<MaterializeEncodingInfo> encodingInfo =
       typeConverter.getEncodingInfo(boundTensorType);
   if (failed(encodingInfo)) {
     return failure();
   }
+
+  // TODO: Need to relayout it back to the original layout and re-apply the
+  // packing that the execution target device wants. Those operations should be
+  // derived by an interface method.
+  if (*encodingInfo != *srcLayout) {
+    return failure();
+  }
+
   if (typeConverter.getTransposeNarrowN() &&
       isNarrowNResult(IREE::Encoding::getEncodingAttr(boundTensorType))) {
     transposeInPlace(*encodingInfo);

compiler/src/iree/compiler/Codegen/Common/TileSwizzle.h

@@ -51,6 +51,11 @@ struct TileSwizzle {
 
     // The size of the dimension.
     int16_t size = 0;
+
+    bool operator==(const Dim &other) const {
+      return kind == other.kind && size == other.size;
+    }
+    bool operator!=(const Dim &other) const { return !(*this == other); }
   };
 
   using ExpandShapeDimVectorType = llvm::SmallVector<Dim, 4>;
@@ -68,6 +73,11 @@ struct TileSwizzle {
   // example, permutation[0] dictates which of the expanded dimensions becomes
   // the leading dimension of the layout.
   llvm::SmallVector<int64_t> permutation;
+
+  bool operator==(const TileSwizzle &other) const {
+    return expandShape == other.expandShape && permutation == other.permutation;
+  }
+  bool operator!=(const TileSwizzle &other) const { return !(*this == other); }
 };
 
 DictionaryAttr serializeTileSwizzle(MLIRContext *ctx, TileSwizzle swizzle);

compiler/src/iree/compiler/Dialect/Encoding/IR/EncodingAttrs.cpp

@@ -14,6 +14,7 @@
 #include "mlir/Dialect/Utils/StructuredOpsUtils.h"
 #include "mlir/IR/Attributes.h"
 #include "mlir/IR/Builders.h"
+#include "mlir/IR/BuiltinAttributes.h"
 #include "mlir/IR/TypeUtilities.h"
 #include "mlir/Interfaces/InferTypeOpInterface.h"
 #include "mlir/Support/LLVM.h"
@@ -114,8 +115,8 @@ EncodingAttr EncodingAttr::clone(AffineMap bcastMap) {
 EncodingAttr EncodingAttr::cloneWithTargets(ArrayRef<Attribute> targets) {
   MLIRContext *ctx = getContext();
   return get(ctx, getOperandIndex(), getOpType(), getElementTypes(),
-             getUserIndexingMaps(), getBcastMap(), getRoundDimsTo(),
-             ArrayAttr::get(ctx, targets));
+             getUserIndexingMaps(), /*bcast_map=*/AffineMapAttr(),
+             /*round_dims_to=*/DenseArrayAttr(), ArrayAttr::get(ctx, targets));
 }
 
 MatmulNarrowDim getMatmulNarrowDim(EncodingAttr encoding) {