[HAL] Renumber export ordinals after pruning (#24162)

Backend serializers index the flatbuffer `exports` vector by ordinal and
assume a dense [0, N) range. When PruneExecutablesPass removed unused
exports it left ordinals sparse (e.g. [0, 2, 3, 4] with 1 pruned),
causing the serializer to write past the end of its vector and produce
invalid flatbuffers. This showed up as flatcc verification failures at
runtime for models with unused `hal.executable.source` exports.
Renumbering to a dense [0, N) range here restores the invariant for all
backends without touching any serializer or runtime code.

Signed-off-by: Jorn Tuyls <jorn.tuyls@gmail.com>

Author: jtuyls

compiler/src/iree/compiler/Dialect/HAL/Transforms/PruneExecutables.cpp

@@ -174,6 +174,23 @@ struct PruneExecutablesPass
     eraseOps(exportRefAttrs.getArrayRef(), referenceMap);
     eraseOps(variantRefAttrs, referenceMap);
     eraseOps(executableRefAttrs, referenceMap);
+
+    // Renumber export ordinals to be contiguous from 0. Backend serializers
+    // index their flatbuffer `exports` vector by ordinal and assume a dense
+    // [0, N) range; leaving gaps from pruned exports produces invalid
+    // flatbuffers. Export symbol references are resolved by name via
+    // ResolveExportOrdinalsPass, so renumbering is a safe rewrite.
+    Builder builder(moduleOp.getContext());
+    for (auto executableOp : moduleOp.getOps<IREE::HAL::ExecutableOp>()) {
+      for (auto variantOp :
+           executableOp.getOps<IREE::HAL::ExecutableVariantOp>()) {
+        int64_t nextOrdinal = 0;
+        for (auto exportOp :
+             variantOp.getOps<IREE::HAL::ExecutableExportOp>()) {
+          exportOp.setOrdinalAttr(builder.getIndexAttr(nextOrdinal++));
+        }
+      }
+    }
   }
 };
 

compiler/src/iree/compiler/Dialect/HAL/Transforms/test/prune_executables.mlir

@@ -99,3 +99,35 @@ hal.executable private @exe {
 util.func private @user() attributes {
   some.ref = @exe::@variant::@used_export
 }
+
+// -----
+
+// Tests that after pruning unused exports the remaining exports are
+// renumbered to have contiguous ordinals starting from 0. Back-end
+// serializers index their flatbuffer exports vector by ordinal; if ordinals
+// are left sparse the downstream vector has gaps that compound into invalid
+// flatbuffers (ExportDef_ref_t at index N overwrites at a smaller index when
+// the vector is sized by pre-prune export count).
+
+#pipeline_layout_renumber = #hal.pipeline.layout<bindings = [
+  #hal.pipeline.binding<storage_buffer>
+]>
+hal.executable private @exe_renumber {
+  hal.executable.variant public @variant target(<"backend", "format">) {
+    // CHECK-NOT: @unused_a
+    hal.executable.export public @unused_a ordinal(1) layout(#pipeline_layout_renumber)
+    // CHECK: hal.executable.export public @used_0 ordinal(0)
+    hal.executable.export public @used_0 ordinal(0) layout(#pipeline_layout_renumber)
+    // CHECK: hal.executable.export public @used_1 ordinal(1)
+    hal.executable.export public @used_1 ordinal(2) layout(#pipeline_layout_renumber)
+    // CHECK-NOT: @unused_b
+    hal.executable.export public @unused_b ordinal(3) layout(#pipeline_layout_renumber)
+    // CHECK: hal.executable.export public @used_2 ordinal(2)
+    hal.executable.export public @used_2 ordinal(4) layout(#pipeline_layout_renumber)
+  }
+}
+util.func private @multi_user() attributes {
+  ref_0 = @exe_renumber::@variant::@used_0,
+  ref_1 = @exe_renumber::@variant::@used_1,
+  ref_2 = @exe_renumber::@variant::@used_2
+}