Fix transform test

Signed-off-by: hanhanW <hanhan0912@gmail.com>

Author: hanhanW

compiler/src/iree/compiler/Codegen/Common/TransformExtensions/CommonExtensions.cpp

@@ -12,6 +12,7 @@
 #include "iree/compiler/Codegen/Common/GPU/GPUVectorDistribution.h"
 #include "iree/compiler/Codegen/Common/GPU/Passes.h"
 #include "iree/compiler/Codegen/Common/Transforms.h"
+#include "iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenOps.h"
 #include "iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUAttrs.h"
 #include "iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUOps.h"
 #include "iree/compiler/Codegen/Dialect/GPU/Transforms/Transforms.h"
@@ -673,11 +674,42 @@ transform_dialect::PopulateWorkgroupCountRegionUsingNumThreadsSliceOp::
   }
 
   auto funcOp = forAllOp->getParentOfType<mlir::FunctionOpInterface>();
-  if (failed(
-          lowerWorkgroupCountFromSliceOp(rewriter, funcOp, workgroupCount))) {
+
+  // Resolve the workgroup_count_from_slice in the dispatch_config op.
+  // The dispatch_config is the codegen-level representation of the workgroup
+  // count — codegen passes no longer touch the hal.executable.export directly.
+  // PropagateDispatchConfig copies the result to the export later.
+  auto moduleOp = funcOp->getParentOfType<ModuleOp>();
+  if (!moduleOp) {
+    return mlir::emitDefiniteFailure(state.getTopLevel(),
+                                     "could not find parent module");
+  }
+  IREE::Codegen::DispatchConfigOp configOp;
+  for (auto op : moduleOp.getOps<IREE::Codegen::DispatchConfigOp>()) {
+    if (op.getFunctionRef() == funcOp.getName()) {
+      configOp = op;
+      break;
+    }
+  }
+  if (!configOp) {
     return mlir::emitDefiniteFailure(state.getTopLevel(),
-                                     "failed to lower workgroup count region");
+                                     "could not find dispatch_config for '")
+           << funcOp.getName() << "'";
   }
+  IREE::TensorExt::DispatchWorkgroupCountFromSliceOp fromSliceOp;
+  configOp.getBody().walk(
+      [&](IREE::TensorExt::DispatchWorkgroupCountFromSliceOp fs) {
+        fromSliceOp = fs;
+        return WalkResult::interrupt();
+      });
+  if (fromSliceOp) {
+    if (failed(lowerWorkgroupCountFromSliceOp(rewriter, fromSliceOp, funcOp,
+                                              workgroupCount))) {
+      return mlir::emitDefiniteFailure(
+          state.getTopLevel(), "failed to lower dispatch_config count region");
+    }
+  }
+
   return DiagnosedSilenceableFailure::success();
 }
 

compiler/src/iree/compiler/Codegen/Common/TransformExtensions/CommonExtensionsOps.td

