[Codegen] Check for workgroup level tile sizes in workgroup tiling (#…

…18538)

TileDispatchUsingForall relies on lowering configurations having
workgroup level tile sizes, so this PR adds the additional check that
the tilableOp has workgroup level tile sizes. It also adds verification
that there is only one op with a workgroup tiling level.

---------

Signed-off-by: Max Dawkins <max.dawkins@gmail.com>

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

@@ -53,11 +53,18 @@ struct TilingInfo {
 static FailureOr<TilingInfo>
 getTiledAndDistributionInfo(RewriterBase &rewriter,
                             ArrayRef<Operation *> computeOps) {
+  // It is expected that at most one compute op has a workgroup tiling level.
   Operation *tilableOp = nullptr;
   for (Operation *op : llvm::reverse(computeOps)) {
     if (getLoweringConfig(op)) {
+      if (!getLoweringConfig(op).hasWorkgroupTilingLevel()) {
+        continue;
+      }
+      if (tilableOp) {
+        return op->emitOpError("expected only one op with a workgroup tiling"
+                               "level.");
+      }
       tilableOp = op;
-      break;
     }
   }
   if (!tilableOp) {

compiler/src/iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenAttrs.cpp

@@ -323,6 +323,10 @@ bool LoweringConfigAttr::hasTilingLevel(unsigned level) const {
   return !getTileSizeVals(level).empty();
 }
 
+bool LoweringConfigAttr::hasWorkgroupTilingLevel() const {
+  return !getWorkgroupTileSizes().empty();
+}
+
 LogicalResult
 LoweringConfigAttr::verify(function_ref<InFlightDiagnostic()> emitError,
                            LoweringConfigTilingLevelsAttr levels,

compiler/src/iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenAttrs.td

@@ -296,6 +296,7 @@ def IREECodegen_LoweringConfigAttr :
         "getStaticTilingLevelSizes",
         "getTilingLevelSizes",
         "hasTilingLevel",
+        "hasWorkgroupTilingLevel",
       ]>
     ]> {
   let mnemonic = "lowering_config";

compiler/src/iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenInterfaces.td

@@ -60,6 +60,19 @@ def IREECodegen_LoweringConfigAttrInterface :
         return false;
       }]
     >,
+    InterfaceMethod<
+      /*desc=*/[{
+        Returns true if the lowering config specifies tile sizes for the
+        workgroup tiling level.
+      }],
+      /*retTy=*/"bool",
+      /*methodName=*/"hasWorkgroupTilingLevel",
+      /*args=*/(ins),
+      /*methodBody=*/"",
+      /*defaultImplementation=*/[{
+        return false;
+      }]
+    >,
     InterfaceMethod<
       /*desc=*/[{
         Returns the tile sizes for the specified tiling level. The

compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUAttrs.cpp

@@ -1348,6 +1348,10 @@ bool LoweringConfigAttr::hasTilingLevel(unsigned level) const {
               .empty();
 }
 
+bool LoweringConfigAttr::hasWorkgroupTilingLevel() const {
+  return !getWorkgroupTileSizes().empty();
+}
+
 constexpr StringLiteral kMmaKindName = "mma_kind";
 
 IREE::GPU::MmaInterfaceAttr LoweringConfigAttr::getMmaKind() const {

compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUAttrs.td

@@ -39,6 +39,7 @@ def IREEGPU_LoweringConfigAttr :
         "getStaticTilingLevelSizes",
         "getTilingLevelSizes",
         "hasTilingLevel",
+        "hasWorkgroupTilingLevel",
       ]>
     ]> {
   let mnemonic = "lowering_config";