Skip to content

Wrap scalar int as NineToothedTensor in launcher#152

Merged
voltjia merged 1 commit into
masterfrom
wrap-int-as-tensor-in-launcher
May 9, 2026
Merged

Wrap scalar int as NineToothedTensor in launcher#152
voltjia merged 1 commit into
masterfrom
wrap-int-as-tensor-in-launcher

Conversation

@voltjia
Copy link
Copy Markdown
Collaborator

@voltjia voltjia commented May 9, 2026

Summary

  • Treat positional int arguments at tensor-parameter positions as scalar tensors in the AOT launcher, mirroring the existing float handling (wrapped via _ArgumentTensor.from_scalar with ctypes.c_int64).
  • Parse the generated <kernel>.h header to count NineToothedTensor parameters, so the scalar-to-tensor wrapping only applies to true tensor-arg slots; trailing arguments (e.g. dtype index strings) keep their prior conversion path.

Testing

pytest output:

============================= test session starts ==============================
platform linux -- Python 3.10.16, pytest-9.0.2, pluggy-1.6.0
rootdir: /home/huangjiacheng/ninetoothed
configfile: pyproject.toml
plugins: anyio-4.12.1, xdist-3.8.0, cov-7.0.0, typeguard-4.4.4
collected 213 items

tests/test_add.py .                                                      [  0%]
tests/test_addmm.py ..                                                   [  1%]
tests/test_aot.py .........                                              [  5%]
tests/test_aot_auto_tuning.py ....                                       [  7%]
tests/test_attention.py ........                                         [ 11%]
tests/test_auto_tuner.py ....                                            [ 13%]
tests/test_clone.py ....                                                 [ 15%]
tests/test_conv2d.py ....                                                [ 16%]
tests/test_data_ptr.py .                                                 [ 17%]
tests/test_debugging.py .                                                [ 17%]
tests/test_dropout.py .                                                  [ 18%]
tests/test_eval.py ........                                              [ 22%]
tests/test_expand.py .                                                   [ 22%]
tests/test_generation.py ............................................... [ 44%]
.............................                                            [ 58%]
tests/test_getitem.py ..........                                         [ 62%]
tests/test_ipynb.py .                                                    [ 63%]
tests/test_jagged.py ................                                    [ 70%]
tests/test_matmul.py ..                                                  [ 71%]
tests/test_max_pool2d.py ..                                              [ 72%]
tests/test_naming.py .......                                             [ 76%]
tests/test_pad.py ................................................       [ 98%]
tests/test_pow.py .                                                      [ 99%]
tests/test_softmax.py .                                                  [ 99%]
tests/test_unsqueeze.py .                                                [100%]

======================= 213 passed in 2813.42s (0:46:53) =======================

@voltjia voltjia merged commit b1c554b into master May 9, 2026
8 checks passed
@voltjia voltjia deleted the wrap-int-as-tensor-in-launcher branch May 9, 2026 02:33
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