[DT][NFC] Collapse MaterializeScaledContractionOp into generic pattern. (#22340)

The op lowering is driven by interface methods -- `lowerOp`, so we don't
need a separate pattern for scaled matmul cases.


https://github.com/iree-org/iree/blob/d4d74cb2f1f71fe8b3aaf7a41e5d5895cb726413/compiler/src/iree/compiler/Codegen/ExternalInterfaces/GPUEncodingExternalModels.cpp#L509-L523

It is a step towards #20160

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

Author: hanhanW

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

@@ -13,7 +13,6 @@
 #include "iree/compiler/Codegen/Utils/EncodingUtils.h"
 #include "iree/compiler/Codegen/Utils/Utils.h"
 #include "iree/compiler/Dialect/Encoding/IR/EncodingOps.h"
-#include "iree/compiler/Dialect/HAL/IR/HALTypes.h"
 #include "iree/compiler/Dialect/LinalgExt/Utils/MatchUtils.h"
 #include "iree/compiler/Dialect/TensorExt/IR/TensorExtOps.h"
 #include "iree/compiler/Dialect/Util/IR/UtilOps.h"
@@ -858,9 +857,13 @@ class MaterializeContractionOp
   LogicalResult
   matchAndRewrite(linalg::LinalgOp op, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const override {
-    if (!linalg::isaContractionOpInterface(op)) {
+    // TODO(hanchung): Remove the check after moving other ops, e.g., fill,
+    // generic, etc, lowering patterns to interface implementation.
+    if (!linalg::isaContractionOpInterface(op) &&
+        !IREE::LinalgExt::isaScaledContractionOpInterface(op)) {
       return rewriter.notifyMatchFailure(
-          op, "does not implement ContractionOpInterface");
+          op, "does not match linalg::isaContractionOpInterface and "
+              "LinalgExt::isaScaledContractionOpInterface");
     }
 
     auto converter = static_cast<const MaterializeEncodingTypeConverter *>(
@@ -879,43 +882,6 @@ class MaterializeContractionOp
   }
 };
 
-/// Pattern to convert scaled contraction operations.
-class MaterializeScaledContractionOp
-    : public OpInterfaceConversionPattern<linalg::LinalgOp> {
-public:
-  MaterializeScaledContractionOp(
-      const MaterializeEncodingTypeConverter &typeConverter,
-      MLIRContext *context, PatternBenefit benefit = 1)
-      : OpInterfaceConversionPattern<linalg::LinalgOp>(typeConverter, context,
-                                                       benefit) {}
-
-  LogicalResult
-  matchAndRewrite(linalg::LinalgOp op, ArrayRef<Value> operands,
-                  ConversionPatternRewriter &rewriter) const override {
-    if (!IREE::LinalgExt::isaScaledContractionOpInterface(op)) {
-      return rewriter.notifyMatchFailure(
-          op, "does not implement ScaledContractionOpInterface");
-    }
-
-    auto converter = static_cast<const MaterializeEncodingTypeConverter *>(
-        this->getTypeConverter());
-
-    IREE::Encoding::LayoutMaterializerAttr layoutAttr =
-        converter->getLayoutAttr();
-    SmallVector<Type> convertedResTypes;
-    for (Value init : op.getDpsInits()) {
-      convertedResTypes.push_back(converter->convertType(init.getType()));
-    }
-    Operation *newOp =
-        layoutAttr.lowerOp(rewriter, op, convertedResTypes, operands);
-    if (!newOp) {
-      return failure();
-    }
-    rewriter.replaceOp(op, newOp->getResults());
-    return success();
-  }
-};
-
 static bool isRankedTensorTypeWithEncoding(Type type) {
   auto rankedTensorType = dyn_cast<RankedTensorType>(type);
   if (!rankedTensorType) {
@@ -975,15 +941,15 @@ void populateMaterializeEncodingPatterns(
                          isRankedTensorTypeWithEncoding);
   });
 
-  patterns.insert<
-      MaterializeContractionOp, MaterializeScaledContractionOp,
-      SetEncodingOpLoweringConversion, UnsetEncodingOpLoweringConversion,
-      MaterializeDPSOperation<linalg::FillOp>,
-      MaterializeDPSOperation<linalg::GenericOp>,
-      MaterializeOperation<tensor::EmptyOp>, MaterializeOptimizationBarrierOp,
-      MaterializeTensorExtDispatchTensorLoadOp,
-      MaterializeTensorExtDispatchTensorStoreOp,
-      MaterializeInterfaceBindingEncoding, MaterializeFuncReturnOp>(
+  patterns.insert<MaterializeContractionOp, SetEncodingOpLoweringConversion,
+                  UnsetEncodingOpLoweringConversion,
+                  MaterializeDPSOperation<linalg::FillOp>,
+                  MaterializeDPSOperation<linalg::GenericOp>,
+                  MaterializeOperation<tensor::EmptyOp>,
+                  MaterializeOptimizationBarrierOp,
+                  MaterializeTensorExtDispatchTensorLoadOp,
+                  MaterializeTensorExtDispatchTensorStoreOp,
+                  MaterializeInterfaceBindingEncoding, MaterializeFuncReturnOp>(
       typeConverter, context);
 };