Skip to content

[BugFix][TIRx] Fix bad-optional-access in BF16/FP8 legalize passes for target-less PrimFuncs#19383

Merged
tlopex merged 1 commit into
apache:mainfrom
swjng:fix/tirx-dtype-legalize-no-target
Apr 10, 2026
Merged

[BugFix][TIRx] Fix bad-optional-access in BF16/FP8 legalize passes for target-less PrimFuncs#19383
tlopex merged 1 commit into
apache:mainfrom
swjng:fix/tirx-dtype-legalize-no-target

Conversation

@swjng

@swjng swjng commented Apr 10, 2026

Copy link
Copy Markdown
Contributor

Problem

BF16ComputeLegalize, BF16StorageLegalize, FP8ComputeLegalize, and FP8StorageLegalize all call f->GetAttr<Target>(kTarget).value() unconditionally. Host-side helper PrimFuncs produced during Relax lowering (e.g. for reshape, mean) carry no target attribute; .value() on an empty Optional<Target> aborts at runtime:

terminate called after throwing an instance of 'std::bad_optional_access'

This surfaces whenever tvm.compile targets CUDA on a model with BF16/FP8 support compiled in, because the Relax lowering pipeline mixes target-annotated GPU kernels with target-less CPU helpers in the same module.

Fix

Retrieve the target into opt_target and combine the defined() check with CheckDataTypeSupport:

if (opt_target.defined() && CheckDataTypeSupport(opt_target.value(), "tvm.contrib.nvcc.supports_bf16")) {
    return f;
}
  • If the target is absent: legalization runs. This is safe — it is a no-op when no BF16/FP8 types are present, and it ensures correctness for the edge case where a target-less function does contain such types.
  • If the target confirms native support: legalization is skipped as before.
  • If the target lacks native support: legalization runs as before.

This is consistent with how explicit host targets (e.g. llvm) are handled, which also go through legalization. Applied identically to all four legalize pass lambdas.

@gemini-code-assist gemini-code-assist Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Code Review

This pull request updates the BF16 and FP8 legalization passes to safely handle functions where the target attribute is missing, preventing potential crashes. However, the current implementation returns early when a target is undefined, which inadvertently skips legalization for target-less functions. Feedback suggests combining the target existence check with the capability check to ensure consistency with host targets and to allow the legalizer to run on functions without an explicit target.

Comment thread src/tirx/transform/unsupported_dtype_legalize.cc Outdated
Comment thread src/tirx/transform/unsupported_dtype_legalize.cc Outdated
Comment thread src/tirx/transform/unsupported_dtype_legalize.cc Outdated
Comment thread src/tirx/transform/unsupported_dtype_legalize.cc Outdated
…r target-less PrimFuncs

BF16ComputeLegalize, BF16StorageLegalize, FP8ComputeLegalize, and
FP8StorageLegalize all called f->GetAttr<Target>(kTarget).value()
unconditionally. Host-side helper PrimFuncs produced during Relax
lowering carry no target attribute; .value() on an empty
Optional<Target> aborts at runtime with a bad optional access.

Check opt_target.defined() and return f unchanged when absent.
Host-only functions do not use BF16/FP8 on device and need no
legalization.
@swjng swjng force-pushed the fix/tirx-dtype-legalize-no-target branch from 532f75b to 7fa7c8f Compare April 10, 2026 15:17
@tlopex tlopex merged commit f4cf9f5 into apache:main Apr 10, 2026
15 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants