[DT] Switch MaterializeEncodingIntoNop to use EncodingNopLayoutAttr.

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

@@ -8,6 +8,7 @@
 #include "iree/compiler/Codegen/Common/EncodingUtils.h"
 #include "iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenDialect.h"
 #include "iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenOps.h"
+#include "iree/compiler/Codegen/Dialect/Codegen/Utils/Utils.h"
 #include "iree/compiler/Codegen/Utils/Utils.h"
 #include "iree/compiler/Dialect/Encoding/IR/EncodingOps.h"
 #include "iree/compiler/Dialect/HAL/Analysis/DeviceAnalysis.h"
@@ -481,7 +482,8 @@ materializeFuncOpEncodings(FunctionOpInterface funcOp,
   // 3. Heuristics for cache-friendly dispatch tiling can get complex on CPU,
   //    so it is nice that they have fewer narrow cases to consider.
   MaterializeEncodingTypeConverter typeConverter(
-      materializeEncodingForTarget, targetAttr, /*transposeNarrowN=*/true);
+      materializeEncodingForTarget, targetAttr, /*transposeNarrowN=*/true,
+      /*layoutAttr=*/{});
   MaterializeEncodingConversionTarget target(*ctx);
   auto materializeEncodingValueFn = getMaterializeEncodingValueFn(targetAttr);
   populateMaterializeEncodingIntoPackUnPackPatterns(

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

@@ -89,9 +89,10 @@ static RankedTensorType transposeIfNarrowNResult(RankedTensorType tensorType) {
 
 MaterializeEncodingTypeConverter::MaterializeEncodingTypeConverter(
     MaterializeEncodingFn materializeEncodingFn,
-    IREE::HAL::ExecutableTargetAttr targetAttr, bool transposeNarrowN)
+    IREE::HAL::ExecutableTargetAttr targetAttr, bool transposeNarrowN,
+    IREE::Codegen::LayoutAttrInterface layoutAttr)
     : materializeEncodingFn(materializeEncodingFn), targetAttr(targetAttr),
-      transposeNarrowN(transposeNarrowN) {
+      transposeNarrowN(transposeNarrowN), layoutAttr(layoutAttr) {
   addConversion([](IntegerType intType) { return intType; });
   addConversion([](IndexType indexType) { return indexType; });
   addConversion([](FloatType floatType) { return floatType; });

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

@@ -7,7 +7,8 @@
 #ifndef IREE_COMPILER_SRC_IREE_COMPILER_CODEGEN_COMMON_ENCODINGUTILS_H_
 #define IREE_COMPILER_SRC_IREE_COMPILER_CODEGEN_COMMON_ENCODINGUTILS_H_
 
-#include "iree/compiler/Codegen/Dialect/Codegen/Utils/Utils.h"
+#include "iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenInterfaces.h"
+#include "iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenTypes.h"
 #include "iree/compiler/Dialect/Encoding/IR/EncodingOps.h"
 #include "iree/compiler/Dialect/HAL/IR/HALTypes.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
@@ -34,9 +35,13 @@ using MaterializeEncodingValueFn =
 /// TypeConverter to use for materializing the encoding.
 class MaterializeEncodingTypeConverter : public TypeConverter {
 public:
-  MaterializeEncodingTypeConverter(MaterializeEncodingFn fn,
-                                   IREE::HAL::ExecutableTargetAttr targetAttr,
-                                   bool transposeNarrowN);
+  MaterializeEncodingTypeConverter(
+      MaterializeEncodingFn fn, IREE::HAL::ExecutableTargetAttr targetAttr,
+      bool transposeNarrowN, IREE::Codegen::LayoutAttrInterface layoutAttr);
+
+  const IREE::Codegen::LayoutAttrInterface &getLayoutAttr() const {
+    return layoutAttr;
+  }
 
   const MaterializeEncodingFn &getMaterializeEncodingFn() const {
     return materializeEncodingFn;
@@ -46,6 +51,9 @@ class MaterializeEncodingTypeConverter : public TypeConverter {
 
   FailureOr<IREE::Codegen::MaterializeEncodingInfo>
   getEncodingInfo(RankedTensorType type) const {
+    if (layoutAttr) {
+      return layoutAttr.getEncodingInfo(type);
+    }
     return materializeEncodingFn(type, targetAttr);
   }
 
@@ -55,6 +63,13 @@ class MaterializeEncodingTypeConverter : public TypeConverter {
   const MaterializeEncodingFn materializeEncodingFn;
   const IREE::HAL::ExecutableTargetAttr targetAttr;
   bool transposeNarrowN = false;
+  // The `layoutAttr` implements the logic of encoding materialization. It has
+  // a higher priority when it is present.
+  // TODO(hanchung): Move the logic that takes `targetAttr` and
+  // `transposeNarrowN` into account to their own attribute implementation. It
+  // is in a transition state, so we have two paths atm. We're incrementally
+  // moving the logic to attributes.
+  const IREE::Codegen::LayoutAttrInterface layoutAttr;
 };
 
 /// Conversion target to use for for materializing the encoding.
@@ -86,17 +101,15 @@ class OpMaterializeEncodingPattern : public OpConversionPattern<OpTy> {
 RankedTensorType dropEncoding(RankedTensorType type);
 
 /// Utility method to convert from `set_encoding` op to `pack` operation.
-/// For now this takes a `paddingValue` as input. The source is also taken
-/// as input so that these could be used with `OpConversionPatterns`.
-FailureOr<tensor::PackOp> lowerSetEncodingOpToPackOp(
+/// NOTE: `source` could be returned when packing is not needed.
+FailureOr<Value> lowerSetEncodingOpToPackOp(
     RewriterBase &rewriter, IREE::Encoding::SetEncodingOp encodingOp,
     Value source, const MaterializeEncodingTypeConverter &typeConverter,
     MaterializeEncodingValueFn materializeEncodingValueFn);
 
 /// Utility method to convert from `unset_encoding` op to `unpack` operation.
-/// The source is taken as input so that these could be used with
-/// `OpConversionPatterns`.
-FailureOr<tensor::UnPackOp> lowerUnsetEncodingToUnpackOp(
+/// NOTE: `packedValue` could be returned when unpacking is not needed.
+FailureOr<Value> lowerUnsetEncodingToUnpackOp(
     RewriterBase &rewriter, IREE::Encoding::UnsetEncodingOp encodingOp,
     Value packedValue, const MaterializeEncodingTypeConverter &typeConverter,
     MaterializeEncodingValueFn materializeEncodingValueFn);

compiler/src/iree/compiler/Codegen/Common/GPU/GPUMaterializeEncoding.cpp

@@ -315,10 +315,10 @@ struct GPUSetEncodingOpLoweringConversion
                   ConversionPatternRewriter &rewriter) const override {
     auto converter = static_cast<const MaterializeEncodingTypeConverter *>(
         getTypeConverter());
-    auto packOp = lowerSetEncodingOpToPackOp(rewriter, encodingOp,
-                                             adaptor.getSource(), *converter,
-                                             this->materializeEncodingValueFn);
-    if (failed(packOp)) {
+    auto packedValue = lowerSetEncodingOpToPackOp(
+        rewriter, encodingOp, adaptor.getSource(), *converter,
+        this->materializeEncodingValueFn);
+    if (failed(packedValue)) {
       Type targetType =
           getTypeConverter()->convertType(encodingOp.getResultType());
       Value result = rewriter.createOrFold<tensor::CastOp>(
@@ -334,7 +334,7 @@ struct GPUSetEncodingOpLoweringConversion
                                          "unhandled result encoding");
     }
     if (!maybeEncodingInfo->swizzle) {
-      rewriter.replaceOp(encodingOp, packOp->getResult());
+      rewriter.replaceOp(encodingOp, packedValue.value());
       return success();
     }
 
@@ -343,7 +343,9 @@ struct GPUSetEncodingOpLoweringConversion
     // Create expand_shape op to tile the innermost two dimensions.
     int origRank = encodingOp.getSourceType().getRank();
     SmallVector<int64_t> expandShapeShape(
-        packOp->getDestType().getShape().take_front(origRank));
+        cast<ShapedType>(packedValue->getType())
+            .getShape()
+            .take_front(origRank));
     expandShapeShape.append(
         getExpandedTileShape(maybeEncodingInfo->swizzle->expandShape));
     RankedTensorType expandShapeType =
@@ -352,7 +354,7 @@ struct GPUSetEncodingOpLoweringConversion
     SmallVector<ReassociationIndices> reassociation = getReassociationIndices(
         origRank, maybeEncodingInfo->swizzle->expandShape);
     auto expandShapeOp = rewriter.create<tensor::ExpandShapeOp>(
-        loc, expandShapeType, packOp->getResult(), reassociation);
+        loc, expandShapeType, packedValue.value(), reassociation);
 
     SmallVector<int64_t> transposePerm =
         llvm::to_vector(llvm::seq<int64_t>(0, origRank));
@@ -433,18 +435,18 @@ struct GPUUnsetEncodingOpLoweringConversion
           loc, unpackSrcType, transposeOp->getResult(0), reassociation);
     }
 
-    auto unPackOp = lowerUnsetEncodingToUnpackOp(
+    auto unpackedValue = lowerUnsetEncodingToUnpackOp(
         rewriter, unsetEncodingOp, unpackSrc, *converter,
         this->materializeEncodingValueFn);
-    if (failed(unPackOp)) {
+    if (failed(unpackedValue)) {
       Type targetType =
           getTypeConverter()->convertType(unsetEncodingOp.getResultType());
       Value result = rewriter.createOrFold<tensor::CastOp>(loc, targetType,
                                                            adaptor.getSource());
       rewriter.replaceOp(unsetEncodingOp, result);
       return success();
     }
-    rewriter.replaceOp(unsetEncodingOp, unPackOp->getResult());
+    rewriter.replaceOp(unsetEncodingOp, unpackedValue.value());
     return success();
   }
 };
@@ -559,7 +561,8 @@ materializeFuncOpEncodings(FunctionOpInterface funcOp,
     // 3. Heuristics for cache-friendly dispatch tiling are internal to the GPU
     //    runtime, so we don't need a simplification at that level either.
     MaterializeEncodingTypeConverter typeConverter(
-        materializeEncodingForTarget, targetAttr, /*transposeNarrowN=*/false);
+        materializeEncodingForTarget, targetAttr, /*transposeNarrowN=*/false,
+        /*layoutAttr=*/{});
     MaterializeEncodingConversionTarget target(*ctx);
     MaterializeEncodingValueFn materializeEncodingValueFn =
         [](RankedTensorType, OpBuilder,

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

@@ -7,6 +7,8 @@
 #include "iree/compiler/Codegen/Common/EncodingUtils.h"
 #include "iree/compiler/Codegen/Common/PassUtils.h"
 #include "iree/compiler/Codegen/Common/Passes.h"
+#include "iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenAttrs.h"
+#include "iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenDialect.h"
 #include "iree/compiler/Dialect/Encoding/IR/EncodingOps.h"
 #include "iree/compiler/Dialect/HAL/IR/HALTypes.h"
 #include "mlir/Dialect/MemRef/Transforms/Transforms.h"
@@ -28,7 +30,8 @@ namespace {
 struct MaterializeEncodingIntoNopPass final
     : impl::MaterializeEncodingIntoNopPassBase<MaterializeEncodingIntoNopPass> {
   void getDependentDialects(DialectRegistry &registry) const override {
-    registry.insert<linalg::LinalgDialect, tensor::TensorDialect>();
+    registry.insert<linalg::LinalgDialect, tensor::TensorDialect,
+                    IREE::Codegen::IREECodegenDialect>();
   }
 
   void runOnOperation() override {
@@ -47,7 +50,8 @@ struct MaterializeEncodingIntoNopPass final
     RewritePatternSet materializeEncodingPattern(context);
     MaterializeEncodingTypeConverter typeConverter(
         materializeEncodingFn, IREE::HAL::ExecutableTargetAttr(),
-        /*transposeNarrowN=*/false);
+        /*transposeNarrowN=*/false,
+        IREE::Codegen::EncodingNopLayoutAttr::get(context));
     MaterializeEncodingConversionTarget target(*context);
     populateMaterializeEncodingIntoPackUnPackPatterns(
         materializeEncodingPattern, typeConverter, materializeEncodingValueFn);

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

@@ -10,12 +10,14 @@
 
 #include "iree/compiler/Codegen/Common/EncodingUtils.h"
 #include "iree/compiler/Codegen/Common/Passes.h"
+#include "iree/compiler/Codegen/Dialect/Codegen/Utils/Utils.h"
 #include "iree/compiler/Codegen/Utils/Utils.h"
 #include "iree/compiler/Dialect/Encoding/IR/EncodingOps.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"
 #include "llvm/ADT/SmallVectorExtras.h"
+#include "llvm/Support/LogicalResult.h"
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/Linalg/IR/Linalg.h"
@@ -119,12 +121,7 @@ static void transposeInPlace(MaterializeEncodingInfo &info) {
 // to `pack` and `unpack` operations respectively.
 //===---------------------------------------------------------------------===//
 
-/// TODO(hanchung): Move the implementation to EncodingUtils.cpp. It is not
-/// moved because it needs some cleanup for this file. E.g., `getPaddingValue`
-/// is no longer needed. Ideally we should move CPU specific patterns (e.g.,
-/// lowerContractionOpWithEncoding, etc) to the CPUMaterializeEncoding file;
-/// move general patterns to EncodingUtils, and retire this file.
-FailureOr<tensor::PackOp> lowerSetEncodingOpToPackOp(
+FailureOr<Value> lowerSetEncodingOpToPackOp(
     RewriterBase &rewriter, IREE::Encoding::SetEncodingOp encodingOp,
     Value source, const MaterializeEncodingTypeConverter &typeConverter,
     MaterializeEncodingValueFn materializeEncodingValueFn) {
@@ -135,6 +132,11 @@ FailureOr<tensor::PackOp> lowerSetEncodingOpToPackOp(
     return rewriter.notifyMatchFailure(encodingOp, "unhandled result encoding");
   }
 
+  // Shortcut to avoid creating new operations.
+  if (IREE::Codegen::isIdentityLayout(encodingInfo.value())) {
+    return source;
+  }
+
   auto encoding = IREE::Encoding::getEncodingAttr(resultType);
   if (!encoding) {
     return failure();
@@ -160,14 +162,14 @@ FailureOr<tensor::PackOp> lowerSetEncodingOpToPackOp(
       encodingInfo->outerDimsPerm);
   auto emptyOp = rewriter.create<tensor::EmptyOp>(loc, resultDims,
                                                   resultType.getElementType());
-  return rewriter.create<tensor::PackOp>(
-      loc, source, emptyOp, encodingInfo->innerDimsPos, *innerTileSizesOfr,
-      paddingValue, encodingInfo->outerDimsPerm);
+  return rewriter
+      .create<tensor::PackOp>(loc, source, emptyOp, encodingInfo->innerDimsPos,
+                              *innerTileSizesOfr, paddingValue,
+                              encodingInfo->outerDimsPerm)
+      .getResult();
 }
 
-/// TODO(hanchung): Move the implementation to EncodingUtils.cpp. See the reason
-/// in the implementation comment of lowerSetEncodingToPackOp method.
-FailureOr<tensor::UnPackOp> lowerUnsetEncodingToUnpackOp(
+FailureOr<Value> lowerUnsetEncodingToUnpackOp(
     RewriterBase &rewriter, IREE::Encoding::UnsetEncodingOp encodingOp,
     Value packedValue, const MaterializeEncodingTypeConverter &typeConverter,
     MaterializeEncodingValueFn materializeEncodingValueFn) {
@@ -177,6 +179,12 @@ FailureOr<tensor::UnPackOp> lowerUnsetEncodingToUnpackOp(
   if (failed(encodingInfo)) {
     return rewriter.notifyMatchFailure(encodingOp, "unhandled source encoding");
   }
+
+  // Shortcut to avoid creating new operations.
+  if (IREE::Codegen::isIdentityLayout(encodingInfo.value())) {
+    return packedValue;
+  }
+
   auto encoding = IREE::Encoding::getEncodingAttr(sourceType);
   if (typeConverter.getTransposeNarrowN() && isNarrowNResult(encoding)) {
     transposeInPlace(*encodingInfo);
@@ -194,9 +202,11 @@ FailureOr<tensor::UnPackOp> lowerUnsetEncodingToUnpackOp(
     return rewriter.notifyMatchFailure(
         encodingOp, "failed to generate runtime tile size query");
   }
-  return rewriter.create<tensor::UnPackOp>(
-      loc, packedValue, emptyOp, encodingInfo->innerDimsPos, *innerTileSizesOfr,
-      encodingInfo->outerDimsPerm);
+  return rewriter
+      .create<tensor::UnPackOp>(loc, packedValue, emptyOp,
+                                encodingInfo->innerDimsPos, *innerTileSizesOfr,
+                                encodingInfo->outerDimsPerm)
+      .getResult();
 }
 
 /// Utility method to convert `tensor.empty` with encoding to a `tensor.empty`
@@ -609,7 +619,7 @@ struct SetEncodingOpToPackOpConversion
       rewriter.replaceOp(encodingOp, result);
       return success();
     }
-    rewriter.replaceOp(encodingOp, packOp->getResult());
+    rewriter.replaceOp(encodingOp, packOp.value());
     return success();
   }
 };
@@ -625,18 +635,18 @@ struct UnsetEncodingOpToUnPackOpConversion
                   ConversionPatternRewriter &rewriter) const override {
     auto converter = static_cast<const MaterializeEncodingTypeConverter *>(
         this->getTypeConverter());
-    auto unpackOp = lowerUnsetEncodingToUnpackOp(
+    auto unpackedValue = lowerUnsetEncodingToUnpackOp(
         rewriter, encodingOp, adaptor.getSource(), *converter,
         this->materializeEncodingValueFn);
-    if (failed(unpackOp)) {
+    if (failed(unpackedValue)) {
       Type targetType =
           getTypeConverter()->convertType(encodingOp.getResultType());
       Value result = rewriter.createOrFold<tensor::CastOp>(
           encodingOp.getLoc(), targetType, adaptor.getSource());
       rewriter.replaceOp(encodingOp, result);
       return success();
     }
-    rewriter.replaceOp(encodingOp, unpackOp->getResult());
+    rewriter.replaceOp(encodingOp, unpackedValue.value());
     return success();
   }
 };
@@ -734,6 +744,18 @@ class MaterializeContractionOp
 
     auto converter = static_cast<const MaterializeEncodingTypeConverter *>(
         this->getTypeConverter());
+
+    if (auto layoutAttr = converter->getLayoutAttr()) {
+      SmallVector<Type> convertedResTypes;
+      for (auto init : op.getDpsInits()) {
+        convertedResTypes.push_back(converter->convertType(init.getType()));
+      }
+      Operation *newOp =
+          layoutAttr.lowerOp(rewriter, op, convertedResTypes, operands);
+      rewriter.replaceOp(op, newOp->getResults());
+      return success();
+    }
+
     // TODO(hanchung): This is a transition state for moving the implementation
     // details to backend attributes. We won't need the function type argument
     // after all the backends that support encodings implement the attribute.

compiler/src/iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenAttrs.cpp

@@ -8,11 +8,13 @@
 
 #include "iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenDialect.h"
 #include "iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenInterfaces.h"
+#include "iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenTypes.h"
 #include "llvm/ADT/TypeSwitch.h"
 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/Arith/Utils/Utils.h"
 #include "mlir/Dialect/Transform/IR/TransformOps.h"
 #include "mlir/Dialect/Utils/StaticValueUtils.h"
+#include "mlir/Dialect/Utils/StructuredOpsUtils.h"
 #include "mlir/IR/DialectImplementation.h"
 #include "mlir/IR/StorageUniquerSupport.h"
 
@@ -460,6 +462,21 @@ int64_t WorkgroupMappingAttr::getRelativeIndex() const {
   return getMappingId();
 }
 
+//===---------------------------------------------------------------------===//
+// iree_codegen.encoding_layout
+//===---------------------------------------------------------------------===//
+
+MaterializeEncodingInfo
+EncodingNopLayoutAttr::getEncodingInfo(RankedTensorType type) const {
+  return MaterializeEncodingInfo{};
+}
+
+Operation *EncodingNopLayoutAttr::lowerOp(OpBuilder &b, Operation *op,
+                                          TypeRange convertedResTypes,
+                                          ValueRange convertedOperands) const {
+  return clone(b, op, convertedResTypes, convertedOperands);
+}
+
 //===----------------------------------------------------------------------===//
 // Initialize attributes
 //===----------------------------------------------------------------------===//