minor PTX cleanups [pr] (#8351)

This commit is contained in:
chenyu
2024-12-20 12:52:53 -05:00
committed by GitHub
parent 59f4b8da95
commit 20eebbc61a
+2 -4
View File
@@ -32,7 +32,7 @@ asm_for_op: dict[Ops, Callable] = {
f"selp.{'b16' if name == 'f16' else name} {d}, {b}, {c}, {a};"
}
supports_half: list[Ops] = [Ops.EXP2, Ops.ADD, Ops.MUL, Ops.MAX, Ops.CMPLT, Ops.WHERE]
supports_half = (Ops.EXP2, Ops.ADD, Ops.MUL, Ops.MAX, Ops.CMPLT, Ops.WHERE)
doesnt_support_half: tuple[Ops, ...] = tuple(op for op in asm_for_op.keys() if op not in supports_half)
ptx_matcher = PatternMatcher([
# bool CMPNE is XOR, bool CMPLT is XOR+AND (universal makes this slow, this is for renderer only)
@@ -137,14 +137,12 @@ class PTXRenderer(Renderer):
.address_size 64
.visible .entry"""
barrier = "bar.sync\t0;"
supports_half = supports_half
# HACK: Use s16 and u16 for int8 and uint8 buffers. This can be wrong in cast.
types: dict[DType, str] = { dtypes.int8: "s16", dtypes.int16: "s16", dtypes.int32: "s32", dtypes.int64: "s64",
dtypes.uint8: "u16", dtypes.uint16: "u16", dtypes.uint32: "u32", dtypes.uint64: "u64",
dtypes.float16: "f16", dtypes.float32: "f32", dtypes.float64: "f64", dtypes.bool: "pred" }
mem_types: dict[DType, str] = types.copy()
mem_types.update({dtypes.int8: "s8", dtypes.uint8: "u8", dtypes.bool: "u8", dtypes.float16: "b16"})
mem_types = {**types, dtypes.int8: "s8", dtypes.uint8: "u8", dtypes.bool: "u8", dtypes.float16: "b16"}
def render_kernel(self, kernel, function_name, bufs, regs) -> str:
def fmt(line): return line if line[0]=="$" else "\t" + line.replace(" ", "\t" if len(line.split(" ")[0]) > 7 else "\t\t", 1)