Repository: triton-lang/triton

Issue #4739: Softmax tutorial crashes (invalid arith.select) when n_cols is a multiple of 16 but <= 128
Created: 2024-09-17T23:01:43Z
triton.__version__ is 3.0.0 for me
The tutorial code 02-fused-softmax.py given in https://triton-lang.org/main/getting-started/tutorials/02-fused-softmax.html fails to compile a kernel during the warmup when n_col is a multiple of 16 that is less than or equal to 128 (i.e., <= 16 * num_warps). Error looks like:
(triton) mantissa@MantissaAmpere:~/junk$ python3 ../Downloads/02-fused-softmax.py 
loc("/home/mantissa/junk/../Downloads/02-fused-softmax.py":97:22): error: 'arith.select' op expected condition type to have the same shape as the result type, expected 'tensor<128xi1, #triton_gpu.blocked<{sizePerThread = [1], threadsPerWarp = [32], warpsPerCTA = [8], order = [0]}>>', but got 'tensor<128xi1, #triton_gpu.blocked<{sizePerThread = [4], threadsPerWarp = [32], warpsPerCTA = [8], order = [0]}>>'
Traceback (most recent call last):
  File "/home/mantissa/junk/../Downloads/02-fused-softmax.py", line 196, in <module>
    y_triton = softmax(x)
  File "/home/mantissa/junk/../Downloads/02-fused-softmax.py", line 144, in softmax
    kernel = softmax_kernel.warmup(y, x, x.stride(0), y.stride(0), n_rows, n_cols, BLOCK_SIZE=BLOCK_SIZE,
  File "/home/mantissa/junk/triton/lib/python3.10/site-packages/triton/runtime/jit.py", line 764, in warmup
    return self.run(grid=grid, warmup=True, *map(MockTensor.wrap_dtype, args), **kwargs)
  File "/home/mantissa/junk/triton/lib/python3.10/site-packages/triton/runtime/jit.py", line 662, in run
    kernel = self.compile(
  File "/home/mantissa/junk/triton/lib/python3.10/site-packages/triton/compiler/compiler.py", line 282, in compile
    next_module = compile_ir(module, metadata)
  File "/home/mantissa/junk/triton/lib/python3.10/site-packages/triton/backends/nvidia/compiler.py", line 317, in <lambda>
    stages["ttgir"] = lambda src, metadata: self.make_ttgir(src, metadata, options, self.capability)
  File "/home/mantissa/junk/triton/lib/python3.10/site-packages/triton/backends/nvidia/compiler.py", line 189, in make_ttgir
    pm.run(mod)
RuntimeError: PassManager::run failed

Repro: modify line 195 in 02-fused-softmax.py from
x = torch.randn(1823, 781, device='cuda')

to
x = torch.randn(1823, 80, device='cuda')

and run.

Duplicate issue #5435: 02-fused-softmax do not suport N<=128
Created: 2024-12-16T07:37:23Z
Describe the bug
this demo do not support N <=128, https://github.com/triton-lang/triton/blob/release/3.1.x/python/tutorials/02-fused-softmax.py#L168
i modify the x = torch.randn(1823, 781, device='cuda') to  x = torch.randn(1823, 128, device='cuda')  or else other value less than 128, it will occur an error
python 02-fused-softmax.py
loc("/root/triton/python/tutorials/02-fused-softmax.py":88:22): error: 'arith.select' op expected condition type to have the same shape as the result type, expected 'tensor<128xi1, #triton_gpu.blocked<{sizePerThread = [1], threadsPerWarp = [32], warpsPerCTA = [8], order = [0]}>>', but got 'tensor<128xi1, #triton_gpu.blocked<{sizePerThread = [4], threadsPerWarp = [32], warpsPerCTA = [8], order = [0]}>>'
Traceback (most recent call last):
  File "/root/triton/python/tutorials/02-fused-softmax.py", line 169, in <module>
    y_triton = softmax(x)
               ^^^^^^^^^^
  File "/root/triton/python/tutorials/02-fused-softmax.py", line 135, in softmax
    kernel = softmax_kernel.warmup(y, x, x.stride(0), y.stride(0), n_rows, n_cols, BLOCK_SIZE=BLOCK_SIZE,
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/root/miniconda3/envs/spconv_cu121/lib/python3.11/site-packages/triton/runtime/jit.py", line 766, in warmup
    return self.run(grid=grid, warmup=True, *map(MockTensor.wrap_dtype, args), **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/root/miniconda3/envs/spconv_cu121/lib/python3.11/site-packages/triton/runtime/jit.py", line 664, in run
    kernel = self.compile(
             ^^^^^^^^^^^^^
  File "/root/miniconda3/envs/spconv_cu121/lib/python3.11/site-packages/triton/compiler/compiler.py", line 282, in compile
    next_module = compile_ir(module, metadata)
                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/root/miniconda3/envs/spconv_cu121/lib/python3.11/site-packages/triton/backends/nvidia/compiler.py", line 317, in <lambda>
    stages["ttgir"] = lambda src, metadata: self.make_ttgir(src, metadata, options, self.capability)
                                            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/root/miniconda3/envs/spconv_cu121/lib/python3.11/site-packages/triton/backends/nvidia/compiler.py", line 189, in make_ttgir
    pm.run(mod)
RuntimeError: PassManager::run failed

Environment details
Triton: <=3.0.0
GPU: rtx3060, rtx3090

Closing comment (2025-01-15T01:20:06Z): @l1351868270 I wasn't able to reproduce the bug on my machine with latest Triton main branch. I have a different NVIDIA GPU,  but if it is the the pipeliner, then the bug has not yet touched the hardware.
Your bug is probably related to #4739 which has been fixed.
@plotfi I think your bug is a different one. I cannot use your easy repro as the IR has evolved. Let me try to make a repro with recent OSS following the uneasy path.
So, can I close this issue as the OP's bug has been fixed in triton/main?