[BUG]examples/dsa_sparse_finetune/indexer_topk_reducesum.py Run Error

[changdong1687]

Required prerequisites

• I have read the documentation https://tilelang.com.
• I have searched the Issue Tracker that this hasn't already been reported. (comment there if it has.)

What version of TileLang are you using?

0.1.7.post2+cuda.gitd6eb5d3d

System information

sys.version and sys.platform 3.12.0 | packaged by Anaconda, Inc. | (main, Oct 2 2023, 17:29:18) [GCC 11.2.0] linux
tilelang version: 0.1.7.post2+cuda.git2d8d3676
torch version: 2.9.0+cu128

Problem description

when run example below(examples/dsa_sparse_finetune/indexer_topk_reducesum.py), it errors:

ValueError: Cannot use and / or / not operator to Expr, hint: use tvm.tir.all / tvm.tir.any instead

Reproducible example code

The Python snippets:

python examples/dsa_sparse_finetune/indexer_topk_reducesum.py 

Traceback

Expected behavior

No response

Additional context

No response

[tzj-fxz]

I run this script successfully on cu129. Could you please double-check it on this latest commit 3144969.

[changdong1687]

I run this script successfully on cu129. Could you please double-check it on this latest commit 3144969.

Ok, thanks, I'm running on RTX4090 with cu128, but I can't checkout to lastest commit 31449690b9f1767dde0bc25cf6392a1707c69ce7 , it errors

fatal: unable to read tree (31449690b9f1767dde0bc25cf6392a1707c69ce7)

And I clone the repo again, using pip install -e . -v, the tilelang version is tilelang-0.1.7.post2+cuda.git2d8d3676
then I run python examples/dsa_sparse_finetune/indexer_topk_reducesum.py it errors

