[DT] Collapse MaterializeEncodingIntoPaddingPass into the generic pass.

compiler/src/iree/compiler/Codegen/Common/BUILD.bazel

@@ -127,7 +127,6 @@ iree_compiler_cc_library(
         "LoweringConfigInterpreter.cpp",
         "MaterializeEncoding.cpp",
         "MaterializeEncodingIntoNop.cpp",
-        "MaterializeEncodingIntoPadding.cpp",
         "MaterializeEncodingPatterns.cpp",
         "MaterializeTuningSpecsPass.cpp",
         "MathTransformPass.cpp",

compiler/src/iree/compiler/Codegen/Common/CMakeLists.txt

@@ -120,7 +120,6 @@ iree_cc_library(
     "LoweringConfigInterpreter.cpp"
     "MaterializeEncoding.cpp"
     "MaterializeEncodingIntoNop.cpp"
-    "MaterializeEncodingIntoPadding.cpp"
     "MaterializeEncodingPatterns.cpp"
     "MaterializeTuningSpecsPass.cpp"
     "MathTransformPass.cpp"

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

@@ -10,11 +10,26 @@
 #include "iree/compiler/Codegen/Dialect/Codegen/Utils/Utils.h"
 #include "iree/compiler/Dialect/Encoding/IR/EncodingOps.h"
 #include "iree/compiler/Dialect/TensorExt/IR/TensorExtTypes.h"
-#include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/Transforms/DialectConversion.h"
 
 namespace mlir::iree_compiler {
 
+//===---------------------------------------------------------------------===//
+// Utilities for testing purpose.
+//===---------------------------------------------------------------------===//
+
+namespace detail {
+/// Kinds of testing resolvers for MaterializeDeviceEncodingPass.
+enum class TestingResolverKind {
+  // iree_gpu.gpu_encoding_resolver<>
+  kGPUDataTiling,
+  // iree_gpu.gpu_padding_resolver<>
+  kGPUPadding,
+  // Do not create any GPU specific resolver, so fallback to identity resolver.
+  kNone
+};
+} // namespace detail
+
 //===---------------------------------------------------------------------===//
 // TypeConverter
 //===---------------------------------------------------------------------===//

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

@@ -10,7 +10,9 @@
 #include "iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenDialect.h"
 #include "iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUAttrs.h"
 #include "iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUDialect.h"
+#include "iree/compiler/Codegen/Dialect/GPU/TargetUtils/KnownTargets.h"
 #include "iree/compiler/Codegen/Utils/GPUUtils.h"
+#include "iree/compiler/Dialect/Encoding/IR/EncodingDialect.h"
 #include "iree/compiler/Dialect/Encoding/IR/EncodingTypes.h"
 #include "iree/compiler/Dialect/HAL/Analysis/DeviceAnalysis.h"
 #include "iree/compiler/Dialect/HAL/IR/HALTypes.h"
@@ -59,42 +61,43 @@ updateFuncSignature(FunctionOpInterface funcOp,
   }
 }
 
-static LogicalResult
-materializeFuncOpEncodings(FunctionOpInterface funcOp,
-                           IREE::HAL::ExecutableTargetAttr targetAttr,
-                           bool testCLGPUTarget = false) {
+static LogicalResult materializeFuncOpEncodings(
+    FunctionOpInterface funcOp, IREE::HAL::ExecutableTargetAttr targetAttr,
+    ::mlir::iree_compiler::detail::TestingResolverKind resolverSource =
+        ::mlir::iree_compiler::detail::TestingResolverKind::kNone) {
   MLIRContext *ctx = funcOp.getContext();
   {
     RewritePatternSet patterns(ctx);
     DictionaryAttr targetConfig =
         targetAttr ? targetAttr.getConfiguration() : nullptr;
-    // Check if the encoding resolver is a GPUPaddingResolverAttr. For padding
-    // encoding materialization, we use a separate pass, so skip materialization
-    // here.
-    // TODO(#20160): Support GPUPaddingResolverAttr materialization through this
-    // pass, and remove the ad-hoc materialization pass for padding.
-    if (targetConfig && targetConfig.getAs<IREE::GPU::GPUPaddingResolverAttr>(
-                            IREE::Encoding::kEncodingResolverAttrName)) {
-      LDBG()
-          << "Found GPUPaddingResolverAttr encoding resolver. Materialization "
-             "will be handled later.";
-      return success();
-    }
-
-    auto getTestTargetOrNopLayout =
+    auto getTestTargetResolverOrIdentityResolver =
         [&]() -> IREE::Encoding::LayoutMaterializerAttr {
-      if (testCLGPUTarget) {
-        LDBG() << "Select GPUEncodingResolverAttr attribute as the layout "
-                  "attribute. (testCLGPUTarget)";
-        SmallVector<NamedAttribute> configItems;
-        // Setting a nullptr to `target` below returns the target from command
-        // line.
-        IREE::GPU::TargetAttr targetAttr =
-            getGPUTargetAttr(ctx, /*target=*/nullptr);
+      SmallVector<NamedAttribute> configItems;
+      IREE::GPU::TargetAttr targetAttr =
+          getGPUTargetAttr(ctx, /*target=*/nullptr);
+      if (targetAttr) {
         addConfigGPUTarget(ctx, targetAttr, configItems);
-        return cast<IREE::Encoding::LayoutMaterializerAttr>(
-            IREE::GPU::GPUEncodingResolverAttr::get(
-                ctx, DictionaryAttr::get(ctx, configItems)));
+        switch (resolverSource) {
+        case ::mlir::iree_compiler::detail::TestingResolverKind::
+            kGPUDataTiling: {
+          LDBG() << "Select GPUEncodingResolverAttr attribute as the layout "
+                    "attribute. (kGPUDataTiling)";
+          return cast<IREE::Encoding::LayoutMaterializerAttr>(
+              IREE::GPU::GPUEncodingResolverAttr::get(
+                  ctx, DictionaryAttr::get(ctx, configItems)));
+        }
+        case ::mlir::iree_compiler::detail::TestingResolverKind::kGPUPadding: {
+          LDBG() << "Select GPUPaddingResolverAttr attribute as the layout "
+                    "attribute. (kGPUPadding)";
+          std::optional<IREE::GPU::L1CacheInfo> cache =
+              IREE::GPU::getL1CacheInfo(targetAttr);
+          return cast<IREE::Encoding::LayoutMaterializerAttr>(
+              IREE::GPU::GPUPaddingResolverAttr::get(ctx, cache->cacheLineBytes,
+                                                     cache->cacheSets));
+        }
+        case ::mlir::iree_compiler::detail::TestingResolverKind::kNone:
+          break;
+        }
       }
       LDBG() << "Select IdentityResolverAttr attribute as the layout "
                 "attribute (Encoding resolver unknown or unsupported).";
@@ -107,8 +110,9 @@ materializeFuncOpEncodings(FunctionOpInterface funcOp,
     // so it can access the target info during materialization.
     //
     // If the layoutAttr was not found, or if it does not implement the layout
-    // resolver interface, fall back to the resolver for getCLGPUTarget. If
-    // there is also no test target set, fall back to the nop layout.
+    // resolver interface, fall back to the resolver for TestingResolverKind
+    // resolver. If there is also no test target set or it is kNone, fall back
+    // to the identity resolver.
     IREE::Encoding::LayoutMaterializerAttr layoutAttr =
         targetConfig
             ? targetConfig.getAs<IREE::Encoding::LayoutMaterializerAttr>(
@@ -121,7 +125,7 @@ materializeFuncOpEncodings(FunctionOpInterface funcOp,
         layoutAttr && resolverAttr
             ? cast<IREE::Encoding::LayoutMaterializerAttr>(
                   resolverAttr.cloneWithSimplifiedConfig(targetConfig))
-            : getTestTargetOrNopLayout();
+            : getTestTargetResolverOrIdentityResolver();
 
     LDBG() << "Selected Encoding::LayoutMaterializerAttr with target "
               "configuration: "
@@ -132,10 +136,10 @@ materializeFuncOpEncodings(FunctionOpInterface funcOp,
     populateMaterializeEncodingPatterns(patterns, target, typeConverter);
 
     // Replace any unrealized conversions to tensor.cast ops if they come from
-    // block arguments. The function signature is updated to match the converted
-    // types after the partial conversion. This is used in testing, where
-    // function arguments have encodings to reduce the amount of IR, but we do
-    // not expect function arguments to have encodings in practice.
+    // block arguments. The function signature is updated to match the
+    // converted types after the partial conversion. This is used in testing,
+    // where function arguments have encodings to reduce the amount of IR, but
+    // we do not expect function arguments to have encodings in practice.
     auto castFnArguments = [](OpBuilder &builder, Type resultTy,
                               ValueRange inputs, Location loc) -> Value {
       if (inputs.size() != 1) {
@@ -177,11 +181,11 @@ materializeFuncOpEncodings(FunctionOpInterface funcOp,
         patterns, [](OpOperand *opOperand) {
           Operation *producer = opOperand->get().getDefiningOp();
           Operation *consumer = opOperand->getOwner();
-          // If we have a pack/unpack consumer and a producer that has multiple
-          // uses, this _probably_ means the producer won't get dce'd. If that
-          // is the case, by folding the consumer pack/unpack, we break the
-          // producer consumer chain between them and inhibit fusion later in
-          // the pipeline.
+          // If we have a pack/unpack consumer and a producer that has
+          // multiple uses, this _probably_ means the producer won't get
+          // dce'd. If that is the case, by folding the consumer pack/unpack,
+          // we break the producer consumer chain between them and inhibit
+          // fusion later in the pipeline.
           if (isa<linalg::PackOp, linalg::UnPackOp>(consumer) &&
               isa_and_nonnull<TilingInterface>(producer) &&
               !producer->hasOneUse())
@@ -296,15 +300,16 @@ struct MaterializeDeviceEncodingPass final
 
   void getDependentDialects(DialectRegistry &registry) const override {
     registry.insert<arith::ArithDialect, tensor::TensorDialect,
-                    vector::VectorDialect, IREE::Codegen::IREECodegenDialect,
+                    vector::VectorDialect, IREE::Encoding::IREEEncodingDialect,
+                    IREE::Codegen::IREECodegenDialect,
                     IREE::CPU::IREECPUDialect, IREE::GPU::IREEGPUDialect>();
   }
 
   void runOnOperation() override {
     FunctionOpInterface funcOp = getOperation();
     auto executableTargetAttr = IREE::HAL::ExecutableTargetAttr::lookup(funcOp);
     if (failed(materializeFuncOpEncodings(funcOp, executableTargetAttr,
-                                          testCLGPUTarget))) {
+                                          testGPUEncodingResolver))) {
       return signalPassFailure();
     }
   }

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

@@ -13,6 +13,7 @@
 #define IREE_COMPILER_CODEGEN_COMMON_PASSES_H_
 
 #include "iree/compiler/Codegen/Common/CombineLayoutTransformation.h"
+#include "iree/compiler/Codegen/Common/EncodingUtils.h"
 #include "iree/compiler/Codegen/Common/PassUtils.h"
 #include "iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenAttrs.h"
 #include "iree/compiler/Codegen/Utils/Utils.h"
@@ -49,9 +50,6 @@ void addIREEComprehensiveBufferizePasses(
 /// Populate Encoding to Nop pass and canonicalizer pass to the pipeline.
 void addEncodingToNopPasses(FunctionLikeNest &passManager);
 
-/// Populate Encoding to padding pass and canonicalizer pass to the pipeline.
-void addEncodingToPaddingPasses(FunctionLikeNest &passManager);
-
 /// Links nested transform dialect tuning specs named sequences into a single
 /// entry point. Returns the new named sequence op (inserted into the `module`)
 /// that includes the nested tuning specs, or a null op when no nested named

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

@@ -616,8 +616,18 @@ def MaterializeDeviceEncodingPass :
     InterfacePass<"iree-codegen-materialize-device-encoding", "mlir::FunctionOpInterface"> {
   let summary = "Materialize the encoding for tensor as specified by the backend.";
   let options = [
-    Option<"testCLGPUTarget", "test-cl-gpu-target", "bool", /*default=*/"false",
-           "Flag used for lit-testing GPU target only. Not for general usage">,
+    Option<"testGPUEncodingResolver", "test-gpu-encoding-resolver", "::mlir::iree_compiler::detail::TestingResolverKind",
+           /*default=*/"detail::TestingResolverKind::kNone",
+           "Flag used for lit-testing GPU target only. Not for general usage",
+           [{llvm::cl::values(
+             clEnumValN(::mlir::iree_compiler::detail::TestingResolverKind::kGPUDataTiling, "gpu_data_tiling",
+                        "Use iree_gpu.gpu_encoding_resolver<>."),
+             clEnumValN(::mlir::iree_compiler::detail::TestingResolverKind::kGPUPadding, "gpu_padding",
+                        "Use iree_gpu.gpu_padding_resolver<>."),
+             clEnumValN(::mlir::iree_compiler::detail::TestingResolverKind::kNone, "none",
+                        "Ignore testing resolvers.")
+      )}]>,
+
   ];
 }
 
@@ -627,17 +637,6 @@ def MaterializeEncodingIntoNopPass :
   let summary = "Drop the encodings from tensor types with encodings.";
 }
 
-def MaterializeEncodingIntoPaddingPass :
-    InterfacePass<"iree-codegen-materialize-encoding-into-padding", "mlir::FunctionOpInterface"> {
-  let summary = "Materialize `#iree_encoding.padding` attributes.";
-  let description = [{
-    Handles padding introduced by `padding` encoding layouts, which
-    requires `iree_tensor_ext.dispatch.tensor.load`/`.store` to be adjusted to account for
-    padding regions.
-    Materializes any other encoding layouts into nop.
-  }];
-}
-
 def MaterializeTuningSpecsPass : Pass<"iree-codegen-materialize-tuning-specs", "ModuleOp"> {
   let summary =
       "Load tuning spec transform dialect libraries and encode them in the module";