[Metal] Carry export name separately from MSL entry point for name lookup (#24642)
spirv-cross rewrites SPIR-V entry point names to be legal MSL/C
identifiers when serializing Metal executables: '.' characters (codegen
dispatch names like _dot_general_matmul_i32.i32.i32_...) become '_', and
reserved '__' prefixes (stream builtins like __builtin_splat_i64) are
dropped. The MetalSPIRV target stored only the rewritten MSL function
name in PipelineDef.entry_point, and the runtime loaded pipeline->name
from that same field -- so hal.executable.lookup.function, which
references the original un-rewritten hal.executable.export symbol name,
could never match for any export whose name spirv-cross rewrites. This
broke every metal-spirv e2e test using such exports
(dot/dot_general/convolution matmul kernels and
compare/dynamic_slice/gather stream builtins).
Add a dedicated always-present name field to PipelineDef carrying the
canonical hal.executable.export symbol name (populated from
exportOp.getName() in the compiler), appended at the end of the table
for backward-compatible vtable layout. The runtime now sources
pipeline->name (used by lookup_function_by_name) from this field,
keeping entry_point as the MTLFunction name for -[MTLLibrary
newFunctionWithName:]. The loader falls back to entry_point when name is
absent for compatibility.
Verified: dot_general (8/8), compare (15/15), and the full Metal CTS
suite
(command_buffer_tests/queue_tests/dispatch_tests/queue_dispatch_tests,
4/4) pass. Signed-off-by: Alex Vasile
<48962821+Alex-Vasile@users.noreply.github.com>
Signed-off-by: Alex Vasile <48962821+Alex-Vasile@users.noreply.github.com>
diff --git a/compiler/plugins/target/MetalSPIRV/MetalSPIRVTarget.cpp b/compiler/plugins/target/MetalSPIRV/MetalSPIRVTarget.cpp
index 5618a5b..feb61b2 100644
--- a/compiler/plugins/target/MetalSPIRV/MetalSPIRVTarget.cpp
+++ b/compiler/plugins/target/MetalSPIRV/MetalSPIRVTarget.cpp
@@ -157,9 +157,11 @@
".spv", spvBinary);
}
- // The runtime use ordinals instead of names but Metal requires function
- // names for constructing pipeline states. Get an ordered list of the entry
- // point names.
+ // Collect the SPIR-V entry point function names. These are cross-compiled
+ // to MSL and used to construct Metal pipeline states; they are the MSL
+ // function names (which spirv-cross may rewrite to be MSL-legal) and are
+ // distinct from the hal.executable.export symbol names the runtime uses to
+ // look up exports by name.
SmallVector<StringRef, 8> spirvEntryPointNames;
spvModuleOp.walk([&](spirv::EntryPointOp exportOp) {
spirvEntryPointNames.push_back(exportOp.getFn());
@@ -262,6 +264,7 @@
llvm::zip_equal(llvm::seq(mslShaders.size()), mslShaders,
mslEntryPointNames, exportOps)) {
auto entryPointRef = builder.createString(entryPoint);
+ auto exportNameRef = builder.createString(exportOp.getName());
iree_hal_metal_ThreadgroupSize_t threadgroupSize = {
shader.threadgroupSize.x,
@@ -286,6 +289,7 @@
iree_hal_metal_PipelineDef_start(builder);
iree_hal_metal_PipelineDef_library_ordinal_add(builder, i);
iree_hal_metal_PipelineDef_entry_point_add(builder, entryPointRef);
+ iree_hal_metal_PipelineDef_name_add(builder, exportNameRef);
iree_hal_metal_PipelineDef_threadgroup_size_add(builder,
&threadgroupSize);
// TODO: embed additional metadata on threadgroup info if available.
diff --git a/runtime/src/iree/hal/drivers/metal/executable.m b/runtime/src/iree/hal/drivers/metal/executable.m
index 87c6447..f4a990a 100644
--- a/runtime/src/iree/hal/drivers/metal/executable.m
+++ b/runtime/src/iree/hal/drivers/metal/executable.m
@@ -424,18 +424,23 @@
iree_hal_metal_ExecutableDef_pipelines_get(executable_def);
iree_host_size_t pipeline_count = flatbuffers_string_vec_len(pipelines_vec);
- // Calculate the total number of characters across all entry point names so
- // that we can store copies of the names as the flatbuffer storing the strings
- // may be released while the executable is still live.
+ // Calculate the total number of characters across all export names so that we
+ // can store copies of the names as the flatbuffer storing the strings may be
+ // released while the executable is still live.
iree_host_size_t total_pipeline_name_length = 0;
iree_host_size_t total_debug_info_length = 0;
for (iree_host_size_t i = 0; i < pipeline_count; ++i) {
iree_hal_metal_PipelineDef_table_t pipeline_def =
iree_hal_metal_PipelineDef_vec_at(pipelines_vec, i);
- if (IREE_UNLIKELY(!iree_host_size_checked_add(
- total_pipeline_name_length,
- flatbuffers_string_len(iree_hal_metal_PipelineDef_entry_point_get(pipeline_def)),
- &total_pipeline_name_length))) {
+ flatbuffers_string_t export_name = iree_hal_metal_PipelineDef_name_get(pipeline_def);
+ if (flatbuffers_string_len(export_name) == 0) {
+ // Fall back to the entry point name for executables that don't carry an
+ // explicit export name (equivalent to the pre-name-field behavior).
+ export_name = iree_hal_metal_PipelineDef_entry_point_get(pipeline_def);
+ }
+ if (IREE_UNLIKELY(!iree_host_size_checked_add(total_pipeline_name_length,
+ flatbuffers_string_len(export_name),
+ &total_pipeline_name_length))) {
return iree_make_status(IREE_STATUS_OUT_OF_RANGE, "pipeline name storage size overflow");
}
IREE_TRACE({
@@ -487,11 +492,14 @@
status = iree_hal_metal_create_pipeline(device, library, pipeline_def, pipeline);
if (!iree_status_is_ok(status)) break;
- flatbuffers_string_t entry_point = iree_hal_metal_PipelineDef_entry_point_get(pipeline_def);
- const iree_host_size_t entry_point_length = flatbuffers_string_len(entry_point);
- pipeline->name = iree_make_string_view(pipeline_name_ptr, entry_point_length);
- memcpy(pipeline_name_ptr, entry_point, entry_point_length);
- pipeline_name_ptr += entry_point_length;
+ flatbuffers_string_t export_name = iree_hal_metal_PipelineDef_name_get(pipeline_def);
+ if (flatbuffers_string_len(export_name) == 0) {
+ export_name = iree_hal_metal_PipelineDef_entry_point_get(pipeline_def);
+ }
+ const iree_host_size_t export_name_length = flatbuffers_string_len(export_name);
+ pipeline->name = iree_make_string_view(pipeline_name_ptr, export_name_length);
+ memcpy(pipeline_name_ptr, export_name, export_name_length);
+ pipeline_name_ptr += export_name_length;
IREE_TRACE({
iree_hal_debug_export_info_t* export_info = (iree_hal_debug_export_info_t*)export_info_ptr;
diff --git a/runtime/src/iree/schemas/metal_executable_def.fbs b/runtime/src/iree/schemas/metal_executable_def.fbs
index fbee2ad..550b7e8 100644
--- a/runtime/src/iree/schemas/metal_executable_def.fbs
+++ b/runtime/src/iree/schemas/metal_executable_def.fbs
@@ -31,7 +31,11 @@
// libraries list.
library_ordinal:uint32;
- // String name of the entry point MTLFunction in the library.
+ // String name of the entry point MTLFunction in the library. This is the MSL
+ // function name produced by spirv-cross, which rewrites symbols that are
+ // illegal or reserved in MSL/C (for example '.' characters become '_' and a
+ // leading reserved '__' prefix is dropped). Used to construct the
+ // MTLComputePipelineState via -[MTLLibrary newFunctionWithName:].
entry_point:string;
// The maximum number of threads in a threadgroup that can be dispatched to
@@ -55,6 +59,12 @@
// Optional debug information related to the export.
debug_info:iree.hal.debug.ExportDef;
+
+ // Canonical hal.executable.export symbol name used by the runtime to look up
+ // the export by name (hal.executable.lookup.function). This carries the
+ // original, un-rewritten name and may differ from entry_point because MSL
+ // cannot represent all legal symbol names; dispatch lookups must match it.
+ name:string;
}
// MSL (Metal Shading Language) source code.
