[DT][NFC] Refactor linalg.fill/generic op lowering to interface implementation. (iree-org#22343)

The revision deletes two patterns and moves the implementation to
interface methods:
- MaterializeDPSOperation<linalg::FillOp>
- MaterializeDPSOperation<linalg::GenericOp>

The lowerGenericOpWithResolvedLayouts method is very similar to the
original lowerGenericOpWithEncoding. The new method takes the same input
arguments like `lowerOp` method with an additional
`LayoutMaterializerAttr` attribute, because it needs the packing info to
generate indexing maps. In the new implementation, it uses
`getEncodingInfoFromLayout` method directly while the method in the type
converter is just a wrapper. I.e., it reduces the dependency from the
type converter. See
https://gist.github.com/hanhanW/e89aa9e5052f8db37b2543fa368ed605 for the
diff.

It is a step towards iree-org#20160, as
the padding resolver no longer rely on "fallback" path to drop the
encodings. It has its own implementation that uses `clone`.

---------

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

Author: hanhanW

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

@@ -44,6 +44,7 @@ iree_compiler_cc_library(
         "@llvm-project//mlir:DialectUtils",
         "@llvm-project//mlir:IR",
         "@llvm-project//mlir:LinalgDialect",
+        "@llvm-project//mlir:LinalgTransforms",
         "@llvm-project//mlir:TensorDialect",
     ],
 )

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

@@ -30,6 +30,7 @@ iree_cc_library(
     MLIRArithDialect
     MLIRIR
     MLIRLinalgDialect
+    MLIRLinalgTransforms
     MLIRTensorDialect
     iree::compiler::Codegen::Dialect::CPU::IR::IREECPUDialect
     iree::compiler::Codegen::Dialect::Codegen::IR::IREECodegenDialect

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

@@ -43,7 +43,6 @@
 #include "iree/compiler/Codegen/ExternalInterfaces/Utils.h"
 #include "iree/compiler/Codegen/Utils/CPUUtils.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/Encoding/Utils/Utils.h"
 #include "llvm/Support/DebugLog.h"
@@ -284,11 +283,11 @@ TileMxNxK chooseMatmulTile(ArrayRef<TileMxNxK> enumeratedTiles,
   return bestRatedTile;
 }
 
-FailureOr<Operation *> lowerContractionOpWithEncoding(
+Operation *lowerContractionOpWithEncoding(
     OpBuilder &builder, linalg::LinalgOp linalgOp, ValueRange operands,
     IREE::Encoding::LayoutMaterializerAttr layoutAttr) {
   if (!linalgOp.hasPureTensorSemantics()) {
-    return failure();
+    return nullptr;
   }
 
   auto inputs = linalgOp.getDpsInputOperands();
@@ -301,14 +300,14 @@ FailureOr<Operation *> lowerContractionOpWithEncoding(
   auto rhsEncoding = IREE::Encoding::getEncodingAttr(rhsType);
   auto resultEncoding = IREE::Encoding::getEncodingAttr(resultType);
   if (!lhsEncoding || !rhsEncoding || !resultEncoding) {
-    return failure();
+    return nullptr;
   }
 
   if (lhsEncoding.getOperandIndex().getValue() != IREE::Encoding::MATMUL_LHS ||
       rhsEncoding.getOperandIndex().getValue() != IREE::Encoding::MATMUL_RHS ||
       resultEncoding.getOperandIndex().getValue() !=
           IREE::Encoding::MATMUL_RESULT) {
-    return failure();
+    return nullptr;
   }
 
   MaterializeEncodingInfo encodingInfo = {};
@@ -725,11 +724,21 @@ struct CPUEncodingResolverMaterializerAttr final
     if (!linalgOp) {
       return nullptr;
     }
-
-    FailureOr<Operation *> newOp = lowerContractionOpWithEncoding(
-        b, linalgOp, convertedOperands,
-        cast<IREE::Encoding::LayoutMaterializerAttr>(layoutAttr));
-    return newOp.value_or(nullptr);
+    if (auto fillOp = dyn_cast<linalg::FillOp>(op)) {
+      return lowerFillOpWithResolvedLayouts(b, fillOp, convertedResTypes,
+                                            convertedOperands);
+    }
+    if (linalg::isaContractionOpInterface(linalgOp)) {
+      return lowerContractionOpWithEncoding(
+          b, linalgOp, convertedOperands,
+          cast<IREE::Encoding::LayoutMaterializerAttr>(layoutAttr));
+    }
+    if (auto genericOp = dyn_cast<linalg::GenericOp>(op)) {
+      return lowerGenericOpWithResolvedLayouts(
+          b, genericOp, convertedResTypes, convertedOperands,
+          cast<IREE::Encoding::LayoutMaterializerAttr>(attr));
+    }
+    return nullptr;
   }
 };
 
@@ -869,11 +878,21 @@ struct VMVXEncodingResolverMaterializerAttr final
     if (!linalgOp) {
       return nullptr;
     }
-
-    FailureOr<Operation *> newOp = lowerContractionOpWithEncoding(
-        b, linalgOp, convertedOperands,
-        cast<IREE::Encoding::LayoutMaterializerAttr>(layoutAttr));
-    return newOp.value_or(nullptr);
+    if (auto fillOp = dyn_cast<linalg::FillOp>(op)) {
+      return lowerFillOpWithResolvedLayouts(b, fillOp, convertedResTypes,
+                                            convertedOperands);
+    }
+    if (linalg::isaContractionOpInterface(linalgOp)) {
+      return lowerContractionOpWithEncoding(
+          b, linalgOp, convertedOperands,
+          cast<IREE::Encoding::LayoutMaterializerAttr>(layoutAttr));
+    }
+    if (auto genericOp = dyn_cast<linalg::GenericOp>(op)) {
+      return lowerGenericOpWithResolvedLayouts(
+          b, genericOp, convertedResTypes, convertedOperands,
+          cast<IREE::Encoding::LayoutMaterializerAttr>(attr));
+    }
+    return nullptr;
   }
 };
 

