merge main

Author: xurui1995

cmake/llvm-version.txt

@@ -1 +1 @@
-8345289ded788f4df79f59df76df0c0437c3df64
+f06563a5c0d239a6b98f74db522681613254ad08

include/gc/Analysis/MatmulConfigAnalysis.h

@@ -10,6 +10,7 @@
 #define MLIR_ANALYSIS_MATMULCONFIGANALYSIS_H
 
 #include "gc/Dialect/Linalgx/LinalgxOps.h"
+#include "gc/Dialect/Linalgx/Utils.h"
 #include "mlir/Dialect/DLTI/DLTI.h"
 #include "mlir/Dialect/Linalg/IR/Linalg.h"
 #include "mlir/Interfaces/DataLayoutInterfaces.h"
@@ -55,13 +56,17 @@ getOprandDimType(linalg::LinalgOp &linalgOp) {
         SmallVector<DimType>{DimType::M, DimType::K},
         SmallVector<DimType>{DimType::K, DimType::N},
         SmallVector<DimType>{DimType::M, DimType::N}};
-  } else if (llvm::isa<linalgx::Mm2DVnniOp>(linalgOp)) {
+  } else if (linalgx::isGenericPackedMatmulOp(
+                 linalgOp.getOperation(), linalgx::PackingType::VNNI_MM2D) ||
+             llvm::isa<linalgx::Mm2DVnniOp>(linalgOp)) {
     return SmallVector<SmallVector<DimType>>{
         SmallVector<DimType>{DimType::M, DimType::K},
         SmallVector<DimType>{DimType::N, DimType::K, DimType::K, DimType::N,
                              DimType::K},
         SmallVector<DimType>{DimType::M, DimType::N, DimType::M, DimType::N}};
-  } else if (llvm::isa<linalgx::Mm4DVnniOp>(linalgOp)) {
+  } else if (linalgx::isGenericPackedMatmulOp(
+                 linalgOp.getOperation(), linalgx::PackingType::VNNI_MM4D) ||
+             llvm::isa<linalgx::Mm4DVnniOp>(linalgOp)) {
     return SmallVector<SmallVector<DimType>>{
         SmallVector<DimType>{DimType::M, DimType::K, DimType::M, DimType::K},
         SmallVector<DimType>{DimType::N, DimType::K, DimType::K, DimType::N,
@@ -92,18 +97,34 @@ getOprandDimType(linalg::LinalgOp &linalgOp) {
         SmallVector<DimType>{DimType::Batch, DimType::M, DimType::K},
         SmallVector<DimType>{DimType::Batch, DimType::N, DimType::K},
         SmallVector<DimType>{DimType::Batch, DimType::M, DimType::N}};
+  } else if (linalgx::isGenericPackedMatmulOp(linalgOp.getOperation(),
+                                              linalgx::PackingType::MM4D)) {
+    return SmallVector<SmallVector<DimType>>{
+        SmallVector<DimType>{DimType::M, DimType::K, DimType::M, DimType::K},
+        SmallVector<DimType>{DimType::N, DimType::K, DimType::K, DimType::N},
+        SmallVector<DimType>{DimType::M, DimType::N, DimType::M, DimType::N}};
   }
   return failure();
 }
 
 // The analysis to extract the matmul configuration from the given linalg op
 struct MatmulConfigAnalysis {
 public:
-  explicit MatmulConfigAnalysis(Operation *root);
-  MatmulConfig getConfig() { return config; }
+  // Extract the matmul configuration from the given linalg op
+  MatmulConfigAnalysis(Operation *root) : root(root){};
+
+  // Get the matmul configuration
+  MatmulConfig getConfig();
+
+  void setAllowIndivisibleInnerBlock(bool allow) {
+    allowIndivisibleInnerBlock = allow;
+  }
 
 private:
-  MatmulConfig config;
+  MatmulConfig config = MatmulConfig{1, 1, 1, 1, 1, 1, 1, 1, 1};
+  Operation *root;
+  bool hasConfig = false;
+  bool allowIndivisibleInnerBlock = true;
 };
 
 } // namespace gc

include/gc/Analysis/TargetDescriptionAnalysis.h

@@ -9,10 +9,12 @@
 #ifndef MLIR_ANALYSIS_TARGETDESCRIPTIONANALYSIS_H
 #define MLIR_ANALYSIS_TARGETDESCRIPTIONANALYSIS_H
 
-#include "gc/Dialect/Linalgx/LinalgxOps.h"
 #include "mlir/Dialect/DLTI/DLTI.h"
-#include "mlir/Dialect/Linalg/IR/Linalg.h"
+#include "mlir/IR/BuiltinDialect.h"
+#include "mlir/IR/BuiltinTypes.h"
+#include "mlir/IR/PatternMatch.h"
 #include "mlir/Interfaces/DataLayoutInterfaces.h"
+#include "mlir/Support/LLVM.h"
 #include "llvm/ADT/StringRef.h"
 
 namespace mlir {

include/gc/Transforms/Passes.td

@@ -78,16 +78,17 @@ def IterativeTilingAndFusion : Pass<"iterative-tiling-and-fusion",
            "Set default TileSize for the certain type of op, saying `matmul:{32,32}`">,
     ];
 }
-def DeepTileContractionNamedOp
-    : Pass<"deep-tile-contraction-named-op", "func::FuncOp"> {
-  let summary = "Tile linalg contraction named operation deeply";
+def DeepTileContractionOp
+    : Pass<"deep-tile-contraction-op", "func::FuncOp"> {
+  let summary = "Tile linalg contraction operation deeply";
   let description =
-      [{The pass tries to tile the linalg contraction named op deeply.}];
+      [{The pass tries to tile the linalg contraction op deeply.}];
   let dependentDialects = [
     "func::FuncDialect",
     "arith::ArithDialect",
     "tensor::TensorDialect",
     "linalg::LinalgDialect",
+    "linalgx::LinalgxDialect",
   ];
 }