[DT] Unify encoding materialization pass into a single pass. (iree-or…

…g#19454)

The revision creates a generic materialization pass and uses it for
backends that implement data-tiling. After months of development, we
identify that the needs of GPU is a superset of the needs of CPU. To be
more specific, it has the additional "swizzle" field in terms of layout.
It means that the GPU set_encoding/unset_encoding lowering patterns
cover the needs of CPU path. The lowering of contraction ops is
different. CPU lowers it to mmt4d op, while GPU lowers it to multi_mma
op. However, the lowering of contraction is implemented through
attribute interface. Thus, we can have a generic pattern to lower
contraction ops.

To make the review process much easier, the revision is created by 5
commits.

1. It directly creates the MaterializeEncoding pass and copy-paste the
GPU patterns: SetEncodingOpLoweringConversion,
UnSetEncodingOpLoweringConversion, and MaterializeContractionOp. In the
first commit, it also updates the GPU tests to use the new pass.
2. The GPU data-tiling does not support element-wise generic op lowering
atm. The second commit moves the pattern to shared pattern set and bail
out when swizzle is present. This is an NFC for both pipelines.
3. The third commit replaces the existing materialization pass with the
generic pass, and deletes all the legacy passes.
4. The four commit moves the lit tests from `Common/[CPU|GPU]/test` to
`Common/test`.
5. Now there are duplicate patterns for set_encoding, unset_encoding,
and contraction ops lowering. The last commit deletes the legacy
patterns, and move the patterns from MaterializeEncoding.cpp to where
the legacy patterns locate. Furthermore, it renames the file as
`MaterializeEncodingPatterns.cpp`.

The revision retains the MaterializeEncodingIntoNop pass, and add a TODO
item. Because it is still used by MaterializeHomogeneousEncoding pass.
It can be deleted once we deprecate the early materialization path.

---------

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

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

@@ -125,8 +125,9 @@ iree_compiler_cc_library(
         "LinkTuningSpecsPass.cpp",
         "LowerExecutableUsingTransformDialect.cpp",
         "LowerUKernelsToCalls.cpp",
+        "MaterializeEncoding.cpp",
         "MaterializeEncodingIntoNop.cpp",
-        "MaterializeEncodingIntoPackUnPack.cpp",
+        "MaterializeEncodingPatterns.cpp",
         "MaterializeTuningSpecsPass.cpp",
         "MemrefCopyToLinalg.cpp",
         "NormalizeLoopBounds.cpp",
@@ -173,8 +174,10 @@ iree_compiler_cc_library(
         ":PassHeaders",
         ":PassesIncGen",
         "//compiler/src/iree/compiler/Codegen/Common:FoldTensorExtractOpIncGen",
+        "//compiler/src/iree/compiler/Codegen/Dialect/CPU/IR:IREECPUDialect",
         "//compiler/src/iree/compiler/Codegen/Dialect/Codegen/IR:IREECodegenDialect",
         "//compiler/src/iree/compiler/Codegen/Dialect/Codegen/Utils",
+        "//compiler/src/iree/compiler/Codegen/Dialect/GPU/IR:IREEGPUDialect",
         "//compiler/src/iree/compiler/Codegen/Dialect/VectorExt/IR:IREEVectorExtDialect",
         "//compiler/src/iree/compiler/Codegen/Interfaces:BufferizationInterfaces",
         "//compiler/src/iree/compiler/Codegen/Interfaces:PartitionableLoopsInterface",
@@ -183,9 +186,11 @@ iree_compiler_cc_library(
         "//compiler/src/iree/compiler/Codegen/Utils",
         "//compiler/src/iree/compiler/Dialect/Encoding/IR",
         "//compiler/src/iree/compiler/Dialect/Flow/IR",
+        "//compiler/src/iree/compiler/Dialect/HAL/Analysis",
         "//compiler/src/iree/compiler/Dialect/HAL/IR",
         "//compiler/src/iree/compiler/Dialect/LinalgExt/IR",
         "//compiler/src/iree/compiler/Dialect/LinalgExt/Transforms",
+        "//compiler/src/iree/compiler/Dialect/Stream/Analysis",
         "//compiler/src/iree/compiler/Dialect/Util/Analysis",
         "//compiler/src/iree/compiler/Dialect/Util/IR",
         "//compiler/src/iree/compiler/Utils",

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

@@ -117,8 +117,9 @@ iree_cc_library(
     "LinkTuningSpecsPass.cpp"
     "LowerExecutableUsingTransformDialect.cpp"
     "LowerUKernelsToCalls.cpp"
+    "MaterializeEncoding.cpp"
     "MaterializeEncodingIntoNop.cpp"
-    "MaterializeEncodingIntoPackUnPack.cpp"
+    "MaterializeEncodingPatterns.cpp"
     "MaterializeTuningSpecsPass.cpp"
     "MemrefCopyToLinalg.cpp"
     "NormalizeLoopBounds.cpp"
@@ -203,8 +204,10 @@ iree_cc_library(
     MLIRVectorTransforms
     MLIRViewLikeInterface
     iree::compiler::Codegen::Common::FoldTensorExtractOpIncGen
+    iree::compiler::Codegen::Dialect::CPU::IR::IREECPUDialect
     iree::compiler::Codegen::Dialect::Codegen::IR::IREECodegenDialect
     iree::compiler::Codegen::Dialect::Codegen::Utils
+    iree::compiler::Codegen::Dialect::GPU::IR::IREEGPUDialect
     iree::compiler::Codegen::Dialect::VectorExt::IR::IREEVectorExtDialect
     iree::compiler::Codegen::Interfaces::BufferizationInterfaces
     iree::compiler::Codegen::Interfaces::PartitionableLoopsInterface
@@ -213,9 +216,11 @@ iree_cc_library(
     iree::compiler::Codegen::Utils
     iree::compiler::Dialect::Encoding::IR
     iree::compiler::Dialect::Flow::IR
+    iree::compiler::Dialect::HAL::Analysis
     iree::compiler::Dialect::HAL::IR
     iree::compiler::Dialect::LinalgExt::IR
     iree::compiler::Dialect::LinalgExt::Transforms
+    iree::compiler::Dialect::Stream::Analysis
     iree::compiler::Dialect::Util::Analysis
     iree::compiler::Dialect::Util::IR
     iree::compiler::Utils

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

@@ -45,7 +45,6 @@ iree_compiler_cc_library(
     name = "CommonCPUPasses",
     srcs = [
         "CPULowerToUKernels.cpp",
-        "CPUMaterializeEncodings.cpp",
         "CPUPrepareUkernels.cpp",
         "Passes.cpp",
     ],
@@ -56,16 +55,13 @@ iree_compiler_cc_library(
         ":PassHeaders",
         ":PassesIncGen",
         "//compiler/src/iree/compiler/Codegen/Common",
-        "//compiler/src/iree/compiler/Codegen/Dialect/CPU/IR:IREECPUDialect",
         "//compiler/src/iree/compiler/Codegen/Dialect/Codegen/IR:IREECodegenDialect",
         "//compiler/src/iree/compiler/Codegen/Dialect/Codegen/Utils",
         "//compiler/src/iree/compiler/Codegen/Interfaces:UKernelOpInterface",
         "//compiler/src/iree/compiler/Codegen/Transforms",
         "//compiler/src/iree/compiler/Codegen/Utils",
         "//compiler/src/iree/compiler/Dialect/Encoding/IR",
-        "//compiler/src/iree/compiler/Dialect/HAL/Analysis",
         "//compiler/src/iree/compiler/Dialect/HAL/IR",
-        "//compiler/src/iree/compiler/Dialect/Stream/Analysis",
         "//runtime/src/iree/builtins/ukernel:exported_bits",
         "@llvm-project//llvm:Support",
         "@llvm-project//mlir:AffineDialect",

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

@@ -42,7 +42,6 @@ iree_cc_library(
     "Passes.h"
   SRCS
     "CPULowerToUKernels.cpp"
-    "CPUMaterializeEncodings.cpp"
     "CPUPrepareUkernels.cpp"
     "Passes.cpp"
   DEPS
@@ -78,16 +77,13 @@ iree_cc_library(
     MLIRVectorTransforms
     iree::builtins::ukernel::exported_bits
     iree::compiler::Codegen::Common
-    iree::compiler::Codegen::Dialect::CPU::IR::IREECPUDialect
     iree::compiler::Codegen::Dialect::Codegen::IR::IREECodegenDialect
     iree::compiler::Codegen::Dialect::Codegen::Utils
     iree::compiler::Codegen::Interfaces::UKernelOpInterface
     iree::compiler::Codegen::Transforms
     iree::compiler::Codegen::Utils
     iree::compiler::Dialect::Encoding::IR
-    iree::compiler::Dialect::HAL::Analysis
     iree::compiler::Dialect::HAL::IR
-    iree::compiler::Dialect::Stream::Analysis
   PUBLIC
 )
 

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

@@ -13,26 +13,6 @@ include "mlir/Pass/PassBase.td"
 // Common Passes used for CPU-like backends (keep alphabetical)
 //===---------------------------------------------------------------------===//
 
-def CPUMaterializeHostEncodingPass :
-    Pass<"iree-codegen-cpu-materialize-host-encoding", "mlir::ModuleOp"> {
-  let summary = "Convert encoding-specific operations based on target attributes.";
-  let description = [{
-    Examples:
-      encoding.set_encoding   -> tensor.pack
-      encoding.unset_encoding -> tensor.unpack
-      linalg.matmul             -> linalg.mmt4d  "}];
-}
-
-def CPUMaterializeDeviceEncodingPass :
-    InterfacePass<"iree-codegen-cpu-materialize-device-encoding", "mlir::FunctionOpInterface"> {
-  let summary = "Convert encoding-specific operations based on target attributes.";
-  let description = [{
-    Examples:
-      encoding.set_encoding   -> tensor.pack
-      encoding.unset_encoding -> tensor.unpack
-      linalg.matmul             -> linalg.mmt4d  "}];
-}
-
 def CPULowerToUKernelsPass :
     Pass<"iree-codegen-cpu-lower-to-ukernels", ""> {
   let summary =

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

@@ -19,10 +19,8 @@ iree_lit_test_suite(
     srcs = enforce_glob(
         # keep sorted
         [
-            "llvmcpu_materialize_encoding.mlir",
             "lower_to_ukernel_ops.mlir",
             "prepare_ukernels.mlir",
-            "vmvx_materialize_encoding.mlir",
         ],
         include = ["*.mlir"],
     ),

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

@@ -14,10 +14,8 @@ iree_lit_test_suite(
   NAME
     lit
   SRCS
-    "llvmcpu_materialize_encoding.mlir"
     "lower_to_ukernel_ops.mlir"
     "prepare_ukernels.mlir"
-    "vmvx_materialize_encoding.mlir"
   TOOLS
     FileCheck
     iree-opt

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

@@ -93,17 +93,9 @@ FailureOr<Value> lowerUnsetEncodingToUnpackOp(
     Value packedValue, const MaterializeEncodingTypeConverter &typeConverter,
     MaterializeEncodingValueFn materializeEncodingValueFn);
 
-/// Pouplates the set of patterns that lowers set_encoding, unset_encoding, and
-/// upstream dialect ops with encoding types to pack/unpack ops.
-void populateMaterializeEncodingIntoPackUnPackPatterns(
-    RewritePatternSet &patterns,
-    MaterializeEncodingTypeConverter &typeConverter,
-    MaterializeEncodingValueFn materializeEncodingValueFn);
-
-/// Pouplates the set of patterns that lowers shape-like operations (e.g., Flow
-/// ops, Hal ops, tensor.empty, linalg.fill, etc) with encoding types to the
-/// same op with materialized shapes.
-void populateShapeIndependentMaterializeEncodingPatterns(
+/// Pouplates the set of patterns that lowers operations with encoding types to
+/// operations without encodings.
+void populateMaterializeEncodingPatterns(
     RewritePatternSet &patterns, MaterializeEncodingConversionTarget &target,
     MaterializeEncodingTypeConverter &typeConverter,
     MaterializeEncodingValueFn materializeEncodingValueFn);

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

@@ -65,7 +65,6 @@ iree_compiler_cc_library(
         "GPUGreedilyDistributeToThreads.cpp",
         "GPUInferMemorySpace.cpp",
         "GPULowerToUKernels.cpp",
-        "GPUMaterializeEncoding.cpp",
         "GPUMultiBuffering.cpp",
         "GPUNestedLayoutDistributionPatterns.cpp",
         "GPUPackToIntrinsics.cpp",
@@ -107,10 +106,7 @@ iree_compiler_cc_library(
         "//compiler/src/iree/compiler/Codegen/Transforms",
         "//compiler/src/iree/compiler/Codegen/Utils",
         "//compiler/src/iree/compiler/Codegen/Utils:VectorOpUtils",
-        "//compiler/src/iree/compiler/Dialect/Encoding/IR",
-        "//compiler/src/iree/compiler/Dialect/HAL/Analysis",
         "//compiler/src/iree/compiler/Dialect/HAL/IR",
-        "//compiler/src/iree/compiler/Dialect/Stream/Analysis",
         "//compiler/src/iree/compiler/Utils",
         "@llvm-project//llvm:Support",
         "@llvm-project//mlir:AMDGPUDialect",

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

@@ -63,7 +63,6 @@ iree_cc_library(
     "GPUGreedilyDistributeToThreads.cpp"
     "GPUInferMemorySpace.cpp"
     "GPULowerToUKernels.cpp"
-    "GPUMaterializeEncoding.cpp"
     "GPUMultiBuffering.cpp"
     "GPUNestedLayoutDistributionPatterns.cpp"
     "GPUPackToIntrinsics.cpp"
@@ -140,10 +139,7 @@ iree_cc_library(
     iree::compiler::Codegen::Transforms
     iree::compiler::Codegen::Utils
     iree::compiler::Codegen::Utils::VectorOpUtils
-    iree::compiler::Dialect::Encoding::IR
-    iree::compiler::Dialect::HAL::Analysis
     iree::compiler::Dialect::HAL::IR
-    iree::compiler::Dialect::Stream::Analysis
     iree::compiler::Utils
   PUBLIC
 )