from triton.language import core
@core.extern
def globaltimer(_builder=None):
return core.inline_asm_elementwise("mov.u64 $0, %globaltimer;", "=l", [], dtype=core.int64, is_pure=False, pack=1,
_builder=_builder)
@core.extern
def smid(_builder=None):
return core.inline_asm_elementwise("mov.u32 $0, %smid;", "=r", [], dtype=core.int32, is_pure=True, pack=1,
_builder=_builder)
@core.builtin
def num_threads(_builder=None):
return core.constexpr(_builder.options.num_warps * 32)
@core.builtin
def num_warps(_builder=None):
return core.constexpr(_builder.options.num_warps)
# ----- FP8E4M3B15 ------
# This data-type is a variant of the standard FP8E4M3 format.
# It was designed for fast software conversion to FP16 on
# nvidia GPUs that do not support it natively.
# This is the same format as FP8E4M3Nv, but:
# - the exponent bias is 15 instead of 7
# - 0xff and 0x7f are mapped to +-1.750 instead of +-nan
@core.builtin
def convert_fp8e4b15_to_float16(arg, _builder=None):
return core.inline_asm_elementwise(
"{ \n"
".reg .b32 a<2>, b<2>; \n"
"prmt.b32 a0, 0, $2, 0x5746; \n"
"and.b32 b0, a0, 0x7f007f00; \n"
"and.b32 b1, a0, 0x00ff00ff; \n"
"and.b32 a1, a0, 0x00800080; \n"
"shr.b32 b0, b0, 1; \n"
"add.u32 b1, b1, a1; \n"
"lop3.b32 $0, b0, 0x80008000, a0, 0xf8; \n"
"shl.b32 $1, b1, 7; \n"
"} \n", "=r,=r,r", [arg], dtype=core.float16, is_pure=True, pack=4,
_builder=_builder)
@core.builtin
def convert_float16_to_fp8e4b15(arg, has_minx2, _builder=None):
asm = """{
.reg .pred p<4>;
.reg .b32 a<2>, b<2>;
.reg .b16 c<4>;
.reg .b16 max_val_f16;
.reg .b32 max_val_f16x2;
mov.b16 max_val_f16, 0x3F00;
mov.b32 max_val_f16x2, 0x3F003F00;
and.b32 a0, $1, 0x7fff7fff;
and.b32 a1, $2, 0x7fff7fff;"""
if has_minx2:
asm += """min.f16x2 a0, a0, max_val_f16x2;
min.f16x2 a1, a1, max_val_f16x2;"""
else:
asm += """setp.lt.f16x2 p0|p1, a0, max_val_f16x2;
setp.lt.f16x2 p2|p3, a1, max_val_f16x2;
mov.b32 {c0, c1}, a0;
mov.b32 {c2, c3}, a1;
selp.b16 c0, c0, max_val_f16, p0;
selp.b16 c1, c1, max_val_f16, p1;
selp.b16 c2, c2, max_val_f16, p2;
selp.b16 c3, c3, max_val_f16, p3;
mov.b32 a0, {c0, c1};
mov.b32 a1, {c2, c3};"""
asm += """mad.lo.u32 a0, a0, 2, 0x00800080;
mad.lo.u32 a1, a1, 2, 0x00800080;
lop3.b32 b0, $1, 0x80008000, a0, 0xea;
lop3.b32 b1, $2, 0x80008000, a1, 0xea;
prmt.b32 $0, b0, b1, 0x7531;
}"""
return core.inline_asm_elementwise(asm, "=r,r,r", [arg], dtype=core.float8e4b15, is_pure=True, pack=4,
_builder=_builder)
@core.builtin
def convert_custom_float8(arg, dst_ty, fp_downcast_rounding, has_minx2, _builder=None):
if arg.type.scalar.is_fp8e4b15():
upcast_val = convert_fp8e4b15_to_float16(arg, _builder=_builder)
if dst_ty.scalar.is_fp32():
upcast_val = upcast_val.to(core.float32, _builder=_builder)
return upcast_val
assert arg.type.scalar.is_fp16() or arg.type.scalar.is_fp32()
downcast_val = arg
if arg.type.scalar.is_fp32():
downcast_val = downcast_val.to(core.float16, fp_downcast_rounding="rtz", _builder=_builder)
downcast_val = convert_float16_to_fp8e4b15(downcast_val, has_minx2=has_minx2, _builder=_builder)
return downcast_val
@core.builtin
def convert_custom_float8_sm80(arg, dst_ty, fp_downcast_rounding=None, _builder=None):
return convert_custom_float8(arg, dst_ty, fp_downcast_rounding, has_minx2=True, _builder=_builder)
@core.builtin
def convert_custom_float8_sm70(arg, dst_ty, fp_downcast_rounding=None, _builder=None):
return convert_custom_float8(arg, dst_ty, fp_downcast_rounding, has_minx2=False, _builder=_builder)