From 20eebbc61a323f389a15e302bfba206db3549ab0 Mon Sep 17 00:00:00 2001 From: chenyu Date: Fri, 20 Dec 2024 12:52:53 -0500 Subject: [PATCH] minor PTX cleanups [pr] (#8351) --- tinygrad/renderer/ptx.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/tinygrad/renderer/ptx.py b/tinygrad/renderer/ptx.py index 0503c0b386..e7fb4244d0 100644 --- a/tinygrad/renderer/ptx.py +++ b/tinygrad/renderer/ptx.py @@ -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)