[HAL] Memoize allocator selects in program order (#25243) MemoizeDeviceSelection processed its select groups in DenseMap order, which follows attribute addresses, so the names of the generated globals (`__allocator_select_device` vs `_device_0`, etc) changed between runs and `.vmfb` files of multi-device modules were not reproducible. This does not affect correctness, only the reproducibility of `.vmfb` files. The fix uses a MapVector to process the groups in the order of their first select. The lit test now checks exact names instead of capturing them with `.*` to avoid matching only the prefix part of the name. Assisted-by: Claude Code Signed-off-by: Zmicier Prybysh <zprybysh@baylibre.com>
diff --git a/compiler/src/iree/compiler/Dialect/HAL/Transforms/MemoizeDeviceSelection.cpp b/compiler/src/iree/compiler/Dialect/HAL/Transforms/MemoizeDeviceSelection.cpp index 56f8d98..be0de8c 100644 --- a/compiler/src/iree/compiler/Dialect/HAL/Transforms/MemoizeDeviceSelection.cpp +++ b/compiler/src/iree/compiler/Dialect/HAL/Transforms/MemoizeDeviceSelection.cpp
@@ -8,6 +8,7 @@ #include "iree/compiler/Dialect/HAL/IR/HALDialect.h" #include "iree/compiler/Dialect/HAL/IR/HALOps.h" #include "iree/compiler/Dialect/Util/IR/UtilDialect.h" +#include "llvm/ADT/MapVector.h" #include "mlir/Dialect/Arith/IR/Arith.h" #include "mlir/IR/Builders.h" #include "mlir/IR/Diagnostics.h" @@ -30,7 +31,7 @@ IntegerAttr memoryType; IntegerAttr bufferUsage; - // Ideally the SelectOpOperands should be a valid key to the DenseMap, so we + // Ideally the SelectOpOperands should be a valid key to the MapVector, so we // could use it directly. However, we just take the shortcut of // converting to an ArrayAttr and using that as the key. ArrayAttr toArrayAttr(MLIRContext *context) const { @@ -158,8 +159,9 @@ // Gather all select ops in the program and bucket by unique key. // For each bucket the first op will be the first that appears in the // module for that given bucket. - DenseMap<ArrayAttr, std::pair<SelectOpOperands, - SmallVector<IREE::HAL::AllocatorSelectOp>>> + llvm::MapVector< + ArrayAttr, + std::pair<SelectOpOperands, SmallVector<IREE::HAL::AllocatorSelectOp>>> selectOps; for (auto callableOp : moduleOp.getOps<mlir::CallableOpInterface>()) { // TODO(benvanik): an interface for when we have other select ops. For now
diff --git a/compiler/src/iree/compiler/Dialect/HAL/Transforms/test/memoize_device_selection.mlir b/compiler/src/iree/compiler/Dialect/HAL/Transforms/test/memoize_device_selection.mlir index 5358e1b..609df54 100644 --- a/compiler/src/iree/compiler/Dialect/HAL/Transforms/test/memoize_device_selection.mlir +++ b/compiler/src/iree/compiler/Dialect/HAL/Transforms/test/memoize_device_selection.mlir
@@ -7,8 +7,8 @@ util.global private @device_b : !hal.device util.global private @device_c : !hal.device -// CHECK: util.global private @__allocator_select_device[[DEVICE_0:.*]] : !hal.device -// CHECK: util.global private @__allocator_select_affinity[[AFFINITY_0:.*]] : i64 +// CHECK: util.global private @__allocator_select_device : !hal.device +// CHECK: util.global private @__allocator_select_affinity : i64 // CHECK: util.initializer // CHECK: %[[I0_DEVICE_A:.+]] = util.global.load @device_a // CHECK: %[[I0_AFFINITY_A:.+]] = arith.constant -1 @@ -20,14 +20,14 @@ // CHECK: (%[[I0_DEVICE_A]], %[[I0_AFFINITY_A]] : !hal.device, i64), // CHECK: (%[[I0_DEVICE_B]], %[[I0_AFFINITY_B]] : !hal.device, i64) // CHECK: ]) type(%[[I0_MEMORY_TYPE]]) usage(%[[I0_BUFFER_USAGE]]) -// CHECK: util.global.store %[[I0_DEVICE]], @__allocator_select_device[[DEVICE_0]] -// CHECK: util.global.store %[[I0_AFFINITY]], @__allocator_select_affinity[[AFFINITY_0]] +// CHECK: util.global.store %[[I0_DEVICE]], @__allocator_select_device : !hal.device +// CHECK: util.global.store %[[I0_AFFINITY]], @__allocator_select_affinity : i64 // CHECK: @fn1 util.func public @fn1() -> (!hal.device, i64) { // CHECK-NOT: hal.allocator.select - // CHECK: %[[FN1_DEVICE:.+]] = util.global.load @__allocator_select_device[[DEVICE_0]] - // CHECK: %[[FN1_AFFINITY:.+]] = util.global.load @__allocator_select_affinity[[AFFINITY_0]] + // CHECK: %[[FN1_DEVICE:.+]] = util.global.load @__allocator_select_device : !hal.device + // CHECK: %[[FN1_AFFINITY:.+]] = util.global.load @__allocator_select_affinity : i64 %memory_type = hal.memory_type<"HostLocal"> : i32 %buffer_usage = hal.buffer_usage<"Transfer"> : i32 %device_a = util.global.load @device_a : !hal.device @@ -45,8 +45,8 @@ // CHECK: @fn2 util.func public @fn2() -> (!hal.device, i64) { // CHECK-NOT: hal.allocator.select - // CHECK: %[[FN2_DEVICE:.+]] = util.global.load @__allocator_select_device[[DEVICE_0]] - // CHECK: %[[FN2_AFFINITY:.+]] = util.global.load @__allocator_select_affinity[[AFFINITY_0]] + // CHECK: %[[FN2_DEVICE:.+]] = util.global.load @__allocator_select_device : !hal.device + // CHECK: %[[FN2_AFFINITY:.+]] = util.global.load @__allocator_select_affinity : i64 %memory_type = hal.memory_type<"HostLocal"> : i32 %buffer_usage = hal.buffer_usage<"Transfer"> : i32 %device_a = util.global.load @device_a : !hal.device @@ -61,15 +61,15 @@ util.return %device, %queue_affinity : !hal.device, i64 } -// CHECK: util.global private @__allocator_select_device[[DEVICE_1:.*]] : !hal.device -// CHECK: util.global private @__allocator_select_affinity[[AFFINITY_1:.*]] : i64 +// CHECK: util.global private @__allocator_select_device_0 : !hal.device +// CHECK: util.global private @__allocator_select_affinity_1 : i64 // CHECK: util.initializer // CHECK: @fn3 util.func public @fn3() -> (!hal.device, i64) { // CHECK-NOT: hal.allocator.select - // CHECK: %[[FN3_DEVICE:.+]] = util.global.load @__allocator_select_device[[DEVICE_1]] - // CHECK: %[[FN3_AFFINITY:.+]] = util.global.load @__allocator_select_affinity[[AFFINITY_1]] + // CHECK: %[[FN3_DEVICE:.+]] = util.global.load @__allocator_select_device_0 : !hal.device + // CHECK: %[[FN3_AFFINITY:.+]] = util.global.load @__allocator_select_affinity_1 : i64 %memory_type = hal.memory_type<"HostLocal"> : i32 %buffer_usage = hal.buffer_usage<"DispatchStorage"> : i32 %device_a = util.global.load @device_a : !hal.device