@@ -500,15 +500,13 @@ def PopulateWorkgroupCountRegionUsingNumThreadsSliceOp :
        TransformOpInterface,
        ReportTrackingListenerFailuresOpTrait]> {
   let description = [{
-    Populate the workgroup_count region on the `hal.executable.export` op.
-
-    The default dispatch region formation expects that the workgroup count
-    be computed from within the dispatch by using a program slice. The utility
-    method `lowerWorkgroupCountFromSliceOp` handles populating the
-    workgroup count region given the values in the dispatch that represent the
-    number of workgroups. This transform op calls the underlying function using
-    the `num_threads` value from the `scf.for_all` op that distributes the work
-    to different workgroups.
+    Populate the `iree_codegen.dispatch_config` workgroup count region.
+
+    Finds the `dispatch_config` op for the parent function and resolves the
+    `workgroup_count_from_slice` placeholder in its body using the
+    `num_threads` values from the given `scf.forall` op.
+    `PropagateDispatchConfig` later copies the result to the
+    `hal.executable.export` count region.
   }];
 
   let arguments = (ins TransformHandleTypeInterface:$for_all_op);

compiler/src/iree/compiler/Codegen/Transforms/Transforms.cpp

@@ -521,35 +521,6 @@ LogicalResult lowerWorkgroupCountFromSliceOp(
   return success();
 }
 
-LogicalResult lowerWorkgroupCountFromSliceOp(
-    RewriterBase &rewriter, mlir::FunctionOpInterface entryPointFn,
-    ArrayRef<OpFoldResult> workgroupCount, int maxWorkgroupParallelDims) {
-  std::optional<IREE::HAL::ExecutableExportOp> exportOp =
-      getEntryPoint(entryPointFn);
-  if (!exportOp) {
-    return success();
-  }
-  Block *body = exportOp->getWorkgroupCountBody();
-  if (!body) {
-    return success();
-  }
-  auto countOps =
-      body->getOps<IREE::TensorExt::DispatchWorkgroupCountFromSliceOp>();
-  if (countOps.empty()) {
-    // If there are no `flow.dispatch.workgroup_count_default` operations
-    // do nothing.
-    return success();
-  }
-  if (!llvm::hasSingleElement(countOps)) {
-    return exportOp->emitOpError(
-        "unexpected multiple flow.dispatch.workgroup_count_default operations "
-        "in body");
-  }
-  return lowerWorkgroupCountFromSliceOp(rewriter, *countOps.begin(),
-                                        entryPointFn, workgroupCount,
-                                        maxWorkgroupParallelDims);
-}
-
 LogicalResult createWorkgroupCountHint(RewriterBase &rewriter, Location loc,
                                        ArrayRef<OpFoldResult> workgroupCount,
                                        int maxWorkgroupParallelDims,

compiler/src/iree/compiler/Codegen/Transforms/Transforms.h

@@ -145,35 +145,18 @@ FailureOr<SmallVector<OpFoldResult>> materializeWorkgroupCountComputation(
     ArrayRef<OpFoldResult> workgroupCount, ValueRange workloadVals);
 
 /// Lower the workgroup count region for the default code-generation path in
-/// IREE. Given the list `workgroupCount` (fastest varying dimension innermost)
-/// as computed within the `entryPointFn`, clones a backward slice of the
-/// computation starting at these values and ending with
-/// `flow.dispatch.constant_ordinal` into the workgroup count region on the
-/// `hal.executable.export` op corresponding to the `entryPointFn`. Also removes
-/// the `flow.dispatch.constant_ordinal` operations from within the
-/// `entryPointFn`. Expects the workgroup count region of the corresponding
-/// `hal.executable.export` to contain the
-/// `flow.dispatch.workgroup_count_slice` operation as a placeholder for the
-/// computation to compute the number of workgroups. In absence of this
-/// operation, this method does nothing assuming that the workgroup count
-/// computation has already been resolved.
+/// Replaces a `workgroup_count_from_slice` placeholder with a materialized
+/// workgroup count computation. The `workgroupCount` list (fastest varying
+/// dimension innermost) provides the desired counts. If there are more
+/// dimensions than `maxWorkgroupParallelDims`, excess dimensions are folded
+/// into the last parallel dimension. Remaining dimensions are padded with 1.
 LogicalResult lowerWorkgroupCountFromSliceOp(
     RewriterBase &rewriter,
     IREE::TensorExt::DispatchWorkgroupCountFromSliceOp workgroupCountOp,
     mlir::FunctionOpInterface entryPointFn,
     ArrayRef<OpFoldResult> workgroupCount,
     int maxWorkgroupParallelDims = kNumMaxParallelDims);
 
-/// Wrapper around `lowerWorkgroupCountFromSliceOp` method that
-/// takes the `iree_tensor_ext.dispatch.workgroup_count_from_slice` op
-/// as an argument. Looks up the `hal.executable.export` operation
-/// and finds the `iree_tensor_ext.dispatch.workgroup_count_from_slice` op to
-/// lower.
-LogicalResult lowerWorkgroupCountFromSliceOp(
-    RewriterBase &rewriter, mlir::FunctionOpInterface entryPointFn,
-    ArrayRef<OpFoldResult> workgroupCount,
-    int maxWorkgroupParallelDims = kNumMaxParallelDims);
-
 /// Creates an `iree_codegen.workgroup_count_hint` op at the current insertion
 /// point with the provided operands. If there are more operands provided than
 /// |maxWorkgroupParallelDims| the outermost sizes are linearized into the