Traceback (most recent call last):
  File "/home/tilelang/examples/dsa_sparse_finetune/indexer_topk_reducesum.py", line 6, in <module>
    import tilelang as tl
  File "/home/tilelang/tilelang/__init__.py", line 133, in <module>
    import tvm
  File "/home/tilelang/3rdparty/tvm/python/tvm/__init__.py", line 27, in <module>
    from .base import TVMError, __version__, _RUNTIME_ONLY
  File "/home/tilelang/3rdparty/tvm/python/tvm/base.py", line 58, in <module>
    _LIB, _LIB_NAME = _load_lib()
                      ^^^^^^^^^^^
  File "/home/tilelang/3rdparty/tvm/python/tvm/base.py", line 45, in _load_lib
    lib = ctypes.CDLL(lib_path[0], ctypes.RTLD_GLOBAL | os.RTLD_LAZY)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tilelang/tilelang/__init__.py", line 117, in lazy_init
    return old_init(self, name, mode | os.RTLD_LAZY, *args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/surd/miniconda3/envs/tilelang/lib/python3.12/ctypes/__init__.py", line 379, in __init__
    self._handle = _dlopen(self._name, mode)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^
OSError: /home/tilelang/build/lib/libtvm.so: undefined symbol: cuLaunchKernel

I run ldd /home/tilelang/build/lib/libtvm.so | grep cuda, it shows

libnvrtc.so.12 => /home/cuda-12/lib64/libnvrtc.so.12 (0x00007f7f3aa00000)
libcudart.so.12 => /home/cuda-12/lib64/libcudart.so.12 (0x00007f7f3a600000)

[Rachmanino]

I run this script successfully on cu129. Could you please double-check it on this latest commit 3144969.

Ok, thanks, I'm running on RTX4090 with cu128, but I can't checkout to lastest commit 31449690b9f1767dde0bc25cf6392a1707c69ce7 , it errors

fatal: unable to read tree (31449690b9f1767dde0bc25cf6392a1707c69ce7)

And I clone the repo again, using pip install -e . -v, the tilelang version is tilelang-0.1.7.post2+cuda.git2d8d3676 then I run python examples/dsa_sparse_finetune/indexer_topk_reducesum.py it errors

Traceback (most recent call last):
  File "/home/tilelang/examples/dsa_sparse_finetune/indexer_topk_reducesum.py", line 6, in <module>
    import tilelang as tl
  File "/home/tilelang/tilelang/__init__.py", line 133, in <module>
    import tvm
  File "/home/tilelang/3rdparty/tvm/python/tvm/__init__.py", line 27, in <module>
    from .base import TVMError, __version__, _RUNTIME_ONLY
  File "/home/tilelang/3rdparty/tvm/python/tvm/base.py", line 58, in <module>
    _LIB, _LIB_NAME = _load_lib()
                      ^^^^^^^^^^^
  File "/home/tilelang/3rdparty/tvm/python/tvm/base.py", line 45, in _load_lib
    lib = ctypes.CDLL(lib_path[0], ctypes.RTLD_GLOBAL | os.RTLD_LAZY)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tilelang/tilelang/__init__.py", line 117, in lazy_init
    return old_init(self, name, mode | os.RTLD_LAZY, *args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/surd/miniconda3/envs/tilelang/lib/python3.12/ctypes/__init__.py", line 379, in __init__
    self._handle = _dlopen(self._name, mode)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^
OSError: /home/tilelang/build/lib/libtvm.so: undefined symbol: cuLaunchKernel

I run ldd /home/tilelang/build/lib/libtvm.so | grep cuda, it shows

libnvrtc.so.12 => /home/cuda-12/lib64/libnvrtc.so.12 (0x00007f7f3aa00000)
libcudart.so.12 => /home/cuda-12/lib64/libcudart.so.12 (0x00007f7f3a600000)

@changdong1687 No need to checkout to 3144969, which doesn't exist in TileLang repo codebase. Simply re-install the latest tilelang and try again, please. Ensure using the correct LD_LIBRARY_PATH containing cuda libs

[changdong1687]

I run this script successfully on cu129. Could you please double-check it on this latest commit 3144969.

Ok, thanks, I'm running on RTX4090 with cu128, but I can't checkout to lastest commit 31449690b9f1767dde0bc25cf6392a1707c69ce7 , it errors

fatal: unable to read tree (31449690b9f1767dde0bc25cf6392a1707c69ce7)

And I clone the repo again, using pip install -e . -v, the tilelang version is tilelang-0.1.7.post2+cuda.git2d8d3676 then I run python examples/dsa_sparse_finetune/indexer_topk_reducesum.py it errors

Traceback (most recent call last):
  File "/home/tilelang/examples/dsa_sparse_finetune/indexer_topk_reducesum.py", line 6, in <module>
    import tilelang as tl
  File "/home/tilelang/tilelang/__init__.py", line 133, in <module>
    import tvm
  File "/home/tilelang/3rdparty/tvm/python/tvm/__init__.py", line 27, in <module>
    from .base import TVMError, __version__, _RUNTIME_ONLY
  File "/home/tilelang/3rdparty/tvm/python/tvm/base.py", line 58, in <module>
    _LIB, _LIB_NAME = _load_lib()
                      ^^^^^^^^^^^
  File "/home/tilelang/3rdparty/tvm/python/tvm/base.py", line 45, in _load_lib
    lib = ctypes.CDLL(lib_path[0], ctypes.RTLD_GLOBAL | os.RTLD_LAZY)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/tilelang/tilelang/__init__.py", line 117, in lazy_init
    return old_init(self, name, mode | os.RTLD_LAZY, *args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/surd/miniconda3/envs/tilelang/lib/python3.12/ctypes/__init__.py", line 379, in __init__
    self._handle = _dlopen(self._name, mode)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^
OSError: /home/tilelang/build/lib/libtvm.so: undefined symbol: cuLaunchKernel

I run ldd /home/tilelang/build/lib/libtvm.so | grep cuda, it shows

libnvrtc.so.12 => /home/cuda-12/lib64/libnvrtc.so.12 (0x00007f7f3aa00000)
libcudart.so.12 => /home/cuda-12/lib64/libcudart.so.12 (0x00007f7f3a600000)

@changdong1687 No need to checkout to 3144969, which doesn't exist in TileLang repo codebase. Simply re-install the latest tilelang and try again, please. Ensure using the correct LD_LIBRARY_PATH containing cuda libs

ok, after install the newest version it can run, but it shows:

[18:16:48] /home/tilelang/src/transform/verify_parallel_loop.cc:72: Warning: Data race detected: `topk_index_shared[j]`is written by multiple threads in loop [i], Example:
  Offsets[i_b + 1] = 0
  Offsets[i_b + 1]$1 = 0
  Offsets[i_b] = 0
  Offsets[i_b]$1 = 0
  TokenIndices[bx, 0] = 0
  TokenIndices[bx, 0]$1 = 0
  TokenIndices[bx, 1] = 0
  TokenIndices[bx, 1]$1 = 0
  bk_i = 0
  bx = 0
  i = 1
  i_OTHER_ = 0
  offset[0] = 0
  offset[0]$1 = 1
  seq_len = 1
  tx = 0
If you believe this is a false positive, pass `PassKey.TL_DISABLE_DATA_RACE_CHECK` to pass key to disable this check.
[18:16:48] /home/tilelang/src/transform/verify_parallel_loop.cc:72: Warning: Data race detected: `topk_value_shared[j]`is written by multiple threads in loop [i], Example:
  Offsets[i_b + 1] = 0
  Offsets[i_b + 1]$1 = 0
  Offsets[i_b] = 0
  Offsets[i_b]$1 = 0
  TokenIndices[bx, 0] = 0
  TokenIndices[bx, 0]$1 = 0
  TokenIndices[bx, 1] = 0
  TokenIndices[bx, 1]$1 = 0
  bk_i = 0
  bx = 0
  i = 0
  i_OTHER_ = 1
  offset[0] = 0
  offset[0]$1 = (- 1)
  seq_len = 1
  tx = 0
If you believe this is a false positive, pass `PassKey.TL_DISABLE_DATA_RACE_CHECK` to pass key to disable this check.

and for some rows, it only choose 63 indices, and topk is 64

If I need using the 0.1.7.post2+cuda.git2d8d3676, how can I deal this bug, because my other kernels is developed in this version, and the newset tilelang will cause many errors that the older version not cause

[ShaobinChen-AH]

Tested on the latest main branch （8d67d67c151e0a2722da0224121a90aec7e1c2cf）, the example runs successfully.

The original issues have been resolved:

1. and/or operator errors - no longer occurs
2. Python 3.9 compatibility (| None and strict= syntax) - already fixed in latest code

This issue can be closed.