error when call tl.sort for data of torch.bfloat16

[horrorChen]

When I call tl.sort with input of dtype bf16 I get the error.

 'llvm.fcmp' op operand #0 must be floating point LLVM type or LLVM dialect-compatible vector of floating point LLVM type, but got 'i16'

I use triton 3.0.0. To reproduce

import torch
import triton
import triton.language as tl

@triton.jit
def test_kernel(
    input,
    output,
    row_size: tl.constexpr,
    BLOCK_SIZE: tl.constexpr
):
    pid = tl.program_id(0)
    off = tl.arange(0, BLOCK_SIZE)
    row = tl.load(input + pid * row_size + off, mask=off<row_size, other=0)
    row = tl.sort(row)
    # row = row * row
    tl.store(output + pid * row_size + off, row, mask=off<row_size)

if __name__ == "__main__":
    a = torch.rand(8, 16, device='cuda', dtype=torch.bfloat16)
    b = torch.empty(8, 16, device='cuda', dtype=torch.bfloat16)
    test_kernel[(8, )](a, b, 16, 16)

It is confused that fp16 and fp32 work.

FYI, I think the error code is located in function _compare_and_swap of standard.py.

[horrorChen]

In another env I get a different error:

Traceback (most recent call last):
  File "/usr/local/lib/python3.10/dist-packages/triton/language/core.py", line 33, in wrapper
    return fn(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/triton/language/core.py", line 912, in to
    return semantic.bitcast(self, dtype, _builder)
  File "/usr/local/lib/python3.10/dist-packages/triton/language/semantic.py", line 737, in bitcast
    raise ValueError("Cannot bitcast data-type of size " + str(src_bits) + " to "
ValueError: Cannot bitcast data-type of size 32 to data-type of size 16

The above exception was the direct cause of the following exception:

triton.compiler.errors.CompilationError: at 13:12:
    n_outer: core.constexpr = x.numel >> n_dims
    shape: core.constexpr = [n_outer * 2**i, 2, 2**(n_dims - i - 1)]
    y = core.reshape(x, shape)
    # slice left/right with 'stride' 2**(n_dims - i - 1)
    mask = core.arange(0, 2)[None, :, None]
    left = core.broadcast_to(sum(y * (1 - mask), 1)[:, None, :], shape)
    right = core.broadcast_to(sum(y * mask, 1)[:, None, :], shape)
    left = core.reshape(left, x.shape)
    right = core.reshape(right, x.shape)
    # actual compare-and-swap
    idtype = core.dtype(f'int{core.constexpr(x.dtype.primitive_bitwidth)}')
    ileft = left.to(idtype, bitcast=True)

I think this error comes from below code in triton/language/standard.py where mask is in default type of int32, thus left and right turn out to be fp32 rather than bf16 like y.

    mask = core.arange(0, 2)[None, :, None]
    left = core.broadcast_to(sum(y * (1 - mask), 1)[:, None, :], shape)
    right = core.broadcast_to(sum(y * mask, 1)[:, None, :], shape)

So I try to cast left and right to dtype of y like this:

    mask = core.arange(0, 2)[None, :, None]
    left = core.broadcast_to(sum(y * (1 - mask), 1)[:, None, :], shape).to(y.dtype)
    right = core.broadcast_to(sum(y * mask, 1)[:, None, :], shape).to(y.dtype)

Then I get the same error as the beginning:

loc(callsite(callsite(callsite("/usr/local/lib/python3.10/dist-packages/triton/language/standard.py":338:34 at "/usr/local/lib/python3.10/dist-packages/triton/language/standard.py":363:61) at "/usr/local/lib/python3.10/dist-packages/triton/language/standard.py":376:66) at "/workspace2/triton-test/test.py":18:18)): error: 'llvm.fcmp' op operand #0 must be floating point LLVM type or LLVM dialect-compatible vector of floating point LLVM type, but got 'i16'
Traceback (most recent call last):
  File "/workspace2/triton-test/test.py", line 49, in <module>
    test_kernel[(8, )](a, b, 16, 16)
  File "/usr/local/lib/python3.10/dist-packages/triton/runtime/jit.py", line 180, in <lambda>
    return lambda *args, **kwargs: self.run(grid=grid, warmup=False, *args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/triton/runtime/jit.py", line 401, in run
    self.cache[device][key] = compile(
  File "/usr/local/lib/python3.10/dist-packages/triton/compiler/compiler.py", line 268, in compile
    next_module = compile_ir(module, metadata)
  File "/usr/local/lib/python3.10/dist-packages/triton/backends/nvidia/compiler.py", line 265, in <lambda>
    stages["llir"] = lambda src, metadata: self.make_llir(src, metadata, options, self.capability)
  File "/usr/local/lib/python3.10/dist-packages/triton/backends/nvidia/compiler.py", line 178, in make_llir
    pm.run(mod)
RuntimeError: PassManager::run failed

I think the error comes from line 338 where llvm.fcmp is not expected to be generated:

    ret = ix ^ core.where((left > right) ^ flip, ileft ^ iright, zeros_like(ix))

FYI, thx.

[horrorChen]

Fixed with #3975.