[SamJSui]

Summary

• recursively search spec.TupleParam.params when detecting the TVM-FFI environment stream
• preserve the existing left-to-right selection of the first non-CPU tensor
• add a regression test using tuple-nested tensors and a non-default PyTorch stream

Fixes #3444.

Root cause

TVMFFIFunctionBuilder.find_env_stream() only inspected top-level parameters. When every GPU tensor was contained by a spec.TupleParam, the search returned None and wrapper construction raised EnvStream cannot be detected even though the parameter tree contained GPU tensors.

Validation

Tested on an NVIDIA GeForce RTX 4070 Ti SUPER (SM89) against commit 564d267e4c.

Base

python -m pytest test/examples/CuTeDSL/test_tvm_ffi_env_stream.py -q --runtime-sm 89
1 failed in 0.70s

The test failed during cute.compile() with the expected EnvStream cannot be detected error.

Patched

python -m pytest test/examples/CuTeDSL/test_tvm_ffi_env_stream.py -q --runtime-sm 89
1 passed in 0.42s

Neighboring tests

python -m pytest \
  test/examples/CuTeDSL/test_dataclasses.py \
  test/examples/CuTeDSL/test_for_control_flow.py \
  test/examples/CuTeDSL/test_math.py \
  -q --runtime-sm 89
5 passed in 0.42s

The local editable environment used the preparation-script change from #3417; that unrelated file is not included in this pull request.

[anakinxc]

LGTM

@Junkai-Wu The fix has landed internally, can you help merge this one? Thanks