update matmul config analysis

Author: yifeizh2

include/gc/Analysis/GlobalAnalysis.h

@@ -112,6 +112,15 @@ class OperatorLayout {
     return supportedOutputLayouts[idx];
   }
 
+  bool isPlain() const {
+    for (const auto &layout : llvm::concat<const TensorLayout>(
+             supportedInputLayouts, supportedOutputLayouts)) {
+      if (!layout.isPlainLayout())
+        return false;
+    }
+    return true;
+  }
+
   friend llvm::raw_ostream &operator<<(llvm::raw_ostream &ss,
                                        const OperatorLayout &opLayout);
 

lib/gc/Analysis/GlobalAnalysis.cpp

@@ -8,6 +8,7 @@
 #include <memory>
 
 #include "gc/Analysis/GlobalAnalysis.h"
+#include "gc/Analysis/MatmulConfigAnalysis.h"
 
 namespace mlir {
 namespace gc {
@@ -176,25 +177,26 @@ GlobalAnalysis::GlobalAnalysis(Operation *root) {
       IRRewriter rewriter(linalgOp);
       if (mlir::linalg::isaContractionOpInterface(linalgOp)) {
         // query the cost model
-        // OperatorLayout suggestedLayout = costModel->queryLayout(linalgOp,
-        // curInputLayouts);
-
-        // hardcode one for now
-        // A side layoutCache, [0, 1, 0, 1]; {32, 32}
+        MatmulConfig cfg =
+            MatmulConfigAnalysis(linalgOp.getOperation()).getConfig();
+        uint32_t iim = cfg.innerMostKBlock, iin = cfg.innerMostNBlock,
+                 iik = cfg.innerMostKBlock;
+        // hardcode outer axis for now
+        // A side layoutCache, [0, 1, 0, 1]; {iim, iik}
         TensorLayout A_layout(
             {0, 1}, {0, 1},
-            SmallVector<OpFoldResult>{rewriter.getIndexAttr(32),
-                                      rewriter.getIndexAttr(32)});
-        // B side layoutCache, [1, 0, 0, 1]; {32, 32}
+            SmallVector<OpFoldResult>{rewriter.getIndexAttr(iim),
+                                      rewriter.getIndexAttr(iik)});
+        // B side layoutCache, [1, 0, 0, 1]; {iik, iin}
         TensorLayout B_layout(
             {1, 0}, {0, 1},
-            SmallVector<OpFoldResult>{rewriter.getIndexAttr(32),
-                                      rewriter.getIndexAttr(32)});
-        // C side layoutCache, [0, 1, 0, 1]; {32, 32}
+            SmallVector<OpFoldResult>{rewriter.getIndexAttr(iik),
+                                      rewriter.getIndexAttr(iin)});
+        // C side layoutCache, [0, 1, 0, 1]; {iim, iin}
         TensorLayout C_layout(
             {0, 1}, {0, 1},
-            SmallVector<OpFoldResult>{rewriter.getIndexAttr(32),
-                                      rewriter.getIndexAttr(32)});
+            SmallVector<OpFoldResult>{rewriter.getIndexAttr(iim),
+                                      rewriter.getIndexAttr(iin)});
         OperatorLayout suggestedLayout({A_layout, B_layout}, {C_layout});
         layoutCache[linalgOp] = suggestedLayout;
       } else {