[DT] Collapse MaterializeEncodingIntoPaddingPass into the generic pass.

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

Author: hanhanW

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,23 @@
 #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 {
 
+//===---------------------------------------------------------------------===//
+// Utilties for testing purpose.
+//===---------------------------------------------------------------------===//
+
+namespace detail {
+/// Kinds of testing resolvers for MaterializeDeviceEncodingPass.
+enum class TestingResolverKind {
+  kGPUDataTiling, // iree_gpu.gpu_encoding_resolver<>
+  kGPUPadding,    // iree_gpu.gpu_padding_resolver<>
+  kNone           // Do not create any specific resolver.
+};
+} // namespace detail
+
 //===---------------------------------------------------------------------===//
 // TypeConverter
 //===---------------------------------------------------------------------===//

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

@@ -10,6 +10,7 @@
 #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/EncodingTypes.h"
 #include "iree/compiler/Dialect/HAL/Analysis/DeviceAnalysis.h"
@@ -59,43 +60,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 =
         [&]() -> IREE::Encoding::LayoutMaterializerAttr {
-      if (testCLGPUTarget) {
+      LDBG() << "Select GPUEncodingResolverAttr attribute as the layout "
+                "attribute. (testCLGPUTarget)";
+      SmallVector<NamedAttribute> configItems;
+      IREE::GPU::TargetAttr targetAttr =
+          getGPUTargetAttr(ctx, /*target=*/nullptr);
+      addConfigGPUTarget(ctx, targetAttr, configItems);
+      switch (resolverSource) {
+      case ::mlir::iree_compiler::detail::TestingResolverKind::kGPUDataTiling: {
         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);
-        addConfigGPUTarget(ctx, targetAttr, configItems);
+                  "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).";
       return cast<IREE::Encoding::LayoutMaterializerAttr>(
@@ -132,10 +133,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) {
@@ -174,11 +175,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())
@@ -301,7 +302,7 @@ struct MaterializeDeviceEncodingPass final
     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/MaterializeEncodingIntoPadding.cpp

@@ -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.h

@@ -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";

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

@@ -1,4 +1,4 @@
-// RUN: iree-opt --pass-pipeline="builtin.module(func.func(iree-codegen-materialize-device-encoding{test-cl-gpu-target}))" \
+// RUN: iree-opt --pass-pipeline="builtin.module(func.func(iree-codegen-materialize-device-encoding{test-gpu-encoding-resolver=gpu_data_tiling}))" \
 // RUN:   --iree-gpu-test-target=gfx1100 \
 // RUN:   --split-input-file %s | FileCheck %s
 

compiler/src/iree/compiler/Codegen/Common/test/materialize_encoding_gfx1100.mlir

@@ -1,4 +1,4 @@
-// RUN: iree-opt --pass-pipeline="builtin.module(func.func(iree-codegen-materialize-device-encoding{test-cl-gpu-target}))" \
+// RUN: iree-opt --pass-pipeline="builtin.module(func.func(iree-codegen-materialize-device-encoding{test-gpu-encoding-resolver=gpu_data_tiling}))" \
 // RUN:   --iree-gpu-test-target=gfx908 \
 // RUN:   --split-input-file %s | FileCheck %s
 

compiler/src/iree/compiler/Codegen/Common/test/materialize_encoding_gfx908.mlir

@@ -1,4 +1,4 @@
-// RUN: iree-opt --pass-pipeline="builtin.module(func.func(iree-codegen-materialize-device-encoding{test-cl-gpu-target}))" \
+// RUN: iree-opt --pass-pipeline="builtin.module(func.func(iree-codegen-materialize-device-encoding{test-gpu-encoding-resolver=gpu_data_tiling}))" \
 // RUN:   --iree-gpu-test-target=gfx90a \
 // RUN:   --split-input-file %s | FileCheck %s