[Codegen][PCF] Add convertForallToGenericNest transform

Add a transform that converts scf.forall operations into multi-level
pcf.generic nests. This is required in cases where a single scf.forall
mapping type needs to map to multiple scopes. The immediate case this
arises from is converting thread mapped scf.forall ops to combined
subgroup + lane scopes. In this case, we can't simply convert to
pcf.loop because of the way that automatic redistribution of pcf.loop's
lowering works. We need to redistribute the iterations of the scf.forall
to *all* workers across both scopes. This means that if we make one of
the two scopes (subgroup or lane) a pcf.loop this fails to create the
correct IR structure. If we make it loop over subgroups, this fails to
predicate the lanes, and if we make it loop over lanes, we end up with
more thread divergence than normal (normally only a single subgroup
may exhibit divergence).

The transform creates an outer loop over workers and an inner
scf.forall over the per-worker iteration range. For multi-dimensional
cases, affine.linearize_index and affine.delinearize_index are used
to flatten/unflatten indices appropriately.

Additionally adds a new method getNativeNumProcessorIds to ScopeAttrInterface.
This is needed to query for the number of ids to generate.

Author: qedawkins

compiler/src/iree/compiler/Codegen/Common/BUILD.bazel

@@ -161,6 +161,7 @@ iree_compiler_cc_library(
         "TensorToVectorVectorizePad.cpp",
         "TestExecutablePreprocessing.cpp",
         "TestPartitionableLoopsInterface.cpp",
+        "ConvertForallToGenericNestWorkgroup.cpp",
         "TileAndDistributeToWorkgroupsPass.cpp",
         "TileAndFuseUtils.cpp",
         "TileDispatchUsingForall.cpp",

compiler/src/iree/compiler/Codegen/Common/CMakeLists.txt

@@ -154,6 +154,7 @@ iree_cc_library(
     "TensorToVectorVectorizePad.cpp"
     "TestExecutablePreprocessing.cpp"
     "TestPartitionableLoopsInterface.cpp"
+    "ConvertForallToGenericNestWorkgroup.cpp"
     "TileAndDistributeToWorkgroupsPass.cpp"
     "TileAndFuseUtils.cpp"
     "TileDispatchUsingForall.cpp"

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

@@ -0,0 +1,70 @@
+// Copyright 2026 The IREE Authors
+//
+// Licensed under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+#include "iree/compiler/Codegen/Common/Passes.h"
+#include "iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenAttrs.h"
+#include "iree/compiler/Codegen/Dialect/PCF/IR/PCF.h"
+#include "iree/compiler/Codegen/Dialect/PCF/Transforms/Transforms.h"
+#include "mlir/Dialect/SCF/IR/SCF.h"
+#include "mlir/IR/PatternMatch.h"
+
+namespace mlir::iree_compiler {
+
+#define GEN_PASS_DEF_CONVERTFORALLTOGENERICNESTWORKGROUPPASS
+#include "iree/compiler/Codegen/Common/Passes.h.inc"
+
+namespace {
+
+/// Returns true if the forall op has WorkgroupMappingAttr mapping attributes.
+static bool hasWorkgroupMapping(scf::ForallOp forallOp) {
+  std::optional<ArrayAttr> mapping = forallOp.getMapping();
+  if (!mapping || mapping->empty()) {
+    return false;
+  }
+  return llvm::all_of(mapping.value(),
+                      llvm::IsaPred<IREE::Codegen::WorkgroupMappingAttr>);
+}
+
+struct ConvertForallToGenericNestWorkgroupPass final
+    : public impl::ConvertForallToGenericNestWorkgroupPassBase<
+          ConvertForallToGenericNestWorkgroupPass> {
+  using Base::Base;
+
+  void runOnOperation() override {
+    MLIRContext *ctx = &getContext();
+
+    // Always use linearized workgroup scope (1 id).
+    // Interface is implemented via external models hence the cast.
+    auto scope = cast<IREE::PCF::ScopeAttrInterface>(
+        IREE::Codegen::WorkgroupScopeAttr::get(ctx, /*linearize=*/true));
+
+    SmallVector<IREE::PCF::ScopeAttrInterface> scopes = {scope};
+
+    IRRewriter rewriter(ctx);
+    SmallVector<scf::ForallOp> forallOps;
+    getOperation()->walk([&](scf::ForallOp forallOp) {
+      // Only convert foralls with workgroup mapping attributes.
+      if (hasWorkgroupMapping(forallOp)) {
+        forallOps.push_back(forallOp);
+      }
+    });
+
+    for (scf::ForallOp forallOp : forallOps) {
+      rewriter.setInsertionPoint(forallOp);
+      FailureOr<IREE::PCF::GenericOp> result =
+          IREE::PCF::convertForallToGenericNest(rewriter, forallOp, scopes);
+      if (failed(result)) {
+        forallOp.emitError("failed to convert forall to generic nest");
+        return signalPassFailure();
+      }
+      // Replace forall results with generic results.
+      rewriter.replaceOp(forallOp, result->getResults());
+    }
+  }
+};
+
+} // namespace
+} // namespace mlir::iree_compiler

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

@@ -53,7 +53,8 @@ ConvertWorkgroupForall::matchAndRewrite(scf::ForallOp op,
   auto scope = cast<IREE::PCF::ScopeAttrInterface>(
       IREE::Codegen::WorkgroupScopeAttr::get(rewriter.getContext(),
                                              /*linearize=*/true));
-  FailureOr<IREE::PCF::LoopOp> res = convertForallToPCF(rewriter, op, scope, 1);
+  FailureOr<IREE::PCF::LoopOp> res =
+      convertForallToPCFLoop(rewriter, op, scope, 1);
   if (failed(res)) {
     return failure();
   }

compiler/src/iree/compiler/Codegen/Common/Passes.h

@@ -18,6 +18,7 @@
 #include "iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenAttrs.h"
 #include "iree/compiler/Codegen/Utils/Utils.h"
 #include "iree/compiler/Dialect/LinalgExt/IR/LinalgExtOps.h"
+#include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h"
 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
 #include "mlir/Dialect/Transform/IR/TransformOps.h"

compiler/src/iree/compiler/Codegen/Common/Passes.td

@@ -155,6 +155,23 @@ def ConvertWorkgroupForallToPCFPass
   let dependentDialects = ["iree_compiler::IREE::PCF::PCFDialect"];
 }
 
+def ConvertForallToGenericNestWorkgroupPass
+    : InterfacePass<"iree-codegen-convert-forall-to-generic-nest-workgroup",
+                    "mlir::FunctionOpInterface"> {
+  let summary = "Converts scf.forall ops with workgroup mapping to pcf.generic";
+  let description = [{
+    Converts `scf.forall` ops with `#iree_codegen.workgroup_mapping` attributes
+    to a `pcf.generic` op using workgroup scope. The pass always linearizes
+    workgroup IDs to a single dimension.
+  }];
+  let dependentDialects = [
+    "iree_compiler::IREE::PCF::PCFDialect",
+    "arith::ArithDialect",
+    "affine::AffineDialect",
+    "scf::SCFDialect"
+  ];
+}
+
 def CombineLayoutTransformationPass :
     InterfacePass<"iree-codegen-combine-layout-transformation", "mlir::FunctionOpInterface"> {
   let summary =

compiler/src/iree/compiler/Codegen/Common/test/BUILD.bazel

@@ -33,6 +33,7 @@ iree_lit_test_suite(
             "convert_accgemm_to_gemm.mlir",
             "convert_bf16_arith_to_f32.mlir",
             "convert_bf16_to_uint16_buffers.mlir",
+            "convert_forall_to_generic_nest_workgroup.mlir",
             "convert_hal_descriptor_type_to_gpu_address_space.mlir",
             "convert_to_destination_passing_style.mlir",
             "convert_unsupported_float_arith.mlir",

compiler/src/iree/compiler/Codegen/Common/test/CMakeLists.txt

@@ -28,6 +28,7 @@ iree_lit_test_suite(
     "convert_accgemm_to_gemm.mlir"
     "convert_bf16_arith_to_f32.mlir"
     "convert_bf16_to_uint16_buffers.mlir"
+    "convert_forall_to_generic_nest_workgroup.mlir"
     "convert_hal_descriptor_type_to_gpu_address_space.mlir"
     "convert_to_destination_passing_style.mlir"
     "convert_unsupported_float_arith.mlir"

compiler/src/iree/compiler/Codegen/Common/test/convert_forall_to_generic_nest_workgroup.mlir

@@ -0,0 +1,29 @@
+// RUN: iree-opt %s --pass-pipeline="builtin.module(func.func(iree-codegen-convert-forall-to-generic-nest-workgroup))" --allow-unregistered-dialect --split-input-file | FileCheck %s
+
+// Test that workgroup scope creates 1 id/count pair with linearized scope.
+
+// CHECK-LABEL: func.func @test_workgroup_scope
+// CHECK:       pcf.generic
+// CHECK-SAME:    scope(#iree_codegen.workgroup_scope<linearize>)
+// CHECK:         execute(%{{.*}})[%[[ID:.+]]: index, %[[COUNT:.+]]: index]
+// Chunk size computed from total iterations / worker count.
+// CHECK:         %[[CHUNK:.+]] = arith.ceildivui
+// Start = id * chunk_size.
+// CHECK:         %[[START:.+]] = arith.muli %[[ID]], %[[CHUNK]]
+// End = min(start + chunk_size, total).
+// CHECK:         %[[END_RAW:.+]] = arith.addi %[[START]], %[[CHUNK]]
+// CHECK:         %[[END:.+]] = arith.minui %[[END_RAW]]
+// CHECK:         scf.forall (%[[IV:.+]]) = (%[[START]]) to (%[[END]])
+// CHECK:           "foo.body"(%[[IV]])
+// CHECK:           pcf.write_slice
+// CHECK:         pcf.return
+func.func @test_workgroup_scope(%init: tensor<64xf32>) -> tensor<64xf32> {
+  %result = scf.forall (%i) in (64) shared_outs(%out = %init) -> tensor<64xf32> {
+    "foo.body"(%i) : (index) -> ()
+    %slice = tensor.extract_slice %out[%i] [1] [1] : tensor<64xf32> to tensor<1xf32>
+    scf.forall.in_parallel {
+      tensor.parallel_insert_slice %slice into %out[%i] [1] [1] : tensor<1xf32> into tensor<64xf32>
+    }
+  } {mapping = [#iree_codegen.workgroup_mapping<x>]}
+  return %result : tensor<64xf32>
+}

compiler/src/iree/compiler/Codegen/Dialect/GPU/ExternalInterfaces/GPUScopeExternalModels.cpp

@@ -83,6 +83,11 @@ struct SubgroupScopeModel
                                         MLIRContext *context) const {
     return gpu::AddressSpaceAttr::get(context, gpu::AddressSpace::Workgroup);
   }
+
+  int64_t getNativeNumProcessorIds(Attribute attr) const {
+    // SubgroupScope natively provides a single 1D processor ID (subgroup_id).
+    return 1;
+  }
 };
 
 /// External model for LaneScopeAttr implementing ScopeAttrInterface.
@@ -133,6 +138,11 @@ struct LaneScopeModel
     // logic to allocate + subview.
     return failure();
   }
+
+  int64_t getNativeNumProcessorIds(Attribute attr) const {
+    // LaneScope natively provides a single 1D processor ID (lane_id).
+    return 1;
+  }
 };
 
 } // namespace