Skip to content

mir-importer: read closure upvars from generic suffix - #212

Merged
nihalpasham merged 1 commit into
NVIDIA:mainfrom
niklebedenko:upstream-closure-upvars
Jun 16, 2026
Merged

nihalpasham merged 1 commit into
NVIDIA:mainfrom
niklebedenko:upstream-closure-upvars

Conversation

@niklebedenko

Copy link
Copy Markdown
Contributor

Closes #211.

Problem

rustc suffix-encodes closure substitutions as [parent_args..., closure_kind, closure_sig, tupled_upvars]. The MIR importer treated substs[2] as the upvars tuple, which is only true when there are no parent generic arguments before the closure suffix.

In a generic kernel that constructs a captured closure and passes it through an out-of-line generic helper, the old importer translated the closure environment as a zero-field struct even though closure construction still had two capture operands. On upstream/main with only the regression applied:

MirConstructStructOp has 2 operands but struct
'generic::kernels::closure_capture::{closure#0}' has 0 fields

Fix

Add a small closure_upvar_tys helper that reads the tupled-upvars type from the last closure generic argument and verifies that it is a tuple. Use that helper in both places that decode closure captures:

  • is_rust_type_zst, so captured closures in generic contexts are not mistaken for zero-sized closures;
  • closure translate_type, so the generated MIR struct fields match rustc's closure environment.

This keeps the compiler rule in one place instead of duplicating a fixed-index assumption.

Diligence

I checked the sibling closure paths for other fixed-index upvar decoding:

  • AggregateKind::Closure builds the closure value from MIR operands and delegates the closure type to types::translate_type.
  • closure function ABI handling uses fn_sig() and does not inspect upvar substitutions.
  • closure call lowering resolves the closure body from the receiver type and substitutions, but does not decode upvars itself.
  • collector closure discovery resolves closure bodies from rustc type information and does not build closure environment fields.

The initial direct-call regression was stale because rustc optimized the closure environment away before import. The committed regression keeps the closure as a real value by passing it through #[inline(never)] apply_unary.

Tests

Extended crates/rustc-codegen-cuda/examples/generic with closure_capture<T>, a generic kernel that captures two generic values in a closure and passes that closure to an out-of-line generic helper:

let transform = move |x: T| (x + bias) * scale;
*out_elem = apply_unary(transform, input[idx_raw]);

The host launches the kernel for both f32 and i32, checking that the captured arithmetic is applied correctly for each monomorphization.

Validation

Verified the regression fails on upstream/main by applying only the example change and running:

CUDA_OXIDE_SHOW_RUSTC_MIR=1 cargo oxide run generic

Verified the patched branch with:

cargo fmt --all -- --check
cargo check -p mir-importer
cargo clippy -p mir-importer -- -D warnings
(cd crates/rustc-codegen-cuda/examples/generic && cargo fmt --all -- --check)
(cd crates/rustc-codegen-cuda/examples/generic && cargo clippy --all-targets -- -D warnings)
(cd crates/rustc-codegen-cuda && cargo fmt --all -- --check)
(cd crates/rustc-codegen-cuda && cargo clippy --all-targets -- -D warnings)
scripts/smoketest.sh --compile-only --no-color --only '^generic$'
cargo oxide run generic
scripts/smoketest.sh --no-color --only '^generic$'

@nihalpasham nihalpasham added bug Something isn't working codegen Device code-generation pipeline (Rust MIR to IR to PTX) labels Jun 16, 2026
@nihalpasham

Copy link
Copy Markdown
Collaborator

Thanks for the PR. This was an interesting catch - suffix vs fixed-index thing.

Merging as-is. Tiny heads up: this one and #206 both add a helper at the same spot in types.rs, so whichever lands second will need a trivial keep-both rebase. 👍🏾

@nihalpasham
nihalpasham merged commit 6dbd4cf into NVIDIA:main Jun 16, 2026
24 checks passed
@niklebedenko
niklebedenko deleted the upstream-closure-upvars branch June 24, 2026 11:40
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

bug Something isn't working codegen Device code-generation pipeline (Rust MIR to IR to PTX)

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Read closure upvars from the closure substitution suffix

2 participants