Skip to content

feat(cutile): support scalar specialization overrides - #227

Open
lucifer1004 wants to merge 1 commit into
NVlabs:mainfrom
lucifer1004:pr/scalar-specialization-overrides
Open

lucifer1004 wants to merge 1 commit into
NVlabs:mainfrom
lucifer1004:pr/scalar-specialization-overrides

Conversation

@lucifer1004

Copy link
Copy Markdown
Contributor

Motivation

The launcher infers divisibility hints (DivHint) for integer scalar and raw-pointer kernel parameters from their runtime values, and those hints feed the JIT specialization key. For truly dynamic scalars this backfires: a value that happens to be divisible by 4/8/16 on some launches (e.g. a decode batch size m that varies per step) silently produces multiple JIT specializations for the same kernel, each triggering a recompile.

What this adds

A scalar_hint(name, hint) builder method on the generated kernel launcher (declared on the TileKernel trait) that replaces the inferred hint for a named parameter — or adds one — before hints are consumed:

  • .scalar_hint("m", DivHint::default()) (divisor 1) disables specialization for a genuinely dynamic scalar, so incidental runtime divisibility no longer creates extra JIT variants.
  • .scalar_hint("n", hint) can also pin a stronger hint than the inferred one when the caller knows a contract the runtime value doesn't reveal.

The override drain runs on both the launch path and the specialization path (the shared statement is pushed into launcher_method and specialization_method), so the two stay consistent.

Example

kernel(ptr, m)
    .grid(grid)
    .scalar_hint("m", DivHint::default()) // m is truly dynamic; do not specialize on it
    .sync()?;

Tests

Added two GPU launch tests in cutile/tests/specialization_bits.rs that dump the generated MLIR and assert:

  • an override replaces the inferred hint (n=12 infers div_by<4>; overriding with div_by<8> emits assume div_by<8> and no div_by<4>);
  • a default-hint override removes the scalar assume entirely while the raw-pointer assume is untouched.

All 4 raw_pointer_launch* tests in that file pass locally on an SM120 GPU.

@copy-pr-bot

copy-pr-bot Bot commented Aug 20, 2026

Copy link
Copy Markdown

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

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.

1 participant