Support FP8 constant

[htyu]

To unblock the compilation of kernels like below which don't operate arithmetically on FP8.

@triton.jit
def triton_poi_fused__scaled_mm__to_copy_constant_pad_nd_lift_fresh_2(in_ptr0, out_ptr0, xnumel, XBLOCK : tl.constexpr):
    xnumel = 400624
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x0 = xindex % 784
    x1 = (xindex // 784)
    x2 = xindex
    tmp0 = x0
    tmp1 = tl.full([1], 769, tl.int64)
    tmp2 = tmp0 < tmp1
    tmp3 = tl.load(in_ptr0 + (x0 + (769*x1)), tmp2 & xmask, other=0.0)
    tmp4 = tmp3.to(tl.float8e4nv)
    tmp5 = tl.full(tmp4.shape, 0.0, tmp4.dtype)
    tmp6 = tl.where(tmp2, tmp4, tmp5)
    tl.store(out_ptr0 + (x2), tmp6, xmask)

[htyu]

I'll add a lit test if this sounds a reasonable fix.

[htyu]

I was hitting

  %2 = "llvm.mlir.constant"() <{value = 0.000000e+00 : f8E4M3FNUZ}> : () -> f8E4M3FNUZ loc(#loc1)

error: 'llvm.mlir.constant' op result #0 must be LLVM dialect-compatible type, but got 'f8E4M3FNUZ'

[ThomasRaoux]

I was hitting

  %2 = "llvm.mlir.constant"() <{value = 0.000000e+00 : f8E4M3FNUZ}> : () -> f8E4M3FNUZ loc(#loc1)

error: 'llvm.mlir.constant' op result #0 must be LLVM dialect-compatible type, but got 'f8E4M3FNUZ'

you're right, makes sense. Does the conversion to LLVM dialect work for scalar constant?

[htyu]

I was hitting

  %2 = "llvm.mlir.constant"() <{value = 0.000000e+00 : f8E4M3FNUZ}> : () -> f8E4M3FNUZ loc(#loc1)

error: 'llvm.mlir.constant' op result #0 must be LLVM dialect-compatible type, but got 'f8E4M3FNUZ'

you're right, makes sense. Does the conversion to LLVM dialect work for scalar constant?

Good point. Probably not. Actually the conversion being changed is for scalar constant, which is then broadcasted to a tensor in registers.

[htyu]

Oh I see your point. The original constant on TTGIR is a tensor
%cst = arith.constant dense<0.000000e+00> : tensor<1024xf8E4M3FNUZ, #blocked> loc(#loc1)

I'll add support for scalar constant.

[ThomasRaoux]

I was hitting

  %2 = "llvm.mlir.constant"() <{value = 0.000000e+00 : f8E4M3FNUZ}> : () -> f8E4M3FNUZ loc(#loc1)

error: 'llvm.mlir.constant' op result #0 must be LLVM dialect-compatible type, but got 'f8E4M3FNUZ'

you're right, makes sense. Does the conversion to LLVM dialect work for scalar constant?

Good point. Probably not. Actually the conversion being changed is for scalar constant, which is then broadcasted to a tensor in registers.

I believe we rely on upstream pattern for scalar arith.constant. It would be worth checking if it works and maybe we need to upstream a fix

[htyu]

Confirmed that scalar const lowering works with LLVM:

`%cst = arith.constant 0.000000e+00 : f8E4M3FNUZ`

=>
%0 = llvm.mlir.constant(0.000000e+00 : f8E4M3FNUZ) : i8

[ThomasRaoux]

so fp8 would not have auto-promote but other types would? Seems a bit odd. I would add support for fp8 in binary_op_type_checking_impl instead

[htyu]

The current logics does not support fp8 because fp8 arithmetic is not available on hardware (except for dot). Do I understand it correctly?

Also, I'm not sure we need promote here, as in the example kernel fp8 is not used for arithmetic, rather, the kernel loads fp32, convert it to fp8, and conditionally stored out.

[htyu]

Or do you think it's safe to not promote anything when x.type == y.type?

[ThomasRaoux]

I meant in general if we have a mix-mode op we promote to the highest format. We could do that for fp8 as well right?

[htyu]

Oh yeah we should do that.

[htyu]

Bypass arithmetic type check when inputs are same-typed.