OpenXLA PR #48532: Why XLA Paired PTX 8.5 With an sm

OpenXLA PR #48532 fixes an invalid pairing between a GPU architecture and the virtual instruction set used to target it.

XLA selected the PTX version from the CUDA toolchain version. In the affected path, CUDA 12.6 mapped to PTX 8.5.

The detected GPU was NVIDIA Blackwell with compute capability 12.0, represented as sm_120. NVIDIA added support for that target in PTX ISA 8.7.

LLVM therefore received an impossible combination:

PTX feature: +ptx85
Target CPU:   sm_120

and stopped compilation with:

LLVM ERROR: ptx87 is required for target sm_120

The patch makes PTX selection depend on the target architecture as well as the CUDA version. When the compute-capability major version is 12 or newer, XLA raises any lower result to PTX 8.7.

Three version numbers were being treated as one decision

The failure is easier to understand when three related identifiers are separated.

CUDA version

A CUDA toolkit or assembler version tells XLA which NVIDIA toolchain capabilities may be available.

XLA already maintained a mapping such as:

CUDA 12.4 → PTX 8.4
CUDA 12.6 → PTX 8.5
CUDA 12.8 → PTX 8.7
CUDA 12.9 → PTX 8.8
CUDA 13.0 → PTX 9.0

That mapping provides a toolchain-derived PTX candidate.

PTX ISA version

PTX is NVIDIA’s virtual instruction set. LLVM emits textual PTX, including a .version directive, before the code is assembled or JIT-compiled for a GPU.

A PTX version determines which target names and instructions the module can legally use.

SM target

The SM target identifies the GPU architecture for which the generated code is intended.

For the reported consumer Blackwell device, that target is:

sm_120

NVIDIA’s PTX ISA 8.7 documentation explicitly introduces support for sm_120 and sm_120a.

CUDA version, PTX ISA version, and SM target are related. They are not interchangeable.

A valid compiler configuration must satisfy both the toolchain environment and the minimum ISA required by the target.

The old function knew only the CUDA side

XLA’s helper was named:

DetermineHighestSupportedPtxVersionFromCudaVersion

Before this change, it accepted only a CUDA semantic version.

For CUDA releases from 12.6 through 12.9, it handled NVIDIA’s shifted mapping by subtracting one from the CUDA minor version:

CUDA 12.6 → PTX 8.5
CUDA 12.8 → PTX 8.7
CUDA 12.9 → PTX 8.8

That result was passed to LLVM as a target feature such as:

+ptx85

Separately, XLA selected an LLVM target CPU from the GPU compute capability:

sm_120

Each choice looked reasonable in isolation. Together they violated LLVM’s target requirements.

The old mapping was not wrong for CUDA 12.6 by itself. It was incomplete once the selected target was sm_120.

Why detecting the GPU was not enough

The linked TensorFlow report came from a system with an RTX 5060 Ti. TensorFlow could list the GPU successfully, and unrelated PyTorch workloads reportedly ran on the same machine.

The failure appeared only when TensorFlow entered XLA JIT compilation.

That distinction matters:

GPU discovery
≠
ability to generate valid code for that GPU

Device enumeration proves that the runtime can identify the hardware.

XLA still has to create a valid compiler target, emit a compatible PTX module, and hand that module to an assembler or driver that can translate it for the device.

The reported environment exposed a mixed-version boundary:

new Blackwell GPU and new driver
        ↓
TensorFlow/XLA path constrained by an older CUDA-derived PTX selection
        ↓
LLVM target set to sm_120
        ↓
PTX feature left at 8.5
        ↓
compiler abort before kernel execution

StarDist also failed because it used TensorFlow as its backend. The application above TensorFlow was not the source of the target mismatch.

LLVM rejected the module before runtime execution

PTX 8.5 does not contain the sm_120 target introduced in PTX 8.7.

LLVM’s NVPTX backend validates the selected processor against the explicitly enabled PTX feature. When the feature is below the architecture’s minimum, code generation stops instead of emitting a module that claims an unsupported target.

This was a useful failure.

Without that check, a lower layer might receive PTX whose header and target declarations disagree, producing a later and less understandable toolchain error.

The fatal diagnostic identified the exact broken contract:

selected target minimum: PTX 8.7
selected XLA feature:    PTX 8.5

The patch adds the target to PTX selection

The shared helper now accepts a second input:

DetermineHighestSupportedPtxVersionFromCudaVersion(
    cuda_version,
    compute_capability_major)

It first performs the existing CUDA-to-PTX mapping.

It then applies the target floor:

if (compute_capability_major >= 12) {
  if (ptx_version < SemanticVersion{8, 7, 0}) {
    ptx_version = {8, 7, 0};
  }
}

For the failing combination:

