Skip to content

feat(cutile): add raw-offset memory ops with generic forwarding support - #229

Open
lucifer1004 wants to merge 2 commits into
NVlabs:mainfrom
lucifer1004:pr/raw-offset-memory-ops
Open

feat(cutile): add raw-offset memory ops with generic forwarding support#229
lucifer1004 wants to merge 2 commits into
NVlabs:mainfrom
lucifer1004:pr/raw-offset-memory-ops

Conversation

@lucifer1004

Copy link
Copy Markdown
Contributor

Motivation

Packed layouts (e.g. sparse/quantized weight formats) keep values and metadata at independent byte offsets from a raw base pointer, outside any Tensor view. Today expressing that access pattern in the DSL requires manually composing pointer_to_tile / reshape_ptr / broadcast_ptr / addptr_tile at every call site.

What this adds

Zero-cost raw-offset memory helpers in cutile::core, built entirely on existing pointer-tile primitives:

  • load_offset(base, offsets, ...) — load base + offsets (offsets in elements of B)
  • load_offset_as(base, offsets, ...) — same addressing, then reinterpret each address as E; a *mut u8 base gives byte-addressed access to packed values with independently placed metadata
  • store_offset(base, offsets, value, ...) — store to base + offsets

All three are variadic_op wrappers forwarding ordering / scope / mask / padding / token / latency through to load_ptr_tko / store_ptr_tko, so they carry no lowering cost of their own (verified in the test: no make_gather_scatter_view, plain offset + load_ptr_tko/store_ptr_tko in the MLIR).

Compiler support required

Because the wrappers receive marker arguments (ordering, scope, Latency<N>) by value and forward them, the compiler previously could not resolve them — primitives like load_ptr_tko only accepted literal type paths at the call site. This PR teaches the compiler to:

  • resolve forwarded ZST type names through local variables and generic instantiations (extract_forwarded_zst_type_name);
  • resolve Latency<N> cycle counts from a variable's type, including const-generic values;
  • compile inlined Some payloads and explicitly typed None::<T> arguments;
  • treat Option<T> as resolvable/fully-known when its payload is, and infer tuple types element-wise.

It also moves the Latency marker into core proper (previously a top-level struct re-exported into core) so the DSL name resolver indexes it for single-segment paths like Latency::<4> — a re-export alone is invisible to the resolver.

Tests

  • cutile/tests/raw_memory.rs (new): compiles kernels using the helpers and asserts the lowered MLIR (offset arithmetic, ordering/scope/latency forwarding, token threading, no gather/scatter view materialization).
  • cutile/tests/type_conversion_ops.rs: F32→E8M0 conversion coverage with positive-infinity rounding (compile-time MLIR check + GPU execution roundtrip).

Full cutile + cutile-compiler test suites pass locally on an SM120 GPU. (Note: the pre-existing cuda_tile_runtime_utils::tests::selects_bytecode_version_from_toolkit_cuda_h is sensitive to the CUTILE_BYTECODE_VERSION env var and fails in shells where it is set — unrelated to this change.)

Add zero-cost load_offset/load_offset_as/store_offset helpers that build
pointer tiles from a raw base pointer and an integer offset tile, for
packed layouts whose values and metadata live outside any tensor view.

The wrappers forward ordering/scope/latency marker arguments by value,
which requires the compiler to resolve forwarded ZST type names,
variable-held Latency const generics, inlined Option payloads, explicit
None::<T> annotations, and tuple element-wise generic inference. Move
the Latency marker into the core module proper so the DSL name resolver
indexes it for single-segment paths.

Also cover F32-to-E8M0 conversion with positive-infinity rounding.
@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.

@elibol elibol modified the milestone: cuTile Rust v0.3.1 Sep 1, 2026
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