Starting support for HAL dispatch specialization. (#12483)
This allows `stream.cmd.dispatch` ops to specify multiple entry points
in executables that can be dispatched. During interface materialization
the entry points are expanded for all materialized variants. HAL
backends are also able to add their own entry points either in existing
variants or new ones during translation (no helpers yet, but it's
possible). When lowering the `stream.cmd.dispatch` into
`hal.command_buffer.dispatch` each export now chooses the condition for
which it should be selected.
This initial work just prepares the op for this conditional dispatch and
also solves an issue with compiler reentrancy where between interface
materialization and HAL conversion we had IR that would not verify
(stream ops still referencing un-materialized exports). Only the
existing target variant matcher condition is supported but we now have
the place where we'd call out to HAL backend-produced logic for
selection (`getExportConditionAttr`).
The `inline-dynamic` HAL model can support this too but for now that's
deferred until we add the backend support - the `hal.device.switch` op
needs a reworking to be used in that context.
Fixes #12476.
diff --git a/compiler/src/iree/compiler/Dialect/HAL/Analysis/BindingLayout.cpp b/compiler/src/iree/compiler/Dialect/HAL/Analysis/BindingLayout.cpp
index 632b0e2..88393be 100644
--- a/compiler/src/iree/compiler/Dialect/HAL/Analysis/BindingLayout.cpp
+++ b/compiler/src/iree/compiler/Dialect/HAL/Analysis/BindingLayout.cpp
@@ -45,9 +45,11 @@
SymbolTable symbolTable(rootOp);
BindingLayoutAnalysis::ExportDispatchMap dispatchMap;
rootOp->walk([&](IREE::Stream::CmdDispatchOp dispatchOp) {
- auto exportOp = symbolTable.lookupNearestSymbolFrom(
- dispatchOp, dispatchOp.getEntryPointAttr());
- dispatchMap[exportOp].push_back(dispatchOp);
+ dispatchOp.forEachEntryPointAttr([&](SymbolRefAttr entryPointAttr) {
+ auto exportOp =
+ symbolTable.lookupNearestSymbolFrom(dispatchOp, entryPointAttr);
+ dispatchMap[exportOp].push_back(dispatchOp);
+ });
});
return dispatchMap;
}
diff --git a/compiler/src/iree/compiler/Dialect/HAL/Conversion/StreamToHAL/Patterns.cpp b/compiler/src/iree/compiler/Dialect/HAL/Conversion/StreamToHAL/Patterns.cpp
index 76f827a..b785220 100644
--- a/compiler/src/iree/compiler/Dialect/HAL/Conversion/StreamToHAL/Patterns.cpp
+++ b/compiler/src/iree/compiler/Dialect/HAL/Conversion/StreamToHAL/Patterns.cpp
@@ -806,6 +806,16 @@
}
};
+// Returns a hal.device.switch match expression that selects the given export.
+static Attribute getExportConditionAttr(
+ IREE::HAL::ExecutableExportOp exportOp) {
+ // TODO(benvanik): customizable selection logic. Today this just checks
+ // whether the variant target is supported but we can also allow
+ // specialization of entry points based on dispatch site parameters.
+ auto variantOp = exportOp->getParentOfType<IREE::HAL::ExecutableVariantOp>();
+ return variantOp.getTarget().getMatchExpression();
+}
+
struct CmdDispatchOpPattern
: public StreamConversionPattern<IREE::Stream::CmdDispatchOp> {
using StreamConversionPattern::StreamConversionPattern;
@@ -821,34 +831,21 @@
auto device = rewriter.create<IREE::HAL::CommandBufferDeviceOp>(
loc, rewriter.getType<IREE::HAL::DeviceType>(), commandBuffer);
- // Get the handle to the executable that is compatible with our device.
- auto executableOp =
- cast<IREE::HAL::ExecutableOp>(SymbolTable::lookupNearestSymbolFrom(
- dispatchOp, dispatchOp.getEntryPoint().getRootReference()));
- assert(executableOp && "dispatch target executable op not found");
-
// Ask each target backend to record their dispatch logic.
IREE::HAL::DeviceSwitchRewriter switchRewriter(loc,
/*resultTypes=*/TypeRange{},
device, rewriter);
- for (auto variantOp :
- executableOp.getOps<IREE::HAL::ExecutableVariantOp>()) {
- auto exportOps = variantOp.getOps<IREE::HAL::ExecutableExportOp>();
- auto exportIt =
- llvm::find_if(exportOps, [&](IREE::HAL::ExecutableExportOp op) {
- return op.getNameAttr() ==
- dispatchOp.getEntryPoint().getLeafReference();
- });
- if (exportIt == exportOps.end()) {
- return variantOp.emitError()
- << "hal.executable.variant is missing the flow entry point for "
- << dispatchOp.getEntryPoint();
- }
- auto exportOp = *exportIt;
+ dispatchOp.forEachEntryPointAttr([&](SymbolRefAttr entryPointAttr) {
+ // NOTE: slow lookup!
+ auto exportOp =
+ SymbolTable::lookupNearestSymbolFrom<IREE::HAL::ExecutableExportOp>(
+ dispatchOp, entryPointAttr);
+ assert(exportOp && "dispatch target export not found");
- auto *region = switchRewriter.addConditionRegion(
- variantOp.getTarget().getMatchExpression());
- auto &entryBlock = region->front();
+ // Setup the case condition for the entry point.
+ auto *caseRegion =
+ switchRewriter.addConditionRegion(getExportConditionAttr(exportOp));
+ auto &entryBlock = caseRegion->front();
auto caseBuilder = OpBuilder::atBlockBegin(&entryBlock);
// Record push constants and buffer bindings.
@@ -856,18 +853,14 @@
exportOp.getLayout(), caseBuilder);
// Dispatch with a target-specific workgroup count.
- auto exportSymRef =
- SymbolRefAttr::get(caseBuilder.getContext(), executableOp.getName(),
- {SymbolRefAttr::get(exportOp->getParentOp()),
- SymbolRefAttr::get(exportOp)});
auto caseWorkgroupCount = exportOp.calculateWorkgroupCount(
loc, device, adaptor.getWorkload(), caseBuilder);
caseBuilder.create<IREE::HAL::CommandBufferDispatchSymbolOp>(
- loc, commandBuffer, exportSymRef, caseWorkgroupCount[0],
+ loc, commandBuffer, entryPointAttr, caseWorkgroupCount[0],
caseWorkgroupCount[1], caseWorkgroupCount[2]);
caseBuilder.create<IREE::HAL::ReturnOp>(loc);
- }
+ });
switchRewriter.build();
rewriter.eraseOp(dispatchOp);
diff --git a/compiler/src/iree/compiler/Dialect/HAL/Conversion/StreamToHAL/test/cmd_ops.mlir b/compiler/src/iree/compiler/Dialect/HAL/Conversion/StreamToHAL/test/cmd_ops.mlir
index ec26860..5d09a21 100644
--- a/compiler/src/iree/compiler/Dialect/HAL/Conversion/StreamToHAL/test/cmd_ops.mlir
+++ b/compiler/src/iree/compiler/Dialect/HAL/Conversion/StreamToHAL/test/cmd_ops.mlir
@@ -232,7 +232,7 @@
// CHECK: hal.command_buffer.dispatch.symbol<%[[CMD]]
// CHECK-SAME: target(@ex::@embedded_elf_x86_64::@dispatch)
// CHECK-SAME: workgroups([%[[X]], %[[YZ]], %[[YZ]]])
- stream.cmd.dispatch @ex::@dispatch[%c1, %c2, %c3](%c4_i32, %c5_i32 : i32, i32) {
+ stream.cmd.dispatch @ex::@embedded_elf_x86_64::@dispatch[%c1, %c2, %c3](%c4_i32, %c5_i32 : i32, i32) {
ro %arg4[%c0 for %c128] : !stream.resource<transient>{%arg1},
wo %arg5[%c0 for %c128] : !stream.resource<external>{%arg3}
} attributes {
diff --git a/compiler/src/iree/compiler/Dialect/HAL/Transforms/DumpExecutableBenchmarks.cpp b/compiler/src/iree/compiler/Dialect/HAL/Transforms/DumpExecutableBenchmarks.cpp
index 440a6d4..2b26df2 100644
--- a/compiler/src/iree/compiler/Dialect/HAL/Transforms/DumpExecutableBenchmarks.cpp
+++ b/compiler/src/iree/compiler/Dialect/HAL/Transforms/DumpExecutableBenchmarks.cpp
@@ -110,22 +110,24 @@
}
// Work around needing a mutable key for the set; C++ was a mistake.
- auto &dispatchParamsSet = map[dispatchOp.getEntryPoint()];
- DispatchParams *dispatchParams = nullptr;
- for (auto &it : dispatchParamsSet) {
- if (it.workload == workload) {
- dispatchParams = ⁢
- break;
+ dispatchOp.forEachEntryPointAttr([&](SymbolRefAttr entryPointAttr) {
+ auto &dispatchParamsSet = map[entryPointAttr];
+ DispatchParams *dispatchParams = nullptr;
+ for (auto &it : dispatchParamsSet) {
+ if (it.workload == workload) {
+ dispatchParams = ⁢
+ break;
+ }
}
- }
- if (!dispatchParams) {
- dispatchParamsSet.push_back({});
- dispatchParams = &dispatchParamsSet.back();
- }
- dispatchParams->locs.push_back(dispatchOp.getLoc());
- dispatchParams->workload = workload;
- dispatchParams->bindings = std::move(bindings);
- dispatchParams->uniformOperands = std::move(uniformOperands);
+ if (!dispatchParams) {
+ dispatchParamsSet.push_back({});
+ dispatchParams = &dispatchParamsSet.back();
+ }
+ dispatchParams->locs.push_back(dispatchOp.getLoc());
+ dispatchParams->workload = workload;
+ dispatchParams->bindings = std::move(bindings);
+ dispatchParams->uniformOperands = std::move(uniformOperands);
+ });
});
}
@@ -390,7 +392,10 @@
for (auto exportOp : variantOp.getOps<IREE::HAL::ExecutableExportOp>()) {
auto symbolRefAttr =
SymbolRefAttr::get(executableOp.getNameAttr(),
- {FlatSymbolRefAttr::get(exportOp.getNameAttr())});
+ {
+ FlatSymbolRefAttr::get(variantOp.getNameAttr()),
+ FlatSymbolRefAttr::get(exportOp.getNameAttr()),
+ });
auto dispatchParamsSet = dispatchParamsMap.find(symbolRefAttr);
if (dispatchParamsSet != dispatchParamsMap.end()) {
for (auto &dispatchParams : dispatchParamsSet->second) {
diff --git a/compiler/src/iree/compiler/Dialect/HAL/Transforms/MaterializeDispatchInstrumentation.cpp b/compiler/src/iree/compiler/Dialect/HAL/Transforms/MaterializeDispatchInstrumentation.cpp
index 5426cc9..764f9b2 100644
--- a/compiler/src/iree/compiler/Dialect/HAL/Transforms/MaterializeDispatchInstrumentation.cpp
+++ b/compiler/src/iree/compiler/Dialect/HAL/Transforms/MaterializeDispatchInstrumentation.cpp
@@ -260,8 +260,19 @@
// Walk dispatches and pass them the ringbuffer and their unique ID.
executeOp.walk([&](IREE::Stream::CmdDispatchOp dispatchOp) {
- auto it = instrumentedExports.find(dispatchOp.getEntryPoint());
- if (it == instrumentedExports.end()) return; // not instrumented
+ // NOTE: we just choose the first instrumented export for attribution
+ // as that's good enough for all current use cases. If we start
+ // specializing really early we may want to fix that.
+ Optional<uint32_t> functionId;
+ for (auto entryPointAttr : dispatchOp.getEntryPointRefs()) {
+ auto it = instrumentedExports.find(entryPointAttr);
+ if (it != instrumentedExports.end()) {
+ // Found the first instrumented export.
+ functionId = it->second;
+ break;
+ }
+ }
+ if (!functionId) return; // not instrumented
// Append dispatch site ID to correlate this op with where it lives in
// the program and what is being dispatched. Note that multiple
@@ -277,7 +288,7 @@
iree_instruments_DispatchSiteDef_start(metadataBuilder);
// TODO(benvanik): source loc to identify the site.
iree_instruments_DispatchSiteDef_function_add(metadataBuilder,
- it->second);
+ *functionId);
dispatchSiteRefs.push_back(
iree_instruments_DispatchSiteDef_end(metadataBuilder));
diff --git a/compiler/src/iree/compiler/Dialect/HAL/Transforms/MaterializeInterfaces.cpp b/compiler/src/iree/compiler/Dialect/HAL/Transforms/MaterializeInterfaces.cpp
index 0e752d8..b07edc3 100644
--- a/compiler/src/iree/compiler/Dialect/HAL/Transforms/MaterializeInterfaces.cpp
+++ b/compiler/src/iree/compiler/Dialect/HAL/Transforms/MaterializeInterfaces.cpp
@@ -33,6 +33,9 @@
namespace HAL {
namespace {
+// Map of original SymbolRefAttr to a list of SymbolRefAttrs in variants.
+using EntryPointExpansions = DenseMap<Attribute, SmallVector<Attribute>>;
+
//===----------------------------------------------------------------------===//
// Linkage utilities
//===----------------------------------------------------------------------===//
@@ -51,9 +54,20 @@
// hal.executable.source materialization
//===----------------------------------------------------------------------===//
+SymbolRefAttr makeExportSymbolRefAttr(IREE::HAL::ExecutableOp executableOp,
+ IREE::HAL::ExecutableVariantOp variantOp,
+ IREE::HAL::ExecutableExportOp exportOp) {
+ return SymbolRefAttr::get(executableOp.getNameAttr(),
+ {
+ FlatSymbolRefAttr::get(variantOp.getNameAttr()),
+ FlatSymbolRefAttr::get(exportOp.getNameAttr()),
+ });
+}
+
static LogicalResult materializeExecutableFromSourceOp(
IREE::HAL::ExecutableSourceOp sourceOp,
- ArrayRef<IREE::HAL::ExecutableTargetAttr> targetAttrs) {
+ ArrayRef<IREE::HAL::ExecutableTargetAttr> targetAttrs,
+ EntryPointExpansions &entryPointExpansions) {
OpBuilder moduleBuilder(sourceOp);
// Create the op that will contain the translated executable.
@@ -77,6 +91,14 @@
OpBuilder variantBuilder(&targetVariantOp.getBlock().back());
for (auto sourceEntryPointOp : sourceEntryPointOps) {
variantBuilder.clone(*sourceEntryPointOp);
+
+ // Map the original export names to the new variant exports.
+ entryPointExpansions[SymbolRefAttr::get(
+ executableOp.getNameAttr(),
+ {FlatSymbolRefAttr::get(
+ sourceEntryPointOp.getNameAttr())})]
+ .push_back(makeExportSymbolRefAttr(executableOp, targetVariantOp,
+ sourceEntryPointOp));
}
// Clone any target-specific object files specified.
@@ -99,7 +121,7 @@
}
static LogicalResult materializeExecutablesFromSourceOps(
- mlir::ModuleOp moduleOp) {
+ mlir::ModuleOp moduleOp, EntryPointExpansions &entryPointExpansions) {
auto sourceOps =
llvm::to_vector<32>(moduleOp.getOps<IREE::HAL::ExecutableSourceOp>());
for (auto sourceOp : sourceOps) {
@@ -112,7 +134,8 @@
<< "no executable targets specified for translation";
}
- if (failed(materializeExecutableFromSourceOp(sourceOp, targetAttrs))) {
+ if (failed(materializeExecutableFromSourceOp(sourceOp, targetAttrs,
+ entryPointExpansions))) {
return failure();
}
}
@@ -234,6 +257,26 @@
return clonedFuncOp;
}
+// Updates the target entry point symbols of |dispatchOp| to the expanded set of
+// variant exports in |entryPointExpansions|.
+static void updateDispatchTargets(
+ IREE::Stream::CmdDispatchOp dispatchOp,
+ const EntryPointExpansions &entryPointExpansions) {
+ SmallVector<Attribute> newAttrs;
+ for (auto oldAttr : dispatchOp.getEntryPointRefs()) {
+ auto it = entryPointExpansions.find(oldAttr);
+ if (it == entryPointExpansions.end()) {
+ newAttrs.push_back(oldAttr); // preserve existing
+ continue;
+ }
+ for (auto newAttr : it->second) {
+ newAttrs.push_back(newAttr);
+ }
+ }
+ dispatchOp.setEntryPointsAttr(
+ ArrayAttr::get(dispatchOp.getContext(), newAttrs));
+}
+
// Annotates |dispatchOp| with resource binding to interface binding mappings.
// TODO(benvanik): have a HAL op with structured information instead.
static void annotateDispatchSite(IREE::Stream::CmdDispatchOp dispatchOp,
@@ -253,7 +296,8 @@
static LogicalResult declareEntryPointOps(
IREE::Stream::ExecutableOp sourceExecutableOp,
IREE::HAL::ExecutableOp targetExecutableOp,
- const BindingLayoutAnalysis &layoutAnalysis) {
+ const BindingLayoutAnalysis &layoutAnalysis,
+ EntryPointExpansions &entryPointExpansions) {
auto sourceModuleOp = sourceExecutableOp.getInnerModule();
auto variantOps =
targetExecutableOp.getBlock().getOps<IREE::HAL::ExecutableVariantOp>();
@@ -294,6 +338,13 @@
/*subgroup_size=*/IntegerAttr{},
/*workgroup_local_memory=*/IntegerAttr{});
+ // Map the original export name to the new variant export.
+ entryPointExpansions[SymbolRefAttr::get(sourceExecutableOp.getNameAttr(),
+ {FlatSymbolRefAttr::get(
+ exportOp.getNameAttr())})]
+ .push_back(makeExportSymbolRefAttr(targetExecutableOp, variantOp,
+ newExportOp));
+
// Clone the workgroup count calculation function.
if (!exportOp.getWorkgroupCount().empty()) {
mlir::IRMapping mapper;
@@ -431,9 +482,12 @@
void runOnOperation() override {
SymbolTable symbolTable(getOperation());
+ EntryPointExpansions entryPointExpansions;
+
// Handle any hand-authored executables; these only need variant expansion
// and no layout analysis as the user specified the layout themselves.
- if (failed(materializeExecutablesFromSourceOps(getOperation()))) {
+ if (failed(materializeExecutablesFromSourceOps(getOperation(),
+ entryPointExpansions))) {
return signalPassFailure();
}
@@ -482,8 +536,8 @@
}
// Define interfaces for each exported function based on analysis.
- if (failed(
- declareEntryPointOps(sourceOp, executableOp, layoutAnalysis))) {
+ if (failed(declareEntryPointOps(sourceOp, executableOp, layoutAnalysis,
+ entryPointExpansions))) {
return signalPassFailure();
}
@@ -501,20 +555,27 @@
// won't have materialized them from the stream ops above. We do expect to
// be able to find the dispatch targets such that we can pull out the
// pipeline layout, though, and any that fall through are errors.
- auto annotateDispatchOp = [&](IREE::Stream::CmdDispatchOp dispatchOp) {
+ auto updateDispatchSites = [&](IREE::Stream::CmdDispatchOp dispatchOp) {
+ // Update the export targets to point at the new variants.
+ updateDispatchTargets(dispatchOp, entryPointExpansions);
+
+ // Annotate the dispatch site with binding information if required.
+ // TODO(benvanik): remove this path; shouldn't be needed in real usage.
+ // Because this is a hack we just look for the first target entry point.
if (dispatchOp->hasAttr("hal.interface.bindings")) {
// Already have bindings defined.
return WalkResult::advance();
}
PipelineResourceMap resourceMap;
- auto exportOp =
+ auto anyEntryPointAttr = *dispatchOp.getEntryPointRefs().begin();
+ auto anyExportOp =
symbolTable.lookupNearestSymbolFrom<IREE::HAL::ExecutableExportOp>(
- dispatchOp, dispatchOp.getEntryPointAttr());
- if (exportOp) {
+ dispatchOp, anyEntryPointAttr);
+ if (anyExportOp) {
// Export found - we can use the pipeline layout defined there to infer
// the bindings. This allows for bindings to be sparse or have
// additional information declared.
- for (auto setLayout : exportOp.getLayoutAttr().getSetLayouts()) {
+ for (auto setLayout : anyExportOp.getLayoutAttr().getSetLayouts()) {
for (auto binding : setLayout.getBindings()) {
resourceMap.emplace_back(setLayout.getOrdinal(),
binding.getOrdinal());
@@ -535,7 +596,7 @@
annotateDispatchSite(dispatchOp, resourceMap);
return WalkResult::advance();
};
- if (getOperation()->walk(annotateDispatchOp).wasInterrupted()) {
+ if (getOperation()->walk(updateDispatchSites).wasInterrupted()) {
return signalPassFailure();
}
}
diff --git a/compiler/src/iree/compiler/Dialect/HAL/Transforms/test/convert_to_hal.mlir b/compiler/src/iree/compiler/Dialect/HAL/Transforms/test/convert_to_hal.mlir
index 4c5513f..2aca6c5 100644
--- a/compiler/src/iree/compiler/Dialect/HAL/Transforms/test/convert_to_hal.mlir
+++ b/compiler/src/iree/compiler/Dialect/HAL/Transforms/test/convert_to_hal.mlir
@@ -2,6 +2,7 @@
// Tests an end-to-end simple single-dispatch `dispatch(arg0, arg1) -> result`.
+#executable_target_embedded_elf_aarch64_ = #hal.executable.target<"llvm-cpu", "embedded-elf-aarch64">
#executable_target_embedded_elf_x86_64_ = #hal.executable.target<"llvm-cpu", "embedded-elf-x86_64">
#device_target_cpu = #hal.device.target<"llvm-cpu", {
executable_targets = [#executable_target_embedded_elf_x86_64_]
@@ -19,6 +20,19 @@
// CHECK: hal.executable private @ex
hal.executable private @ex {
+ hal.executable.variant public @embedded_elf_aarch64, target = #executable_target_embedded_elf_aarch64_ {
+ hal.executable.export public @dispatch ordinal(0) layout(#pipeline_layout) attributes {
+ translation_info = #iree_codegen.translation_info<CPUDefault>
+ } {
+ ^bb0(%device: !hal.device, %arg0: index, %arg1: index, %arg2: index): // no predecessors
+ %c1 = arith.constant 1 : index
+ %0 = affine.apply affine_map<()[s0] -> (s0 ceildiv 4)>()[%arg0]
+ hal.return %0, %c1, %c1 : index, index, index
+ }
+ builtin.module {
+ // Opaque at this point (in some target-specific dialects).
+ }
+ }
hal.executable.variant public @embedded_elf_x86_64, target = #executable_target_embedded_elf_x86_64_ {
hal.executable.export public @dispatch ordinal(0) layout(#pipeline_layout) attributes {
translation_info = #iree_codegen.translation_info<CPUDefault>
@@ -97,7 +111,7 @@
// CHECK-SAME: workgroups([%c1, %c1, %c1])
// CHECK: hal.return
// CHECK: }
- stream.cmd.dispatch @ex::@dispatch[%c4, %c1, %c1] {
+ stream.cmd.dispatch {@ex::@embedded_elf_aarch64::@dispatch, @ex::@embedded_elf_x86_64::@dispatch}[%c4, %c1, %c1] {
ro %arg0_capture[%c0 for %c16] : !stream.resource<external>{%c16},
ro %arg1_capture[%c0 for %c16] : !stream.resource<external>{%c16},
wo %result_capture[%c0 for %c16] : !stream.resource<external>{%c16}
diff --git a/compiler/src/iree/compiler/Dialect/HAL/Transforms/test/dump_executable_benchmarks.mlir b/compiler/src/iree/compiler/Dialect/HAL/Transforms/test/dump_executable_benchmarks.mlir
index d9c1a3e..c13bb59 100644
--- a/compiler/src/iree/compiler/Dialect/HAL/Transforms/test/dump_executable_benchmarks.mlir
+++ b/compiler/src/iree/compiler/Dialect/HAL/Transforms/test/dump_executable_benchmarks.mlir
@@ -124,7 +124,7 @@
%result, %result_timepoint = stream.resource.alloca uninitialized : !stream.resource<transient>{%c128} => !stream.timepoint
%6 = stream.cmd.execute await(%result_timepoint) => with(%result as %result_capture: !stream.resource<transient>{%c128}) {
// Dispatches with static and dynamic args.
- stream.cmd.dispatch @ex0::@dispatch0[%c512](%c100_i32, %c200_i32 : i32, i32) {
+ stream.cmd.dispatch @ex0::@embedded_elf_x86_64::@dispatch0[%c512](%c100_i32, %c200_i32 : i32, i32) {
ro %result_capture[%c0 for %c32] : !stream.resource<transient>{%c128},
rw %result_capture[%c32 for %c32] : !stream.resource<transient>{%c128},
rw %result_capture[%c64 for %c32] : !stream.resource<transient>{%c128}
@@ -135,7 +135,7 @@
]}
// NOTE: today the dynamic args will prevent us from generating
// benchmarks. We could handle this better by tracking alignment and such.
- stream.cmd.dispatch @ex0::@dispatch0[%c512](%c300_i32, %dynamic_arg : i32, i32) {
+ stream.cmd.dispatch @ex0::@embedded_elf_x86_64::@dispatch0[%c512](%c300_i32, %dynamic_arg : i32, i32) {
ro %result_capture[%c0 for %c32] : !stream.resource<transient>{%c128},
rw %result_capture[%c32 for %c32] : !stream.resource<transient>{%c128},
rw %result_capture[%c64 for %c32] : !stream.resource<transient>{%c128}
@@ -147,21 +147,21 @@
// Multiple dispatches to a single entry point.
// Dispatches are deduplicated and the two 128x32x1 should combine.
- stream.cmd.dispatch @ex0::@dispatch1[%c512, %c1] {
+ stream.cmd.dispatch @ex0::@embedded_elf_x86_64::@dispatch1[%c512, %c1] {
ro %result_capture[%c0 for %c64] : !stream.resource<transient>{%c128},
rw %result_capture[%c64 for %c32] : !stream.resource<transient>{%c128}
} attributes {hal.interface.bindings = [
#hal.interface.binding<0, 0>,
#hal.interface.binding<0, 1>
]}
- stream.cmd.dispatch @ex0::@dispatch1[%c128, %c32] {
+ stream.cmd.dispatch @ex0::@embedded_elf_x86_64::@dispatch1[%c128, %c32] {
ro %result_capture[%c0 for %c64] : !stream.resource<transient>{%c128},
rw %result_capture[%c64 for %c32] : !stream.resource<transient>{%c128}
} attributes {hal.interface.bindings = [
#hal.interface.binding<0, 0>,
#hal.interface.binding<0, 1>
]}
- stream.cmd.dispatch @ex0::@dispatch1[%c128, %c32] {
+ stream.cmd.dispatch @ex0::@embedded_elf_x86_64::@dispatch1[%c128, %c32] {
ro %result_capture[%c0 for %c64] : !stream.resource<transient>{%c128},
rw %result_capture[%c64 for %c32] : !stream.resource<transient>{%c128}
} attributes {hal.interface.bindings = [
diff --git a/compiler/src/iree/compiler/Dialect/HAL/Transforms/test/materialize_interfaces.mlir b/compiler/src/iree/compiler/Dialect/HAL/Transforms/test/materialize_interfaces.mlir
index e13bb60..fdc093a 100644
--- a/compiler/src/iree/compiler/Dialect/HAL/Transforms/test/materialize_interfaces.mlir
+++ b/compiler/src/iree/compiler/Dialect/HAL/Transforms/test/materialize_interfaces.mlir
@@ -54,7 +54,7 @@
%c2 = arith.constant 2 : index
%0 = stream.resource.alloc uninitialized : !stream.resource<transient>{%arg2}
%1 = stream.cmd.execute with(%arg0 as %arg4: !stream.resource<constant>{%arg2}, %arg1 as %arg5: !stream.resource<transient>{%arg2}, %0 as %arg6: !stream.resource<transient>{%arg2}) {
- // CHECK: stream.cmd.dispatch @ex_workgroups::@entry
+ // CHECK: stream.cmd.dispatch {@ex_workgroups::@embedded_elf_arm_64::@entry, @ex_workgroups::@embedded_elf_x86_64::@entry}
// CHECK: attributes {
// CHECK-SAME: hal.interface.bindings = [
// CHECK-SAME: #hal.interface.binding<0, 0>,
diff --git a/compiler/src/iree/compiler/Dialect/Stream/IR/StreamOps.cpp b/compiler/src/iree/compiler/Dialect/Stream/IR/StreamOps.cpp
index a14b474..fb7208b 100644
--- a/compiler/src/iree/compiler/Dialect/Stream/IR/StreamOps.cpp
+++ b/compiler/src/iree/compiler/Dialect/Stream/IR/StreamOps.cpp
@@ -2122,29 +2122,66 @@
LogicalResult CmdDispatchOp::verifySymbolUses(
SymbolTableCollection &symbolTable) {
Operation *op = getOperation();
- auto exportOp =
- symbolTable.lookupNearestSymbolFrom<IREE::Stream::ExecutableExportOp>(
- op, getEntryPoint());
- if (!exportOp) {
- // TODO(benvanik): there are a lot of tests that are assuming this is not
- // verified. We'll need to go add dummy executables for all of them. Today
- // we just bail on the verifier if the symbol isn't found.
- //
- // Should be:
- // return op->emitOpError() << "undefined entry point: " << entry_point();
- return success();
+ auto entryPointRefs = getEntryPointRefs();
+ if (entryPointRefs.empty()) {
+ return emitOpError() << "at least one entry point must be defined";
}
+ for (auto entryPointAttr : entryPointRefs) {
+ auto exportOp =
+ symbolTable.lookupNearestSymbolFrom<IREE::Stream::ExecutableExportOp>(
+ op, entryPointAttr);
+ if (!exportOp) {
+ // TODO(benvanik): there are a lot of tests that are assuming this is not
+ // verified. We'll need to go add dummy executables for all of them. Today
+ // we just bail on the verifier if the symbol isn't found.
+ //
+ // Should be:
+ // return op->emitOpError() << "undefined entry point: " <<
+ // entry_point();
+ return success();
+ }
- // Verify that the workload parameters captured match the target export.
- if (failed(verifyDispatchWorkload(op, exportOp, getWorkload()))) {
- return failure();
+ // Verify that the workload parameters captured match the target export.
+ if (failed(verifyDispatchWorkload(op, exportOp, getWorkload()))) {
+ return failure();
+ }
+
+ // TODO(benvanik): verify that the target function has matching operands.
}
-
- // TODO(benvanik): verify that the target function has matching operands.
-
return success();
}
+static ParseResult parseDispatchEntryPoints(OpAsmParser &parser,
+ ArrayAttr &entryPointAttrsArray) {
+ SmallVector<Attribute> entryPointAttrs;
+ if (succeeded(parser.parseOptionalLBrace())) {
+ do {
+ SymbolRefAttr entryPointAttr;
+ if (failed(parser.parseAttribute(entryPointAttr))) return failure();
+ entryPointAttrs.push_back(entryPointAttr);
+ } while (succeeded(parser.parseOptionalComma()));
+ if (failed(parser.parseRBrace())) return failure();
+ } else {
+ SymbolRefAttr entryPointAttr;
+ if (failed(parser.parseAttribute(entryPointAttr))) return failure();
+ entryPointAttrs.push_back(entryPointAttr);
+ }
+ entryPointAttrsArray = parser.getBuilder().getArrayAttr(entryPointAttrs);
+ return success();
+}
+
+static void printDispatchEntryPoints(OpAsmPrinter &p, Operation *op,
+ ArrayAttr entryPointAttrs) {
+ if (entryPointAttrs.size() == 1) {
+ p.printAttribute(entryPointAttrs.getValue().front());
+ } else {
+ p << '{';
+ llvm::interleaveComma(entryPointAttrs, p.getStream(),
+ [&](Attribute attr) { p.printAttribute(attr); });
+ p << '}';
+ }
+}
+
static ParseResult parseDispatchResources(
OpAsmParser &parser,
SmallVectorImpl<OpAsmParser::UnresolvedOperand> &resources,
diff --git a/compiler/src/iree/compiler/Dialect/Stream/IR/StreamOps.td b/compiler/src/iree/compiler/Dialect/Stream/IR/StreamOps.td
index a93b9f1..91c541e 100644
--- a/compiler/src/iree/compiler/Dialect/Stream/IR/StreamOps.td
+++ b/compiler/src/iree/compiler/Dialect/Stream/IR/StreamOps.td
@@ -2481,7 +2481,7 @@
let arguments = (ins
Variadic<Index>:$workload,
- SymbolRefAttr:$entry_point,
+ SymbolRefArrayAttr:$entry_points,
Variadic<Stream_PrimitiveType>:$uniform_operands,
Variadic<Stream_AnyStreamResource>:$resources,
Variadic<Stream_Size>:$resource_sizes,
@@ -2492,7 +2492,7 @@
let results = (outs);
let assemblyFormat = [{
- $entry_point
+ custom<DispatchEntryPoints>($entry_points)
(`[` $workload^ `]`)? ``
(`(` $uniform_operands^ `:` type($uniform_operands) `)`)? `{`
custom<DispatchResources>($resources, type($resources), $resource_sizes,
@@ -2503,6 +2503,13 @@
}];
let extraClassDeclaration = [{
+ auto getEntryPointRefs() {
+ return getEntryPoints().getAsRange<SymbolRefAttr>();
+ }
+ void forEachEntryPointAttr(std::function<void(SymbolRefAttr)> fn) {
+ for (auto entryPointAttr : getEntryPointRefs()) fn(entryPointAttr);
+ }
+
Value getOperandSize(unsigned idx) {
return findValueSizeInList(
idx - getODSOperandIndexAndLength(2).first,
diff --git a/compiler/src/iree/compiler/Dialect/Stream/IR/test/cmd_ops.mlir b/compiler/src/iree/compiler/Dialect/Stream/IR/test/cmd_ops.mlir
index b2ec087..11154de 100644
--- a/compiler/src/iree/compiler/Dialect/Stream/IR/test/cmd_ops.mlir
+++ b/compiler/src/iree/compiler/Dialect/Stream/IR/test/cmd_ops.mlir
@@ -106,11 +106,11 @@
%c5 = arith.constant 5 : index
%c128 = arith.constant 128 : index
%0 = stream.cmd.execute with(%arg0 as %arg4: !stream.resource<transient>{%arg1}, %arg2 as %arg5: !stream.resource<external>{%arg3}) {
- // CHECK: stream.cmd.dispatch @executable::@dispatch[%c1, %c2, %c3](%c4, %c5 : index, index) {
+ // CHECK: stream.cmd.dispatch {@executable::@dispatch0, @executable::@dispatch1}[%c1, %c2, %c3](%c4, %c5 : index, index) {
// CHECK-NEXT: ro %arg4[%c0 for %c128] : !stream.resource<transient>{%arg1},
// CHECK-NEXT: wo %arg5[%c0 for %c128] : !stream.resource<external>{%arg3}
// CHECK-NEXT: }
- stream.cmd.dispatch @executable::@dispatch[%c1, %c2, %c3](%c4, %c5 : index, index) {
+ stream.cmd.dispatch {@executable::@dispatch0, @executable::@dispatch1}[%c1, %c2, %c3](%c4, %c5 : index, index) {
ro %arg4[%c0 for %c128] : !stream.resource<transient>{%arg1},
wo %arg5[%c0 for %c128] : !stream.resource<external>{%arg3}
}
diff --git a/compiler/src/iree/compiler/Dialect/Stream/Transforms/AnnotateDispatchArguments.cpp b/compiler/src/iree/compiler/Dialect/Stream/Transforms/AnnotateDispatchArguments.cpp
index 3ab708f..9de64fd 100644
--- a/compiler/src/iree/compiler/Dialect/Stream/Transforms/AnnotateDispatchArguments.cpp
+++ b/compiler/src/iree/compiler/Dialect/Stream/Transforms/AnnotateDispatchArguments.cpp
@@ -340,9 +340,11 @@
// Find all dispatches and bucket by their target entry point.
rootOp->walk([&](IREE::Stream::CmdDispatchOp dispatchOp) {
- auto exportOp = explorer.getSymbolTables().lookupNearestSymbolFrom(
- dispatchOp, dispatchOp.getEntryPoint());
- entryDispatchMap[exportOp].push_back(dispatchOp);
+ dispatchOp.forEachEntryPointAttr([&](SymbolRefAttr entryPointAttr) {
+ auto exportOp = explorer.getSymbolTables().lookupNearestSymbolFrom(
+ dispatchOp, entryPointAttr);
+ entryDispatchMap[exportOp].push_back(dispatchOp);
+ });
});
}
diff --git a/compiler/src/iree/compiler/Dialect/Stream/Transforms/DumpStatistics.cpp b/compiler/src/iree/compiler/Dialect/Stream/Transforms/DumpStatistics.cpp
index d3b31fc..c9dedea 100644
--- a/compiler/src/iree/compiler/Dialect/Stream/Transforms/DumpStatistics.cpp
+++ b/compiler/src/iree/compiler/Dialect/Stream/Transforms/DumpStatistics.cpp
@@ -78,12 +78,14 @@
}
for (auto executeOp : executeOps) {
executeOp.walk([&](IREE::Stream::CmdDispatchOp dispatchOp) {
- auto exportOp = cast<IREE::Stream::ExecutableExportOp>(
- symbolTable.lookupSymbolIn(moduleOp, dispatchOp.getEntryPoint()));
- assert(exportOp && "missing executable/export");
- auto funcOp = exportOp.lookupFunctionRef();
- assert(funcOp && "missing exported function");
- exportDispatchOps[funcOp].push_back(dispatchOp);
+ dispatchOp.forEachEntryPointAttr([&](SymbolRefAttr entryPointAttr) {
+ auto exportOp = cast<IREE::Stream::ExecutableExportOp>(
+ symbolTable.lookupSymbolIn(moduleOp, entryPointAttr));
+ assert(exportOp && "missing executable/export");
+ auto funcOp = exportOp.lookupFunctionRef();
+ assert(funcOp && "missing exported function");
+ exportDispatchOps[funcOp].push_back(dispatchOp);
+ });
});
}
}
@@ -460,10 +462,10 @@
workloadSum *= dimValue;
}
}
- os << llvm::formatv(R"({0},"dispatch","{1}",,{2},"{3}",{4},{5})",
- depth, op.getEntryPoint(), workloadSum,
- workloadStr, op.getUniformOperands().size(),
- op.getResources().size());
+ os << llvm::formatv(
+ R"({0},"dispatch","{1}",,{2},"{3}",{4},{5})", depth,
+ *op.getEntryPointRefs().begin(), workloadSum, workloadStr,
+ op.getUniformOperands().size(), op.getResources().size());
os << "\n";
});
};
diff --git a/compiler/src/iree/compiler/Dialect/Stream/Transforms/FoldUniformOperands.cpp b/compiler/src/iree/compiler/Dialect/Stream/Transforms/FoldUniformOperands.cpp
index 92d41cf..a18a52c 100644
--- a/compiler/src/iree/compiler/Dialect/Stream/Transforms/FoldUniformOperands.cpp
+++ b/compiler/src/iree/compiler/Dialect/Stream/Transforms/FoldUniformOperands.cpp
@@ -267,9 +267,11 @@
DenseMap<Operation *, SmallVector<IREE::Stream::CmdDispatchOp>>
entryDispatchMap;
getOperation()->walk([&](IREE::Stream::CmdDispatchOp dispatchOp) {
- auto exportOp = symbolTable.lookupNearestSymbolFrom(
- dispatchOp, dispatchOp.getEntryPoint());
- entryDispatchMap[exportOp].push_back(dispatchOp);
+ dispatchOp.forEachEntryPointAttr([&](SymbolRefAttr entryPointAttr) {
+ auto exportOp =
+ symbolTable.lookupNearestSymbolFrom(dispatchOp, entryPointAttr);
+ entryDispatchMap[exportOp].push_back(dispatchOp);
+ });
});
// Optimize each dispatch op.
diff --git a/compiler/src/iree/compiler/Dialect/Stream/Transforms/FuseDispatchBindings.cpp b/compiler/src/iree/compiler/Dialect/Stream/Transforms/FuseDispatchBindings.cpp
index a56e653..66530b3 100644
--- a/compiler/src/iree/compiler/Dialect/Stream/Transforms/FuseDispatchBindings.cpp
+++ b/compiler/src/iree/compiler/Dialect/Stream/Transforms/FuseDispatchBindings.cpp
@@ -299,7 +299,7 @@
OpBuilder builder(dispatchOp);
auto newOp = builder.create<IREE::Stream::CmdDispatchOp>(
dispatchOp.getLoc(), dispatchOp.getWorkload(),
- dispatchOp.getEntryPointAttr(), newOperands, newResources,
+ dispatchOp.getEntryPointsAttr(), newOperands, newResources,
newResourceSizes, newOffsets, newLengths,
builder.getArrayAttr(newAccesses));
(void)newOp;
@@ -433,9 +433,11 @@
DenseMap<Operation *, SmallVector<IREE::Stream::CmdDispatchOp>>
entryDispatchMap;
getOperation()->walk([&](IREE::Stream::CmdDispatchOp dispatchOp) {
- auto exportOp = symbolTable.lookupNearestSymbolFrom(
- dispatchOp, dispatchOp.getEntryPoint());
- entryDispatchMap[exportOp].push_back(dispatchOp);
+ dispatchOp.forEachEntryPointAttr([&](SymbolRefAttr entryPointAttr) {
+ auto exportOp =
+ symbolTable.lookupNearestSymbolFrom(dispatchOp, entryPointAttr);
+ entryDispatchMap[exportOp].push_back(dispatchOp);
+ });
});
// Perform fusion for each executable entry point using all known dispatches
diff --git a/compiler/src/iree/compiler/Dialect/Stream/Transforms/PackDispatchOperands.cpp b/compiler/src/iree/compiler/Dialect/Stream/Transforms/PackDispatchOperands.cpp
index 9973082..7509ca4 100644
--- a/compiler/src/iree/compiler/Dialect/Stream/Transforms/PackDispatchOperands.cpp
+++ b/compiler/src/iree/compiler/Dialect/Stream/Transforms/PackDispatchOperands.cpp
@@ -337,12 +337,15 @@
// Walk the module and update all dispatch operands.
getOperation()->walk([&](IREE::Stream::CmdDispatchOp dispatchOp) {
- auto exportOp =
- symbolTable.lookupNearestSymbolFrom<IREE::Stream::ExecutableExportOp>(
- dispatchOp, dispatchOp.getEntryPoint());
- if (exportOp) {
- updateDispatchOp(dispatchOp, exportOp);
- }
+ dispatchOp.forEachEntryPointAttr([&](SymbolRefAttr entryPointAttr) {
+ auto exportOp =
+ symbolTable
+ .lookupNearestSymbolFrom<IREE::Stream::ExecutableExportOp>(
+ dispatchOp, entryPointAttr);
+ if (exportOp) {
+ updateDispatchOp(dispatchOp, exportOp);
+ }
+ });
return WalkResult::advance();
});
}
diff --git a/compiler/src/iree/compiler/Dialect/Stream/Transforms/ScheduleAllocation.cpp b/compiler/src/iree/compiler/Dialect/Stream/Transforms/ScheduleAllocation.cpp
index 9782741..7cc0ae1 100644
--- a/compiler/src/iree/compiler/Dialect/Stream/Transforms/ScheduleAllocation.cpp
+++ b/compiler/src/iree/compiler/Dialect/Stream/Transforms/ScheduleAllocation.cpp
@@ -693,9 +693,10 @@
}
auto newOp = builder.create<IREE::Stream::CmdDispatchOp>(
- asyncOp.getLoc(), asyncOp.getWorkload(), asyncOp.getEntryPoint(),
- newOperands, newResources, newResourceSizes, newResourceOffsets,
- newResourceLengths, builder.getArrayAttr(newResourceAccesses));
+ asyncOp.getLoc(), asyncOp.getWorkload(),
+ builder.getArrayAttr({asyncOp.getEntryPoint()}), newOperands,
+ newResources, newResourceSizes, newResourceOffsets, newResourceLengths,
+ builder.getArrayAttr(newResourceAccesses));
newOp->setDialectAttrs(asyncOp->getDialectAttrs());
asyncOp.erase();
return success();
diff --git a/compiler/src/iree/compiler/Dialect/Stream/Transforms/SpecializeDispatches.cpp b/compiler/src/iree/compiler/Dialect/Stream/Transforms/SpecializeDispatches.cpp
index 071de01..0b72cfa 100644
--- a/compiler/src/iree/compiler/Dialect/Stream/Transforms/SpecializeDispatches.cpp
+++ b/compiler/src/iree/compiler/Dialect/Stream/Transforms/SpecializeDispatches.cpp
@@ -343,9 +343,11 @@
DenseMap<Operation *, SmallVector<IREE::Stream::CmdDispatchOp>>
entryDispatchMap;
getOperation()->walk([&](IREE::Stream::CmdDispatchOp dispatchOp) {
- auto exportOp = symbolTable.lookupNearestSymbolFrom(
- dispatchOp, dispatchOp.getEntryPoint());
- entryDispatchMap[exportOp].push_back(dispatchOp);
+ dispatchOp.forEachEntryPointAttr([&](SymbolRefAttr entryPointAttr) {
+ auto exportOp =
+ symbolTable.lookupNearestSymbolFrom(dispatchOp, entryPointAttr);
+ entryDispatchMap[exportOp].push_back(dispatchOp);
+ });
});
// Optimize each dispatchable function and its dispatch sites.
diff --git a/compiler/src/iree/compiler/Modules/HAL/Inline/Transforms/InlineExecutables.cpp b/compiler/src/iree/compiler/Modules/HAL/Inline/Transforms/InlineExecutables.cpp
index 496c051..fba157a 100644
--- a/compiler/src/iree/compiler/Modules/HAL/Inline/Transforms/InlineExecutables.cpp
+++ b/compiler/src/iree/compiler/Modules/HAL/Inline/Transforms/InlineExecutables.cpp
@@ -61,14 +61,25 @@
// Annotate all dispatches with the target function.
for (auto funcOp : moduleOp.getOps<mlir::FunctionOpInterface>()) {
- funcOp.walk([&](IREE::Stream::CmdDispatchOp dispatchOp) {
+ auto result = funcOp.walk([&](IREE::Stream::CmdDispatchOp dispatchOp) {
// Specify new target function that conversion can use to make the call.
+ // We only support single variant dispatches when inline.
+ auto entryPointAttrs = dispatchOp.getEntryPoints().getValue();
+ if (entryPointAttrs.size() != 1) {
+ dispatchOp.emitOpError()
+ << "multiple variant targets not supported with the inline HAL";
+ return WalkResult::interrupt();
+ }
auto targetFuncName =
- exportToFuncMap[dispatchOp.getEntryPoint()].cast<StringAttr>();
+ exportToFuncMap[entryPointAttrs.front()].cast<StringAttr>();
assert(targetFuncName && "missing mapping");
dispatchOp->setAttr("hal_inline.target",
FlatSymbolRefAttr::get(targetFuncName));
+ return WalkResult::advance();
});
+ if (result.wasInterrupted()) {
+ return signalPassFailure();
+ }
}
}
diff --git a/compiler/src/iree/compiler/Modules/HAL/Loader/Conversion/StreamToHALLoader/Patterns.cpp b/compiler/src/iree/compiler/Modules/HAL/Loader/Conversion/StreamToHALLoader/Patterns.cpp
index aa659b9..4eece6d 100644
--- a/compiler/src/iree/compiler/Modules/HAL/Loader/Conversion/StreamToHALLoader/Patterns.cpp
+++ b/compiler/src/iree/compiler/Modules/HAL/Loader/Conversion/StreamToHALLoader/Patterns.cpp
@@ -46,10 +46,22 @@
ConversionPatternRewriter &rewriter) const override {
auto loc = dispatchOp.getLoc();
+ // TODO(benvanik): support a lightweight switch builder for picking variants
+ // that doesn't pull in the full HAL dialect - today the
+ // DeviceSwitchRewriter needs a !hal.device and its query methods.
+ // For now we bail if there's multiple.
+ auto entryPointAttrs = dispatchOp.getEntryPoints().getValue();
+ if (entryPointAttrs.size() != 1) {
+ return rewriter.notifyMatchFailure(dispatchOp,
+ "multiple variant targets not yet "
+ "supported in the inline HAL loader");
+ }
+ auto entryPointAttr = entryPointAttrs.front().cast<SymbolRefAttr>();
+
// Get the handle to the executable that is compatible with our device.
auto executableOp =
cast<IREE::HAL::ExecutableOp>(SymbolTable::lookupNearestSymbolFrom(
- dispatchOp, dispatchOp.getEntryPoint().getRootReference()));
+ dispatchOp, entryPointAttr.getRootReference()));
assert(executableOp && "dispatch target executable op not found");
// For now we aren't doing loader support checks. We should, though.
@@ -70,13 +82,12 @@
auto exportOps = variantOp.getOps<IREE::HAL::ExecutableExportOp>();
auto exportIt =
llvm::find_if(exportOps, [&](IREE::HAL::ExecutableExportOp op) {
- return op.getNameAttr() ==
- dispatchOp.getEntryPoint().getLeafReference();
+ return op.getNameAttr() == entryPointAttr.getLeafReference();
});
if (exportIt == exportOps.end()) {
return variantOp.emitError()
<< "hal.executable.variant is missing the entry point for "
- << dispatchOp.getEntryPoint();
+ << entryPointAttr;
}
auto exportOp = *exportIt;
diff --git a/samples/custom_dispatch/cpu/embedded/example_hal.mlir b/samples/custom_dispatch/cpu/embedded/example_hal.mlir
index dbdf89b..e9c9b10 100644
--- a/samples/custom_dispatch/cpu/embedded/example_hal.mlir
+++ b/samples/custom_dispatch/cpu/embedded/example_hal.mlir
@@ -243,9 +243,7 @@
%dim_i32 = arith.index_cast %dim : index to i32
// Dispatch a basic `ret = lhs * rhs` using an external function.
- // This form (@executable::@export) allows for automatic variant selection
- // when multi-targeting (@x86_64 will be chosen if available).
- %0 = flow.dispatch @executable::@simple_mul[%dim](%dim_i32, %arg0, %arg1) {
+ %0 = flow.dispatch @executable::@x86_64::@simple_mul[%dim](%dim_i32, %arg0, %arg1) {
// Bindings are automatically inferred when possible as part of the ABI
// but can be overridden if the user wants to use features such as sparse
// bindings or multiple descriptor sets. To do so the