compiler/src/iree/compiler/Codegen/ExternalInterfaces/CPUEncodingExternalModels.cpp

@@ -33,11 +33,9 @@
 #include "iree/compiler/Codegen/Dialect/GPU/IR/GPUTileSwizzleUtils.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/IR/IREEGPUOps.h"
 #include "iree/compiler/Codegen/Dialect/GPU/TargetUtils/KnownTargets.h"
 #include "iree/compiler/Codegen/ExternalInterfaces/Utils.h"
 #include "iree/compiler/Codegen/Utils/GPUUtils.h"
-#include "iree/compiler/Dialect/Encoding/IR/EncodingOps.h"
 #include "iree/compiler/Dialect/Encoding/IR/EncodingTypes.h"
 #include "iree/compiler/Dialect/Encoding/Utils/Utils.h"
 #include "iree/compiler/Dialect/LinalgExt/Utils/MatchUtils.h"
@@ -335,10 +333,6 @@ static Operation *lowerContractionOrScaledContractionOpToInnerTiledOp(
   if (!linalgOp.hasPureTensorSemantics()) {
     return nullptr;
   }
-  if (!linalg::isaContractionOpInterface(linalgOp) &&
-      !IREE::LinalgExt::isaScaledContractionOpInterface(linalgOp)) {
-    return nullptr;
-  }
 
   SmallVector<Value> inputs = linalgOp.getDpsInputs();
   SmallVector<Value> outputs = linalgOp.getDpsInits();
@@ -517,8 +511,21 @@ struct GPUEncodingResolverMaterializerAttr
     if (!linalgOp) {
       return nullptr;
     }
-    return lowerContractionOrScaledContractionOpToInnerTiledOp(
-        b, linalgOp, convertedOperands, resolverAttr);
+    if (auto fillOp = dyn_cast<linalg::FillOp>(op)) {
+      return lowerFillOpWithResolvedLayouts(b, fillOp, convertedResTypes,
+                                            convertedOperands);
+    }
+    if (linalg::isaContractionOpInterface(linalgOp) ||
+        IREE::LinalgExt::isaScaledContractionOpInterface(linalgOp)) {
+      return lowerContractionOrScaledContractionOpToInnerTiledOp(
+          b, linalgOp, convertedOperands, resolverAttr);
+    }
+    if (auto genericOp = dyn_cast<linalg::GenericOp>(op)) {
+      return lowerGenericOpWithResolvedLayouts(
+          b, genericOp, convertedResTypes, convertedOperands,
+          cast<IREE::Encoding::LayoutMaterializerAttr>(attr));
+    }
+    return nullptr;
   }
 };