Fix ReplicateGlobalsPerAffinity to maintain correct order of globals and initializers (iree-org#22401)

## Summary

Fix ordering issue in ReplicateGlobalsPerAffinity pass where new globals
and their initializers were incorrectly inserted immediately after the
original global, breaking dependency order when the original global had
an initializer.

## Problem

The ReplicateGlobalsPerAffinity pass creates per-affinity copies of
globals by inserting new global operations and their initializers right
after the original global operation. This breaks the correct ordering
when the original global has an initializer, as the new initializers
(which load from the original global) would be placed before the
original global's initializer runs.

**Before (incorrect ordering):**
```
global
global_a          <- inserted here
initializer_a     <- loads from global (not yet initialized!)
initializer       <- initializes global
```

**After (correct ordering):**
```
global
initializer       <- initializes global
global_a          <- inserted after initializer
initializer_a     <- can safely load from global
```

## Solution

Modified `ValuePerAffinityHelper::getOrCreateGlobalForAffinity()` to
track and insert new globals after the last initializer that references
the original global:

1. **Implemented constructor caching**: Pre-computes the insertion point
(last initializer) for each global upfront using GlobalTable
2. **Efficient lookup**: Uses GlobalTable's forEach method with direct
access to storeOps collection after calling rebuild()
3. **Correct ordering**: Gets parent initializer directly from store
operations, ensuring new globals are inserted after all initializers
that reference the original global

## Testing

Added comprehensive test case `unknown_global_device_with_initializer`
that verifies:
- Original global and its initializer appear first
- Replicated globals and their initializers follow in correct dependency
order
- Each replicated initializer can safely load from the initialized
original global

## Performance

- O(N+M) complexity: GlobalTable scans globals and their uses once, then
forEach provides efficient iteration
- No redundant IR walking: uses pre-collected storeOps from GlobalTable
- Efficient lookup: cached insertion points avoid repeated scans

Fixes iree-org#22399

---------

Co-authored-by: copilot-swe-agent[bot] <[email protected]>
Co-authored-by: hanhanW <[email protected]>

Author: 2 people

compiler/src/iree/compiler/Dialect/Stream/Transforms/ReplicateGlobalsPerAffinity.cpp

@@ -9,6 +9,7 @@
 #include "iree/compiler/Dialect/Stream/IR/StreamTypes.h"
 #include "iree/compiler/Dialect/Stream/Transforms/Passes.h"
 #include "iree/compiler/Dialect/Util/Analysis/Explorer.h"
+#include "iree/compiler/Dialect/Util/Analysis/GlobalTable.h"
 #include "iree/compiler/Dialect/Util/IR/UtilDialect.h"
 #include "iree/compiler/Dialect/Util/IR/UtilOps.h"
 #include "llvm/ADT/STLExtras.h"
@@ -35,8 +36,7 @@ namespace {
 // from the new global op for the requested affinity.
 class ValuePerAffinityHelper {
 public:
-  explicit ValuePerAffinityHelper(mlir::ModuleOp moduleOp)
-      : builder(moduleOp), symbolTable(moduleOp) {}
+  explicit ValuePerAffinityHelper(mlir::ModuleOp moduleOp);
   ~ValuePerAffinityHelper() = default;
 
   using OpAffinityPair = std::tuple<Operation *, IREE::Stream::AffinityAttr>;
@@ -70,8 +70,41 @@ class ValuePerAffinityHelper {
   SymbolTable symbolTable;
   DenseMap<OpAffinityPair, IREE::Util::GlobalOpInterface> cachedGlobals;
   DenseMap<ValueAffinityPair, Value> cachedValuePerAffinity;
+
+  // Cache the insertion point (last initializer or the global itself) for each
+  // global for performance.
+  DenseMap<Operation *, Operation *> cachedInsertionPointForGlobal;
 };
 
+ValuePerAffinityHelper::ValuePerAffinityHelper(mlir::ModuleOp moduleOp)
+    : builder(moduleOp), symbolTable(moduleOp) {
+  // Pre-compute the insertion point for each global for performance.
+  // This avoids scanning all initializers multiple times during transformation.
+  IREE::Util::GlobalTable globalTable(moduleOp);
+  globalTable.rebuild();
+  globalTable.forEach([&](IREE::Util::Global &global) {
+    // Initialize with the global op itself as the default insertion point.
+    Operation *insertionPoint = global.op.getOperation();
+
+    // Iterate through store operations to find initializers.
+    for (auto storeOp : global.storeOps) {
+      // Get the parent initializer op if the store is within one.
+      auto initOp = storeOp->getParentOfType<IREE::Util::InitializerOp>();
+      if (!initOp) {
+        continue;
+      }
+
+      // Update the insertion point if this is the latest initializer.
+      if (insertionPoint->isBeforeInBlock(initOp)) {
+        insertionPoint = initOp;
+      }
+    }
+
+    cachedInsertionPointForGlobal[global.op.getOperation()] = insertionPoint;
+    return IREE::Util::GlobalAction::PRESERVE;
+  });
+}
+
 Value ValuePerAffinityHelper::getOrCreateValueForAffinity(
     OpOperand *opOperand, IREE::Stream::AffinityAttr affinityAttr) {
   ValueAffinityPair key = {opOperand->get(), affinityAttr};
@@ -126,9 +159,15 @@ ValuePerAffinityHelper::getOrCreateGlobalForAffinity(
     return cachedGlobals.lookup(key);
   }
 
-  // Create an initializer right after the global op.
+  // Find the insertion point: after the last initializer that references this
+  // global, or after the global itself if no initializers exist.
+  // The cache was already populated in the constructor.
+  Operation *insertionPoint =
+      cachedInsertionPointForGlobal.lookup(globalOp.getOperation());
+
+  // Create the new global and initializer after the insertion point.
   Location loc = globalOp.getLoc();
-  builder.setInsertionPointAfter(globalOp);
+  builder.setInsertionPointAfter(insertionPoint);
   std::string newGlobalName = getNewGlobalName(globalName, affinityAttr);
   auto newGlobalOp = IREE::Util::GlobalOp::create(builder, loc, newGlobalName,
                                                   /*isMutable=*/false,

compiler/src/iree/compiler/Dialect/Stream/Transforms/test/replicate_globals_per_affinity.mlir

@@ -43,6 +43,61 @@ util.func private @unknown_global_device(%arg0: tensor<10xf32>) -> (tensor<10xf3
 
 // -----
 
+// Test case with a global that has an initializer.
+// The new globals should be placed after the original global's initializer.
+
+// CHECK: util.global private @[[$DEVICE_A:.+]] : !hal.device
+// CHECK: util.global private @[[$DEVICE_B:.+]] : !hal.device
+util.global private @device_a : !hal.device
+util.global private @device_b : !hal.device
+
+// CHECK: util.global private @[[$GLOBAL:.+]] : tensor<10xf32>
+// CHECK: util.initializer {
+// CHECK:   %[[CST:.+]] = arith.constant dense<0.0{{.*}}> : tensor<10xf32>
+// CHECK:   util.global.store %[[CST]], @[[$GLOBAL]]
+// CHECK: }
+// CHECK: util.global private @[[$GLOBAL_B:.+]] : tensor<10xf32>
+// CHECK: util.initializer {
+// CHECK:   %[[LOAD:.+]] = util.global.load @[[$GLOBAL]]
+// CHECK:   %[[TRANSFER_B:.+]] = flow.tensor.transfer %[[LOAD]] {{.+}} to #hal.device.affinity<@[[$DEVICE_B]]>
+// CHECK:   util.global.store %[[TRANSFER_B]], @[[$GLOBAL_B]]
+// CHECK: }
+// CHECK: util.global private @[[$GLOBAL_A:.+]] : tensor<10xf32>
+// CHECK: util.initializer {
+// CHECK:   %[[LOAD:.+]] = util.global.load @[[$GLOBAL]]
+// CHECK:   %[[TRANSFER_A:.+]] = flow.tensor.transfer %[[LOAD]] {{.+}} to #hal.device.affinity<@[[$DEVICE_A]]>
+// CHECK:   util.global.store %[[TRANSFER_A]], @[[$GLOBAL_A]]
+// CHECK: }
+util.global private @global : tensor<10xf32>
+util.initializer {
+  %0 = arith.constant dense<0.0> : tensor<10xf32>
+  util.global.store %0, @global : tensor<10xf32>
+  util.return
+}
+
+// CHECK-LABEL: @unknown_global_device_with_initializer(
+// CHECK-SAME:    %[[ARG0:[a-zA-Z0-9]+]]
+util.func private @unknown_global_device_with_initializer(%arg0: tensor<10xf32>) -> (tensor<10xf32>, tensor<10xf32>) {
+  // CHECK: %[[OPERAND_A:.+]] = flow.tensor.transfer %[[ARG0]] {{.+}} to #hal.device.affinity<@[[$DEVICE_A]]>
+  %0 = flow.tensor.transfer %arg0 : tensor<10xf32> to #hal.device.affinity<@device_a>
+
+  // CHECK: %[[LOAD_A:.+]] = util.global.load immutable @[[$GLOBAL_A]]
+  %global = util.global.load immutable @global : tensor<10xf32>
+
+  // CHECK: flow.dispatch @dispatch(%[[OPERAND_A]], %[[LOAD_A]])
+  %1 = flow.dispatch @dispatch(%0, %global) : (tensor<10xf32>, tensor<10xf32>) -> tensor<10xf32>
+
+  // CHECK: %[[OPERAND_B:.+]] = flow.tensor.transfer %[[ARG0]] {{.+}} to #hal.device.affinity<@[[$DEVICE_B]]>
+  %2 = flow.tensor.transfer %arg0 : tensor<10xf32> to #hal.device.affinity<@device_b>
+
+  // CHECK: %[[LOAD_B:.+]] = util.global.load immutable @[[$GLOBAL_B]]
+  // CHECK: flow.dispatch @dispatch(%[[OPERAND_B]], %[[LOAD_B]])
+  %3 = flow.dispatch @dispatch(%2, %global) : (tensor<10xf32>, tensor<10xf32>) -> tensor<10xf32>
+  util.return %1, %3 : tensor<10xf32>, tensor<10xf32>
+}
+
+// -----
+
 // CHECK: util.global private @[[$DEVICE_A:.+]] : !hal.device
 // CHECK: util.global private @[[$DEVICE_B:.+]] : !hal.device
 util.global private @device_a : !hal.device