[BUG] CuTe DSL >=4.6.0: export_to_c does not apply explicit preferred_smem_carveout #3715

Which component has the problem?

CuTe DSL

Bug Report

Description

During AOT compilation an explicit preferred_smem_carveout survives compilation and is stored in the exported object's metadata, but the generated C loader does not apply it. This reproducer exports the same kernel with preferences of 25% and 75%. Both exports set attributes 8 and 14 (via _cudaFuncSetAttribute(...)), but omit attribute 9 (which is PREFERRED_SHARED_MEMORY_CARVEOUT).

This feature was introduced in 4.6.0, link here

Environment

• nvidia-cutlass-dsl and its components: 4.8.0
• cuda-bindings: 13.3.1
• Python 3.10.12, Linux x86-64
• GPU compilation target: sm_90a
• Host object target: x86_64-unknown-linux-gnu, x86-64 CPU
• GCC 11.4.0, GNU binutils 2.38

Reproduce

Point: compile AOT some kernel with some explicit values for preferred_smem_carveout, then compare disassembled object files and see no _cudaFuncSetAttribute(...) with values for attribute 9 (PREFERRED_SHARED_MEMORY_CARVEOUT) but identical loaders.

Create some files with following contents.

requirements.txt:

nvidia-cutlass-dsl==4.8.0
cuda-bindings==13.3.1

repro.py:

import json
from pathlib import Path

import cuda.bindings.driver as cuda
import cutlass
import cutlass.cute as cute
from cutlass.cute.runtime import make_fake_stream, make_fake_tensor
from cutlass.memory import SmemAllocator


@cute.kernel
def rotate(x: cute.Tensor):
    tid = cute.arch.thread_idx()[0]
    shared = SmemAllocator().allocate_tensor(cutlass.Float32, cute.make_layout(128))
    shared[tid] = x[tid]
    cute.arch.sync_threads()
    x[tid] = shared[(tid + 1) % 128]


@cute.jit
def launch(x: cute.Tensor, stream: cuda.CUstream, carveout: cutlass.Constexpr):
    rotate(x).launch(
        grid=(1, 1, 1),
        block=(128, 1, 1),
        stream=stream,
        min_blocks_per_mp=2,
        preferred_smem_carveout=carveout,
    )


out = Path(__file__).parent / "artifacts"
out.mkdir(exist_ok=True)
x = make_fake_tensor(cutlass.Float32, (128,), (1,), assumed_align=16)

for carveout in (25, 75):
    name = f"carveout_{carveout}"
    compiled = cute.compile(
        launch, x, make_fake_stream(), carveout,
        options="--gpu-arch sm_90a --host-target 'llvm -mtriple=x86_64-unknown-linux-gnu -mcpu=x86-64'",
    )
    compiled.export_to_c(str(out), name, name)
    (out / f"{name}.kernel_info.json").write_text(json.dumps(compiled.kernel_info, indent=2) + "\n")
    print(f"{name}: {compiled.kernel_info}")

inspect_artifacts.py:

import json
import re
import subprocess
from pathlib import Path


out = Path(__file__).parent / "artifacts"

for case in (25, 75):
    name = f"carveout_{case}"
    obj = out / f"{name}.o"
    strings = subprocess.check_output(["strings", obj.name], cwd=out, text=True)
    metadata = next(line for line in strings.splitlines() if line.startswith('{"kernel_'))
    (out / f"{name}.embedded_metadata.json").write_text(json.dumps(json.loads(metadata), indent=2) + "\n")
    assembly = subprocess.check_output(["objdump", "-dr", "--no-show-raw-insn", obj.name], cwd=out, text=True)
    (out / f"{name}.disassembly.txt").write_text(assembly)
    functions = re.split(r"(?=^[0-9a-f]+ <)", assembly, flags=re.MULTILINE)
    loaders = [function for function in functions if re.match(rf"[0-9a-f]+ <{name}_cuda_load(?:_to_device)?>:", function)]
    (out / f"{name}.loaders.txt").write_text("".join(loaders))
    print(f"{name}: embedded metadata = {metadata}")

first = (out / "carveout_25.loaders.txt").read_text().replace("carveout_25", "carveout")
second = (out / "carveout_75.loaders.txt").read_text().replace("carveout_75", "carveout")
print(f"Loader disassembly identical after normalizing the export prefix: {first == second}")

record_loader.c

#include <stdio.h>

int _cudaGetDevice(int *device) {
    *device = 0;
    return 0;
}

int _cudaDeviceGetAttribute(int *value, int attribute, int device) {
    *value = 232448;
    return 0;
}

int _cuKernelGetAttribute(int *value, int attribute, void *kernel, int device) {
    *value = 0;
    return 0;
}