CUDA-derived candidate: PTX 8.5
sm_120 minimum:          PTX 8.7
final XLA selection:     PTX 8.7

LLVM now receives a PTX feature that can legally name the selected target.

The patch does not upgrade a CUDA toolkit. It changes the PTX feature XLA asks LLVM to emit.

Whether the resulting PTX can complete the rest of the toolchain still depends on the assembler or driver available in the actual deployment.

Review moved the policy into the shared helper

The original proposal focused on enforcing the minimum near the NVPTX target-machine construction.

Public review asked whether the shared PTX-version helper should own the correction instead.

That change in placement matters because XLA had more than one consumer of the version decision.

The final patch updates both:

LLVM NVPTX target-machine construction

and

MLIR NVIDIA float-lowering configuration

The LLVM path uses the selected PTX version to build the +ptxNN target feature.

The MLIR path passes PTX and compute-capability information into NVIDIA-specific float conversion. That pass uses version information when deciding which low-precision operations can remain native and which require expansion.

A local clamp in only nvptx_backend.cc could have left earlier lowering decisions working from the old PTX assumption.

Moving the rule into one utility and passing the target at both call sites keeps the architecture constraint closer to the version authority.

The selected version now represents two constraints

The helper’s name still emphasizes the CUDA version, but its result now includes a target requirement.

Conceptually, the decision is:

PTX selected from CUDA environment
        ↓
compare with target’s minimum PTX
        ↓
raise to the target minimum when necessary

This is not the same as proving that every component of an older CUDA installation understands the raised PTX version.

If XLA finds an older standalone ptxas, that binary may remain a separate compatibility boundary. A sufficiently new NVIDIA driver can JIT newer PTX even when the framework was built against an older toolkit, but the exact route depends on packaging and runtime configuration.

The merged change repairs the LLVM-side contradiction. It should not be described as a universal upgrade of every CUDA component in a TensorFlow environment.

The existing unit test does not cover the new branch

OpenXLA already has a parameterized test for the CUDA-to-PTX mapping.

It checks values such as:

12.6 → 8.5
12.8 → 8.7
12.9 → 8.8

The function declaration gives compute_capability_major a default value of zero, so those existing one-argument tests continue to exercise only the original mapping.

PR #48532 does not add a case equivalent to:

DetermineHighestSupportedPtxVersionFromCudaVersion(
    {12, 6, 0},
    12)
== {8, 7, 0}

It also does not add a checked-in end-to-end test that constructs an sm_120 target under the CUDA 12.6 mapping and verifies successful PTX emission.

The implementation is small and directly matches the reported LLVM diagnostic. The absence of a target-aware regression test is still an important limitation because the new behavior exists entirely in that second argument.

This is an sm_120 repair, not a complete architecture table

The condition uses:

compute capability major >= 12

and applies a floor of PTX 8.7.

That is sufficient for sm_120, the target named by the issue.

NVIDIA introduced sm_121 in PTX ISA 8.8. OpenXLA’s own target tests recognize sm_121, while the new helper receives only the major number and cannot distinguish 12.0 from 12.1.

A CUDA-derived result below 8.8 paired with sm_121 could therefore require another target-specific rule.

This is a targeted sm_120 fix, not a general table of per-architecture PTX minimums.

A more general design would derive the minimum from the full resolved compute capability or from LLVM’s target metadata rather than encode one major-version floor.

That broader redesign is outside this PR’s stated scope.

What the public evidence establishes

The merged source establishes that XLA now passes compute-capability information into the shared PTX selector at both visible call sites and raises major-12 targets to at least PTX 8.7.

The original PR received public approval after the rule was moved into the utility. Copybara then imported the change into OpenXLA, and a subsequent TensorFlow synchronization imported it into the TensorFlow repository. The linked TensorFlow issue was closed after that synchronization.

A closed issue is evidence of project resolution, not a public rerun on every affected wheel.

The public thread does not contain a post-release confirmation from the original reporter using a newly published TensorFlow package.

The strongest supported conclusion is therefore:

XLA selected PTX 8.5 from a CUDA 12.6 constraint while targeting sm_120, which requires PTX 8.7. PR #48532 adds the missing target floor to the shared version-selection path and removes that specific LLVM configuration mismatch.

The compiler contract restored by the change

A target architecture is not only an optimization preference.

It constrains the minimum language in which valid code for that target can be expressed.

The corrected decision path is:

identify the CUDA toolchain constraint
        ↓
identify the GPU target
        ↓
derive the PTX candidate
        ↓
enforce the target’s minimum PTX
        ↓
use the same decision in lowering and code generation
        ↓
reject or compile at a coherent boundary

A version selector cannot describe only the producer side of an intermediate representation.

It must also account for the consumer architecture that the representation is expected to name.

That is the boundary repaired by OpenXLA PR #48532.