align constraint

Author: yifeizh2

include/gc/Analysis/MatmulConfigAnalysis.h

@@ -33,7 +33,8 @@ struct MatmulConfig {
   uint32_t innerMostMBlock, innerMostNBlock, innerMostKBlock;
 };
 
-bool validateConfig(const MatmulConfig &cfg, ArrayRef<uint32_t> shape = {});
+bool validateConfig(const MatmulConfig &cfg, ArrayRef<uint32_t> shape,
+                    bool allowIndivisibleInnerblock, bool isVNNIMM2D);
 
 enum DimType { Batch, M, N, K };
 

lib/gc/Analysis/MatmulConfigAnalysis.cpp

@@ -37,22 +37,29 @@ static llvm::raw_ostream &operator<<(llvm::raw_ostream &ss,
   return ss;
 }
 
-bool validateConfig(const MatmulConfig &cfg, ArrayRef<uint32_t> shape) {
+bool validateConfig(const MatmulConfig &cfg, ArrayRef<uint32_t> shape,
+                    bool allowIndivisibleInnerblock, bool isVNNIMM2D) {
   if (cfg.MThreads <= 0 || cfg.NThreads <= 0 || cfg.KThreads <= 0 ||
       cfg.MBlock <= 0 || cfg.NBlock <= 0 || cfg.KBlock <= 0 ||
       cfg.innerMostMBlock <= 0 || cfg.innerMostNBlock <= 0 ||
       cfg.innerMostKBlock <= 0)
     return false;
   if (cfg.MBlock % cfg.innerMostMBlock != 0 ||
-      cfg.NBlock % cfg.innerMostNBlock != 0 ||
-      cfg.KBlock % cfg.innerMostKBlock != 0)
+      (shape[0] % cfg.innerMostMBlock != 0 && !allowIndivisibleInnerblock))
+    return false;
+  if (cfg.NBlock % cfg.innerMostNBlock != 0 ||
+      ((shape[1] % cfg.innerMostNBlock != 0) && !allowIndivisibleInnerblock) ||
+      (shape[1] % cfg.NThreads != 0 && isVNNIMM2D &&
+       cfg.NBlock != cfg.innerMostNBlock))
+    return false;
+  if (cfg.KBlock % cfg.innerMostKBlock != 0 ||
+      ((shape[2] / cfg.KThreads % cfg.KBlock != 0 ||
+        shape[2] / cfg.KThreads % cfg.innerMostKBlock != 0) &&
+       !allowIndivisibleInnerblock))
+    return false;
+  // KThreads will not shrink automatically
+  if (llvm::divideCeil(shape[2], cfg.KBlock) < cfg.KThreads)
     return false;
-  if (!shape.empty()) {
-    // KThreads will not shrink automatically
-    // K is shape[2]
-    if (llvm::divideCeil(shape[2], cfg.KBlock) < cfg.KThreads)
-      return false;
-  }
   return true;
 }
 
@@ -185,7 +192,6 @@ double dynamicBufferizationCost(linalg::LinalgOp &linalgOp,
                                 ArrayRef<uint32_t> shape,
                                 const MatmulConfig &config,
                                 CPUTargetDescriptionAnalysis &sysDesc) {
-  assert(validateConfig(config, shape) && "config is invalid");
   assert(shape.size() >= 3 && "shape.size() should >= 3");
   uint32_t M = shape[0], N = shape[1];
   double cost = 0;
@@ -367,8 +373,7 @@ prepareConfigCandidates(Operation *root, CPUTargetDescriptionAnalysis &sysDesc,
 }
 
 // read the config from the attributes for tuning
-bool readConfigFromAttrs(MatmulConfig &config, ArrayRef<NamedAttribute> attrs,
-                         ArrayRef<uint32_t> shape) {
+bool readConfigFromAttrs(MatmulConfig &config, ArrayRef<NamedAttribute> attrs) {
   size_t cfgItemCnt = 0;
   for (const auto &attr : attrs) {
     if (attr.getName() == "KBlock") {
@@ -400,17 +405,28 @@ bool readConfigFromAttrs(MatmulConfig &config, ArrayRef<NamedAttribute> attrs,
       cfgItemCnt++;
     }
   }
-  if (cfgItemCnt != 9) {
-    LLVM_DEBUG(llvm::dbgs() << "The predefined matmul config is incomplete. "
-                               "Default matmul config will be set.\n");
+  return cfgItemCnt == 9;
+}
+
+bool readAndValidateConfig(MatmulConfig &config,
+                           const linalg::LinalgOp &linalgOp,
+                           ArrayRef<uint32_t> shape,
+                           bool allowIndivisibleInnerBlock) {
+  SmallVector<NamedAttribute> attrs(linalgOp->getAttrs());
+  bool fullConfig = readConfigFromAttrs(config, attrs);
+  if (!fullConfig) {
+    LLVM_DEBUG(llvm::dbgs() << "Missing fields in predefined config.\n");
     return false;
   }
-  if (validateConfig(config, shape))
-    return true;
-  else {
-    assert(0 && "config is invalid");
+  bool validConfig =
+      validateConfig(config, shape, allowIndivisibleInnerBlock,
+                     linalgx::isGenericPackedMatmulOp(
+                         linalgOp, linalgx::PackingType::VNNI_MM2D));
+  if (!validConfig) {
+    LLVM_DEBUG(llvm::dbgs() << "Invalid predefined config.\n");
     return false;
   }
+  return true;
 }
 
 // Analyze the workload and system description to generate the default config
@@ -494,13 +510,15 @@ MatmulConfig MatmulConfigAnalysis::getConfig() {
                  << "M: " << M << ", N: " << N << ", K: " << K << "\n");
 
       // try to read the config from the attributes
-      SmallVector<NamedAttribute> attrs(linalgOp->getAttrs());
-      bool hasPredefinedConfig =
-          readConfigFromAttrs(config, attrs, SmallVector<uint32_t>{M, N, K});
+      bool hasValidPredefinedConfig = readAndValidateConfig(
+          config, linalgOp, SmallVector<uint32_t>{M, N, K},
+          allowIndivisibleInnerBlock);
 
       // if there is a given config, skip the cost model
-      if (!hasPredefinedConfig) {
-        LLVM_DEBUG(llvm::dbgs() << "No predefined config\n");
+      if (!hasValidPredefinedConfig) {
+        LLVM_DEBUG(
+            llvm::dbgs()
+            << "No valid predefined config. Setting with default config.\n");
         // TODO: Could add a weight or priority for cost model
         SmallVector<std::tuple<CostModelFn, std::string, double>>
             costModelList = {
@@ -525,7 +543,10 @@ MatmulConfig MatmulConfigAnalysis::getConfig() {
         if (!configCandidates.empty())
           config = configCandidates[0];
 
-        assert(validateConfig(config, shape) && "config is invalid");
+        assert(validateConfig(config, shape, allowIndivisibleInnerBlock,
+                              linalgx::isGenericPackedMatmulOp(
+                                  root, linalgx::PackingType::VNNI_MM2D)) &&
+               "config is invalid");
       }
 
       LLVM_DEBUG(llvm::dbgs()