int _cudaLibraryGetKernel(void **kernel, void *library, const char *name) {
    *kernel = (void *)2;
    return 0;
}

int _cudaFuncSetAttribute(void *kernel, int attribute, int value) {
    printf("  _cudaFuncSetAttribute: attribute=%d, value=%d\n", attribute, value);
    return 0;
}

int _cudaKernelSetAttributeForDevice(void *kernel, int attribute, int value, int device) {
    printf("  _cudaKernelSetAttributeForDevice: attribute=%d, value=%d, device=%d\n", attribute, value, device);
    return 0;
}

int _cudaLibraryLoadData(void *library, ...) {
    return 0;
}

int _cudaLaunchKernelEx(void *config, void *kernel, void *arguments) {
    return 0;
}

int carveout_25_cuda_load(void **library);
int carveout_25_cuda_load_to_device(void **library, int device);
int carveout_75_cuda_load(void **library);
int carveout_75_cuda_load_to_device(void **library, int device);

int main(void) {
    void *library = (void *)1;
    puts("preferred_smem_carveout=25, cuda_load:");
    carveout_25_cuda_load(&library);
    puts("preferred_smem_carveout=25, cuda_load_to_device:");
    carveout_25_cuda_load_to_device(&library, 0);
    puts("preferred_smem_carveout=75, cuda_load:");
    carveout_75_cuda_load(&library);
    puts("preferred_smem_carveout=75, cuda_load_to_device:");
    carveout_75_cuda_load_to_device(&library, 0);
    return 0;
}

Then run in some python venv:

pip install -r requirements.txt
mkdir -p artifacts
CUDA_VISIBLE_DEVICES='' CUTE_DSL_ARCH=sm_90a CUTE_DSL_NO_CACHE=1 \
  CUTE_DSL_KEEP=ir,ptx CUTE_DSL_DUMP_DIR=artifacts \
  python -u repro.py > artifacts/compile.log 2>&1
python inspect_artifacts.py > artifacts/inspection.log
cc record_loader.c artifacts/carveout_25.o artifacts/carveout_75.o -o artifacts/record_loader
./artifacts/record_loader > artifacts/loader_calls.txt
cat artifacts/inspection.log artifacts/loader_calls.txt

repro.py is the complete kernel and AOT compilation example. The kernel rotates 128 floats through shared memory. Its launch uses:

Expected

Loading each export through its generated C header should apply the requested preference. For the loader, this means an attribute setting call equivalent to:

carveout_25: _cudaKernelSetAttributeForDevice(kernel, 9, 25, device)
carveout_75: _cudaKernelSetAttributeForDevice(kernel, 9, 75, device)

The non-device-specific loader should likewise pass attribute 9 and the requested value to _cudaFuncSetAttribute. These names are the runtime wrapper symbols referenced by the emitted object.

This expectation concerns forwarding a preference to CUDA.

Actual

The metadata extracted directly from the exported objects contains the correct values:

{"kernel_cutlass_rotate_tensorptrf32gmemalign16o1281_0": {"9": 25}}
{"kernel_cutlass_rotate_tensorptrf32gmemalign16o1281_0": {"9": 75}}

However, artifacts/loader_calls.txt records:

preferred_smem_carveout=25, cuda_load:
  _cudaFuncSetAttribute: attribute=8, value=232448
  _cudaFuncSetAttribute: attribute=14, value=1
preferred_smem_carveout=25, cuda_load_to_device:
  _cudaKernelSetAttributeForDevice: attribute=8, value=232448, device=0
  _cudaKernelSetAttributeForDevice: attribute=14, value=1, device=0
preferred_smem_carveout=75, cuda_load:
  _cudaFuncSetAttribute: attribute=8, value=232448
  _cudaFuncSetAttribute: attribute=14, value=1
preferred_smem_carveout=75, cuda_load_to_device:
  _cudaKernelSetAttributeForDevice: attribute=8, value=232448, device=0
  _cudaKernelSetAttributeForDevice: attribute=14, value=1, device=0

The two exports' loader disassemblies are identical after normalizing their symbol prefixes. The generated headers call _mlir_*_cuda_load_to_device. They do not read the carveout metadata or apply it themselves.

How the evidence is collected

record_loader.c links the unmodified exported objects against small recording stubs instead of CUDA. It executes both native loader functions on the CPU. All stubs return success. The setter stubs print their actual arguments. There is no early error return hiding an attribute setting call.

The handles and device-query results are synthetic. In particular, 232448 is supplied by the stub, not measured on a GPU. The recorder does not execute the kernel, query a real runtime preference, or measure performance.

inspect_artifacts.py extracts the metadata and the native code with objdump. They show the same missing setter without relying on the recording stubs. The initialization function only loads the library image. It contains no missing carveout setting.