from typing import Literal, Callable import math, sys from collections import defaultdict, Counter from tinygrad.renderer import tc from tinygrad.uop.ops import GroupOp, Ops, UOp, PatternMatcher, UPat, range_str, axis_letters from tinygrad.uop.weak import commit_weak_consts from tinygrad.helpers import strip_parens, getenv, prod, dedup, Target, IMAGE, FLOAT16, is_image_shape from tinygrad.dtype import dtypes, DType, AddrSpace, truncate, to_storage_scalar from tinygrad.renderer import Renderer base_rewrite = PatternMatcher([ # local/reg buffers (UPat(Ops.BUFFER, name="x"), lambda ctx,x: ctx.render_buffer(x)), (UPat(Ops.BINARY, name="x"), lambda ctx,x: f'const unsigned char {ctx[x]}[] = "' + ''.join(f'\\x{b:02x}' for b in x.arg) + '";'), # range/loop/if/endif (UPat(Ops.RANGE, dtypes.void), lambda ctx: "for (;;) {"), (UPat(Ops.RANGE, name="x"), lambda ctx,x: f"for ({ctx.render_dtype(x.dtype)} {ctx[x]} = 0; {ctx[x]} < {ctx[x.src[0]]}; {ctx[x]}++) {{"), (UPat(Ops.BACKEDGE, src=(UPat(), UPat(Ops.RANGE), UPat(name="c", dtype=dtypes.bool))), lambda ctx,c: f" if (!({ctx[c]})) {{ break; }}\n}}"), (UPat(Ops.IF, name="x"), lambda ctx,x: f"if ({ctx[x.src[0]]}) {{"), (UPat((Ops.ENDIF, Ops.END)), lambda ctx: "}"), # const (UPat.cvar("c").cast(dtypes.floats, name="x"), lambda ctx,x,c: None if math.isfinite(v:=c.val) else \ f"({ctx.render_cast(x, ctx.nan if math.isnan(v) else ctx.infinity if v > 0 else f'-{ctx.infinity}')})"), (UPat.cvar("c").cast(dtypes.float), lambda ctx,c: f"{c.val}f"), (UPat.cvar("c").cast(dtypes.int64), lambda ctx,c: f"{c.val}l"), (UPat.cvar("c").cast(dtypes.uint64, name="x"), lambda ctx,x,c: f"{truncate[x.dtype](c.val)}ul"), (UPat.cvar("c").cast(dtypes.uint32, name="x"), lambda ctx,x,c: f"{truncate[x.dtype](c.val)}u"), (UPat.cvar("c").cast(dtypes.bool), lambda ctx,c: "1" if c.val else "0"), # consts are rendered to larger type and casted (UPat.cvar("c").cast((*dtypes.fp8s, dtypes.bfloat16, dtypes.half), name="x"), lambda ctx,x,c: f"({ctx.render_cast(x, f'{c.val}f')})"), (UPat.cvar("c").cast((dtypes.uint8, dtypes.uint16), name="x"), lambda ctx,x,c: f"({ctx.render_cast(x, f'{c.val}u')})"), (UPat.cvar("c").cast((dtypes.int8, dtypes.int16), name="x"), lambda ctx,x,c: f"({ctx.render_cast(x, str(c.val))})"), # default const render (UPat.cvar("c").cast(), lambda ctx,c: str(c.val)), # casting (UPat(Ops.CAST, name="x"), lambda ctx,x: f"({ctx.render_cast(x, ctx[x.src[0]])})"), (UPat(Ops.BITCAST, name="x"), lambda ctx,x: f"(({ctx._render_dtype(x.dtype, addrspace=x.addrspace)})({ctx[x.src[0]]}))" if x.addrspace in (AddrSpace.GLOBAL, AddrSpace.LOCAL) else None), (UPat(Ops.BITCAST, name="x"), lambda ctx,x: f"__builtin_bit_cast({ctx.render_type(x)}, ({ctx.render_type(x.src[0])})({ctx[x.src[0]]}))"), # GPU stuff (UPat(Ops.BARRIER), lambda ctx: ctx.barrier), (UPat(Ops.SPECIAL, name="x"), lambda ctx,x: f"{ctx.code_for_workitem[x.arg[0]](x.arg[-1])}; /* {(x.src[0]).render()} */"), # SHRINK/INDEX (UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var('idx')), name="x"), lambda ctx,**kwargs: ctx.render_index(**kwargs)), (UPat(Ops.SHRINK, src=(UPat.var("buf"), UPat.var('idx'), UPat.cvar().cast()), name="x"), lambda ctx,**kwargs: ctx.render_index(**kwargs)), (UPat(Ops.STACK, name="x"), lambda ctx,x: f"{ctx.float4.replace('float4', ctx.render_type(x))}" + \ f"{ctx.float4_style[0]}{','.join([ctx[y] for y in x.src])}{ctx.float4_style[1]}"), # load/store (UPat(Ops.LOAD, src=(UPat.var('bidx'),)), lambda ctx,bidx: f"({ctx.render_access(bidx)})"), (UPat(Ops.LOAD, src=(UPat.var("bidx"), UPat.var("var"), UPat.var("gate"))), lambda ctx,bidx,var,gate: f"({ctx[gate]}?{ctx.render_access(bidx)}:{ctx[var]})"), (UPat(Ops.STORE, src=(UPat.var('bidx'), UPat.var("var"))), lambda ctx,bidx,var: f"{ctx.render_access(bidx)} = {ctx[var]};"), # alu/gep (UPat(Ops.WMMA, name="x"), lambda ctx,x: f"__{_wmma_name(x)}({ctx[x.src[0]]}, {ctx[x.src[1]]}, {ctx[x.src[2]]})"), (UPat(GroupOp.ALU, name="x"), lambda ctx,x: ctx.code_for_op[x.op]( *([strip_parens(ctx[v]) if v.op == x.op and x.op in {Ops.ADD, Ops.MUL, Ops.XOR, Ops.OR, Ops.AND} else ctx[v] for v in x.src]), x.dtype)), # call an external function by symbol: the CUSTOM_FUNCTION body names the callee, the srcs are the args and set the types (UPat(Ops.CALL, src=(UPat(Ops.CUSTOM_FUNCTION, name="f"),), allow_any_len=True, name="x"), lambda ctx,x,f: f"(({ctx.abi}{ctx.render_dtype(x.dtype)}(*)({', '.join(map(ctx.render_type, x.src[1:]))}))({f.arg.name}))" f"({', '.join(f'({ctx.render_type(y)})({ctx[y]})' for y in x.src[1:])})" + (";" if x.dtype is dtypes.void else "")), (UPat(Ops.CALL, src=(UPat(Ops.LINEAR, name="body"),), allow_any_len=True, name="x"), lambda ctx,x,body: f"{ctx.fn_names[body]}(" f"{', '.join(f'({ctx.param_type(p)}){ctx[x.src[s+1]]}' for s,p in sorted((u.arg.slot,u) for u in body.src if u.op is Ops.PARAM))});"), # custom passes through with format (UPat((Ops.CUSTOM, Ops.CUSTOMI), name="x"), lambda ctx,x: x.arg[0].format(*[ctx[y] for y in x.src])), ]) def create_non_native_float_pats(dts:tuple[DType, ...], casting:bool=True): patterns = PatternMatcher([ # a weak CONST states no width and cannot be restated: commit it at the emulated dtype a sibling src states (UPat(GroupOp.ALU, name="x"), lambda x, dts=dts: commit_weak_consts(x, next((s.dtype for s in x.src if s.dtype in dts), None))), (UPat(GroupOp.ALU-{Ops.WHERE}, dtype=dts, name="x"), lambda x: UOp(x.op, src=tuple(vv.cast(dtypes.float) for vv in x.src), arg=x.arg).cast(x.dtype)), (UPat(GroupOp.ALU, dtypes.bool, name="alu", src=(UPat.var("x", dtype=dts), UPat.var("y", dtype=dts))), lambda alu,x,y: UOp(alu.op, src=(x.cast(dtypes.float), y.cast(dtypes.float)), arg=alu.arg))]) if casting: # add float intermediate casting patterns += PatternMatcher([ (UPat(Ops.CAST, dts, (UPat.var("x"),), name="y"), lambda x,y: x.cast(dtypes.float).cast(y.dtype) if x.dtype!=dtypes.float and x.op is not Ops.CONST else None), (UPat(Ops.CAST, name="x", src=(UPat.var("y", dts),)), lambda x,y: y.cast(dtypes.float).cast(x.dtype) if x.dtype!=dtypes.float else None)]) return patterns def cast_float_to_bf16(x: UOp) -> UOp: assert x.dtype == dtypes.float, "cast float -> bf16 must start with float" x = x.bitcast(dtypes.uint) x = (-x & 0x7f800000).ne(0).where(x + ((x >> 16) & 1) + 0x7fff, (x & 0xffff).ne(0).where((x | 0x10000), x)) return (x >> 16).cast(dtypes.ushort).bitcast(dtypes.bfloat16) # manual bfloat16 casting patterns (shared between LLVM, Clang, and AMD renderers to avoid compiler intrinsics) pm_manual_bf16_cast = PatternMatcher([ (UPat(Ops.CAST, dtypes.float, (UPat.var("x", dtypes.bfloat16),)), lambda x: (x.bitcast(dtypes.ushort).cast(dtypes.uint)<<16).bitcast(dtypes.float)), (UPat(Ops.CAST, dtype=dtypes.bfloat16, src=(UPat.var("x", dtype=dtypes.float),)), cast_float_to_bf16), ]) # a bfloat16 stored as ushort renders its const as the bit pattern pm_bf16_ushort_const = PatternMatcher([(UPat.cvar("c").cast(dtypes.bfloat16), lambda ctx,c: f"{to_storage_scalar(c.val, dtypes.bfloat16)}u")]) def uops_to_dtypes(uops:list[UOp]) -> list[tuple[DType, int]]: return dedup((u.dtype, u.max_numel()) for u in uops if u.addrspace in (AddrSpace.ALU, None) and u.dtype != dtypes.void and u._shape is not None) def _wmma_name(u:UOp) -> str: return f"WMMA_{'_'.join(map(str, u.arg[0]))}_{u.arg[1].name}_{u.dtype.name}" # (name, dims, dtype_in, dtype_out, upcast_sizes) def wmma_args(uops:list[UOp]): return dedup((_wmma_name(uop), uop.arg[0], uop.arg[1], uop.dtype, tuple(x.shape[-1] for x in uop.src)) for uop in uops if uop.op is Ops.WMMA) class CStyleLanguage(Renderer): abi: str = "" kernel_typedef: str = "void" buffer_prefix: str = "" buffer_suffix: str = "" smem_align: str = "" smem_prefix: str = "" smem_prefix_for_cast: bool = True var_prefix: str = "const " var_suffix: str = "" barrier: str = "" code_for_workitem: dict[Literal["g", "l"], Callable] = {} extra_args: list[str] = [] float4: str|None = None float4_style: tuple[str, str] = ('(', ')') gep_arr_threshold: int = 4 type_map: dict[DType, str] = { dtypes.void: "void", dtypes.bool: "bool", dtypes.i8: "signed char", dtypes.u8: "unsigned char", dtypes.i16: "short", dtypes.u16: "unsigned short", dtypes.i32: "int", dtypes.u32: "unsigned int", dtypes.i64: "long", dtypes.u64: "unsigned long", dtypes.f16: "half", dtypes.bf16: "__bf16", dtypes.f32: "float", dtypes.f64: "double"} infinity: str = "INFINITY" nan: str = "NAN" code_for_op: dict = { Ops.SQRT: lambda x,dtype: f"sqrt({x})", Ops.RECIPROCAL: lambda x,dtype: f"(1/{x})", Ops.NEG: lambda x,dtype: f"-{x}", Ops.EXP2: lambda x,dtype: f"exp2({x})", Ops.LOG2: lambda x,dtype: f"log2({x})", Ops.SIN: lambda x,dtype: f"sin({x})", Ops.TRUNC: lambda x,dtype: f"trunc({x})", Ops.AND: lambda a,b,dtype: f"({a}&{b})", Ops.XOR: lambda a,b,dtype: f"({a}^{b})", Ops.OR: lambda a,b,dtype: f"({a}|{b})", Ops.ADD: lambda a,b,dtype: f"({a}+{b})", Ops.SUB: lambda a,b,dtype: f"({a}-{b})", Ops.MUL: lambda a,b,dtype: f"({a}*{b})", Ops.CMOD: lambda a,b,dtype: f"({a}%{b})", Ops.CDIV: lambda a,b,dtype: f"({a}/{b})", Ops.CMPNE: lambda a,b,dtype: f"({a}!={b})", Ops.SHR: lambda a,b,dtype: f"({a}>>{b})", Ops.SHL: lambda a,b,dtype: f"({a}<<{b})", Ops.CMPLT: lambda a,b,dtype: f"({a}<{b})", Ops.WHERE: lambda a,b,c,dtype: f"({a}?{b}:{c})", Ops.CMPEQ: lambda a,b,dtype: f"({a}=={b})"} string_rewrite = base_rewrite def render_kernel(self, function_name:str, kernel:list[str], bufs:list[tuple[str,tuple[UOp,bool]]], uops:list[UOp], prefix=None) -> str: tmp = "" if any(is_image_shape(u._shape) for _,(u,_) in bufs): tmp = "const sampler_t smp = CLK_NORMALIZED_COORDS_FALSE | CLK_ADDRESS_CLAMP | CLK_FILTER_NEAREST;\n" buftypes = [(name, ("volatile " if u.arg.volatile else "")+(self.var_prefix if u.addrspace == AddrSpace.ALU else "")+ self._render_dtype(u.dtype, sz=1, addrspace=u.addrspace, mutable=mutable, shape=u._shape)+ (self.var_suffix if u.addrspace == AddrSpace.ALU else self.buffer_suffix)) for name,(u,mutable) in bufs] local_dims = [u.src[0] for u in uops if u.op is Ops.SPECIAL and u.arg[0] == "l"] launch_bounds = prod([d.vmax for d in local_dims]) prg = ''.join([f"{self.kernel_typedef.format(launch_bounds=launch_bounds)} {function_name}(",] + [', '.join([f'{t} {name}' for name,t in buftypes] + self.extra_args)] + [") {\n" + tmp] + ['\n'.join(kernel), "\n}"]) return prg if prefix is None else "\n".join(prefix)+f"\n{prg}" def render_index(self, x:UOp, buf:UOp, idx:UOp): if buf.addrspace == AddrSpace.ALU: # this is lane access in C if not (idx.op is Ops.CAST and idx.src[0].op is Ops.CONST): return f"({self[buf]})[{self[idx]}]" return self[buf]+(f"[{idx.src[0].val}]" if buf.max_numel() > self.gep_arr_threshold else f".{'xyzwabcd'[idx.src[0].val]}") return f"({self[buf]}+{strip_parens(self[idx]) if idx.arg == Ops.ADD else self[idx]})" def render_buffer(self, x:UOp): prefix = f"{self.smem_align}{self.smem_prefix}" if x.addrspace == AddrSpace.LOCAL else "" suffix = f"[{x.max_numel()}]" return f"{prefix}{self._render_dtype(x.dtype)} {self[x]}{suffix};" def _render_dtype(self, dtype:DType, sz:int=1, addrspace=AddrSpace.ALU, mutable=True, override_ptr=False, shape=None): if is_image_shape(shape): return f"{'write_only' if mutable else 'read_only'} image2d_t" prefix, suffix = "", "" if addrspace in (AddrSpace.LOCAL, AddrSpace.GLOBAL): if addrspace == AddrSpace.LOCAL and self.smem_prefix_for_cast: prefix = self.smem_prefix if addrspace == AddrSpace.GLOBAL: prefix = self.buffer_prefix if addrspace in (AddrSpace.LOCAL, AddrSpace.GLOBAL) or override_ptr: suffix = "*" if sz > 1: return prefix + self.type_map[dtype].replace(" ", "_") + str(sz) + suffix return prefix + self.type_map[dtype] + suffix def render_type(self, u:UOp): return self._render_dtype(u.dtype, u.max_numel(), u.addrspace, shape=u._shape, override_ptr=u.op is Ops.INDEX and u.addrspace is AddrSpace.REG) def render_ptr(self, u:UOp): # the address of an access, vector-cast if the access reads/writes more lanes than the pointer's scalar type if u.max_numel() > 1 or u.dtype != u.src[0].dtype: return f"(({self._render_dtype(u.dtype, u.max_numel(), u.addrspace, override_ptr=True, shape=u._shape)})({self[u]}))" else: return f"{self[u]}" def render_access(self, u:UOp): return f"*{self.render_ptr(u)}" def render_cast(self, u:UOp, val:str) -> str: return f"({self.render_type(u)})({val})" # LEGACY def render_dtype(self, dt:DType, mutable=True) -> str: return self._render_dtype(dt, 1, AddrSpace.REG) def __getitem__(self, key): return self.r[key] # hacky helper def _render(self, uops:list[UOp]) -> tuple[str, list[str], list[tuple[str,tuple[UOp,bool]]]]: r: dict[UOp, str] = {} self.r = r child_count = Counter(v for ru in uops for v in ru.src) # find which PARAMs are stored to with a single toposort writable_params = {u for u in UOp.sink(*[u.src[0] for u in uops if u.op is Ops.STORE]).toposort(lambda u: u.op not in {Ops.END, Ops.BACKEDGE}) if u.op is Ops.PARAM} bufs: dict[UOp, tuple[str, tuple[UOp, bool]]] = {} kernel = [] depth = 1 c: defaultdict[str, int] = defaultdict(int) name = "test" for u in uops: if u.op in {Ops.NOOP, Ops.CONST, Ops.CUSTOM_FUNCTION}: continue if u.op == Ops.STACK and len(u.src) == 0: continue if u.op is Ops.AFTER: if u.dtype is dtypes.void: continue r[u] = r[u.src[0]] continue if u.op is Ops.SINK: if u.arg is not None: name = u.arg.function_name continue if u.op is Ops.PARAM: r[u] = (u.arg.name.replace(":", "_") if u.arg.name is not None else f"data{u.arg.slot}") + \ "_" + '_'.join([str(x) for x in u.shape]) bufs[u] = (r[u], (u, u in writable_params)) continue # naming prefix = None if u.op is Ops.SPECIAL: r[u] = u.arg elif u.op is Ops.RANGE: r[u] = f"{axis_letters[u.axis_type]}idx"+range_str(u) else: prefix = {Ops.WMMA: "wmma", Ops.BUFFER: "buf", Ops.CAST: "cast", Ops.BITCAST: "cast", Ops.STACK: "cast", Ops.INDEX: "bidx", Ops.LOAD: "val"}.get(u.op, "alu") r[u] = f"{prefix}{c[prefix]}" l: str|None = self.string_rewrite.rewrite(u, ctx=self) assert l is not None, f"failed to render {u.op} {u.dtype} {[(x.op,x.dtype) for x in u.src]} {u.arg}" if u.op in {Ops.ENDIF, Ops.END, Ops.BACKEDGE}: depth -= 1 if (u.op is not Ops.CAST or u.max_numel() == 1) and ((u.op is Ops.CAST and u.src[0].op is Ops.CONST) or \ u.op in {Ops.INDEX, Ops.SHRINK, Ops.CUSTOMI} or \ (u.op is Ops.LOAD and u.src[0].addrspace == AddrSpace.REG and child_count[u] == 1) or \ (u.op in {Ops.CAST, Ops.BITCAST} and u.addrspace in (AddrSpace.GLOBAL, AddrSpace.LOCAL)) or \ (u.op in {Ops.STACK, *(GroupOp.ALU-{Ops.WHERE}), Ops.CAST, Ops.BITCAST} and child_count[u] == 1 and not getenv("EXPAND_SSA"))): r[u] = l else: if u.op not in {Ops.RANGE, Ops.BUFFER, Ops.BINARY} and u.dtype != dtypes.void: l = f"{self.render_type(u)} {r[u]} = {l}" + (";" if u.op is not Ops.SPECIAL else "") elif u.op is Ops.CAST: l += ";" # a discarded value renders as a bare statement kernel.append("\n".join(" "*depth + line for line in l.split("\n"))) if prefix: c[prefix] += 1 # if it was used, increment if u.op in {Ops.IF, Ops.RANGE}: depth += 1 del self.r # NOTE: this relies on bufs dict preserving order return (name, kernel, list(bufs.values())) def param_type(self, p:UOp): return "volatile "*p.arg.volatile + self._render_dtype(p.dtype, 1, p.addrspace, True, p.addrspace != AddrSpace.ALU) def render(self, uops:list[UOp]) -> str: prefix, call_bodies, self.fn_names = [], [], dict[UOp, str]() topo = UOp.sink(*uops).toposort(enter_calls=True) prefix += [f"extern void {f}();" for f in dedup(u.arg.name for u in topo if u.op is Ops.CUSTOM_FUNCTION)] # symbols to link for body in (u for u in topo if u.op is Ops.LINEAR): self.fn_names[body] = body.arg + (f"_{n}" if (n:=sum(b.arg == body.arg for b in self.fn_names)) else "") # a name traced with other args _, call, bufs = self._render(body.src) params = ', '.join(f"{self.param_type(p)} {n}" for n,(p,_) in bufs) prefix.append(f"static inline void {self.fn_names[body]}({params}) {{\n" + '\n'.join(call) + "\n}") call_bodies.extend(body.src) return self.render_kernel(*self._render(uops), call_bodies+list(uops), prefix or None) class ClangRenderer(CStyleLanguage): float4 = "(float4)" float4_style = ('{', '}') gep_arr_threshold = 0 has_local = False global_max = (1, 0, 0) infinity = "__builtin_inff()" nan = '__builtin_nanf("")' # language options barrier = "__atomic_thread_fence(__ATOMIC_SEQ_CST);" buffer_suffix = " restrict" type_map = {**CStyleLanguage.type_map, dtypes.bool:"_Bool", dtypes.f16:"__fp16"} code_for_op = {**({k:v for k,v in CStyleLanguage.code_for_op.items() if k not in [Ops.EXP2, Ops.SIN, Ops.LOG2, Ops.RECIPROCAL]}), Ops.SQRT: lambda x,dtype: f"__builtin_sqrt({x})" if dtype == dtypes.float64 else f"__builtin_sqrtf({x})", Ops.TRUNC: lambda x,dtype: f"__builtin_trunc({x})" if dtype == dtypes.float64 else f"__builtin_truncf({x})", Ops.FDIV: lambda a,b,dtype: f"({a}/{b})"} # LLVM legalizes double => half cast on systems that don't support it natively (like x86 cpus without AVX512-FP16) into a compiler-rt libcall. extra_matcher = PatternMatcher([(UPat.var("x", dtypes.float64).cast(dtypes.float16), lambda x: x.cast(dtypes.float32).cast(dtypes.float16))]) \ + create_non_native_float_pats((dtypes.bfloat16,)) + pm_manual_bf16_cast if sys.platform == 'win32': abi = "__attribute__((ms_abi)) " kernel_typedef = abi + "void" def render_vector_prefix(self, dt:DType, count:int) -> str: # round (down) to power of two (this is actually the default clang behavior) alignment = 2**int(math.log2(dt.itemsize * count)) if getenv("ALIGNED", 1) and not dtypes.is_bool(dt) else 1 vec = self._render_dtype(dt, count, AddrSpace.REG) return f"typedef {self.render_dtype(dt)} {vec} __attribute__((aligned({alignment}),ext_vector_type({count})));" def _render_defines(self, uops) -> list[str]: return [self.render_vector_prefix(dt, count) for dt, count in uops_to_dtypes(uops) if count > 1] def _render_body(self, function_name, kernel, bufs, uops, pref=None) -> str: return super().render_kernel(function_name, kernel, bufs, uops, pref) def _render_entry(self, function_name:str, bufs:list[tuple[str,tuple[UOp,bool]]]) -> str: return "" def render_kernel(self, function_name, kernel, bufs, uops, prefix=None) -> str: defines = '\n'.join(self._render_defines(uops)) return defines + "\n" + self._render_body(function_name, kernel, bufs, uops, prefix) + "\n" + self._render_entry(function_name, bufs) def supported_dtypes(self): return {d for d in super().supported_dtypes() if (d != dtypes.bfloat16 or self.target.arch.startswith(("x86", "arm"))) and d not in dtypes.fp8s} def __init__(self, target:Target): super().__init__(target) from tinygrad.runtime.support.compiler_cpu import ClangCompiler self.compiler = ClangCompiler(target.arch.split(",")) class OpenCLRenderer(CStyleLanguage): # language options kernel_typedef = "__kernel void" buffer_prefix = "__global " smem_align = "__attribute__ ((aligned (16))) " smem_prefix = "__local " barrier = "barrier(CLK_LOCAL_MEM_FENCE);" float4 = "(float4)" code_for_workitem = {"g": lambda x: f"get_group_id({x})", "l": lambda x: f"get_local_id({x})"} type_map = {**CStyleLanguage.type_map, dtypes.i8: "char", dtypes.u8: "uchar", dtypes.u32: "uint", dtypes.u16: "ushort", dtypes.u64: "ulong", dtypes.bf16: "ushort"} extra_matcher = create_non_native_float_pats((dtypes.bfloat16,)) + pm_manual_bf16_cast string_rewrite = pm_bf16_ushort_const + PatternMatcher([ (UPat(Ops.BITCAST, name="x"), lambda ctx,x: f"as_{ctx.render_dtype(x.dtype)}(({ctx.render_dtype(x.src[0].dtype)})({ctx[x.src[0]]}))" if x.addrspace not in (AddrSpace.GLOBAL, AddrSpace.LOCAL) else None), # load/store image (OpenCL) (UPat.var('buf').index(UPat.var('idx_y'), UPat.var('idx_x')), lambda ctx,buf,idx_y,idx_x: f"IMAGE<{ctx[buf]}, {ctx[idx_y]}, {ctx[idx_x]}>"), (UPat(Ops.LOAD, dtype=dtypes.float, src=(UPat.var('buf').index(UPat.var('idx_y'), UPat.var('idx_x')), UPat.var("var"), UPat.var("gate"))), lambda ctx,buf,idx_y,idx_x,var,gate: f"({ctx[gate]}?read_imagef({ctx[buf]}, smp, (int2)({ctx[idx_x]},{ctx[idx_y]})):{ctx[var]})"), (UPat(Ops.LOAD, dtype=dtypes.float, src=(UPat.var('buf').index(UPat.var('idx_y'), UPat.var('idx_x')),)), lambda ctx,buf,idx_y,idx_x: f"read_imagef({ctx[buf]}, smp, (int2)({ctx[idx_x]},{ctx[idx_y]}))"), (UPat(Ops.STORE, src=(UPat.var('buf').index(UPat.var('idx_y'), UPat.var('idx_x')), UPat.var("var", dtypes.float))), lambda ctx,buf,idx_y,idx_x,var: f"write_imagef({ctx[buf]}, (int2)({ctx[idx_x]},{ctx[idx_y]}), {ctx[var]});"), ]) + base_rewrite def render_kernel(self, function_name, kernel, bufs, uops, prefix=None) -> str: if any(uop.dtype == dtypes.half for uop in uops): prefix = (["#pragma OPENCL EXTENSION cl_khr_fp16 : enable"] + (prefix or [])) return super().render_kernel(function_name, kernel, bufs, uops, prefix) def supported_dtypes(self): return {d for d in super().supported_dtypes() if (d != dtypes.half or "cl_khr_fp16" in self.target.arch) and (d != dtypes.double or "cl_khr_fp64" in self.target.arch) and d not in dtypes.fp8s} class MetalRenderer(CStyleLanguage): def __init__(self, target:Target): super().__init__(target) from tinygrad.runtime.ops_metal import MetalCompiler self.compiler, self.tensor_cores = MetalCompiler(), tc.metal if target.arch.startswith("Apple") and int(target.arch[5:]) >= 7 else [] # language options kernel_typedef = "kernel void" buffer_prefix = "device " smem_prefix = "threadgroup __attribute__((aligned(16))) " barrier = "threadgroup_barrier(mem_flags::mem_threadgroup);" float4 = "float4" code_for_workitem = {"g": lambda x: f"gid.{chr(120+int(x))}", "l": lambda x: f"lid.{chr(120+int(x))}"} # uint3 used for gid/lid - TODO: this should probably be `ushort3 lid [[thread_position_in_threadgroup]]` extra_args = ['constant args_t& args [[buffer(0)]]', 'uint3 gid [[threadgroup_position_in_grid]]', 'uint3 lid [[thread_position_in_threadgroup]]'] type_map = {**CStyleLanguage.type_map, dtypes.u32: "uint", dtypes.bf16: "bfloat"} # precise::sin code_for_op = {**CStyleLanguage.code_for_op, Ops.SIN: lambda x,dtype: f"precise::sin({x})"} # upcast to float32 all the ops that don't support bfloat16 extra_matcher = PatternMatcher([ # NOTE: this is copied from PTX (UPat((Ops.SQRT, Ops.EXP2, Ops.LOG2, Ops.SIN, Ops.TRUNC), dtype=dtypes.bfloat16, name="x"), lambda x: (UOp(x.op, src=tuple(vv.cast(dtypes.float) for vv in x.src), arg=x.arg).cast(dtypes.bfloat16))), ]) + pm_manual_bf16_cast string_rewrite = PatternMatcher([ (UPat(Ops.BITCAST, name="x"), lambda ctx,x: f"as_type<{ctx.render_dtype(x.dtype)}>(({ctx.render_dtype(x.src[0].dtype)})({ctx[x.src[0]]}))" if x.addrspace not in (AddrSpace.GLOBAL, AddrSpace.LOCAL) else None), ]) + base_rewrite def render_kernel(self, function_name, kernel, bufs, uops, prefix=None): prefix = ["#include ","using namespace metal;"] for name, _, dtype_in, dtype_out, _ in wmma_args(uops): dstr_out, dstr_in = self._render_dtype(dtype_out, 2, AddrSpace.REG), self._render_dtype(dtype_in, 2, AddrSpace.REG) prefix.append( f"""{dstr_out} __{name}({dstr_in} a, {dstr_in} b, {dstr_out} c){{ simdgroup_{self.render_dtype(dtype_in)}8x8 mat_a, mat_b; simdgroup_{self.render_dtype(dtype_out)}8x8 mat_c; mat_a.thread_elements()[0] = a[0]; mat_b.thread_elements()[0] = b[0]; mat_c.thread_elements()[0] = c[0]; mat_a.thread_elements()[1] = a[1]; mat_b.thread_elements()[1] = b[1]; mat_c.thread_elements()[1] = c[1]; simdgroup_multiply_accumulate(mat_c, mat_a, mat_b, mat_c);\n return {dstr_out}(mat_c.thread_elements()[0], mat_c.thread_elements()[1]);\n}}""") # one argument buffer: a struct of the buffer pointers and the scalars, so a binding is a gpu address (an icb offset keeps 32 bits, an address 64) args = [(name, self._render_dtype(u.dtype, addrspace=u.addrspace)) for name, (u, _) in bufs] prefix.append("struct args_t { " + " ".join(f"{t} {n};" for n, t in args) + " };") return super().render_kernel(function_name, [" " + " ".join(f"{t} {n} = args.{n};" for n, t in args)] + kernel, [], uops, prefix) def supported_dtypes(self): return {d for d in super().supported_dtypes() if (d != dtypes.bfloat16 or ((arch:=self.target.arch).startswith("Apple") and int(arch[5:]) >= 6)) and d not in dtypes.fp8s+(dtypes.double,)} _nms = list("xyzwabcdefghijkl") + [f'v{i}' for i in range(16, 32)] class CUDARenderer(CStyleLanguage): global_max = (2147483647, 65535, 65535) local_max = (1024, 1024, 64) shared_max = 49152 def __init__(self, target:Target, use_nvcc=False): super().__init__(target) from tinygrad.runtime.support.compiler_cuda import NVRTCCompiler, NVCCCompiler iface, dev, arch = target.interface, target.device, target.arch self.compiler = (NVCCCompiler if use_nvcc else NVRTCCompiler)(arch, ptx=iface.startswith("MOCK") or dev == "CUDA", cache_key=dev.lower()) self.tensor_cores = tc.get_cuda(arch) # language options # https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html kernel_typedef = "extern \"C\" __global__ void __launch_bounds__({launch_bounds})" smem_prefix = "__shared__ __align__(16) " smem_prefix_for_cast = False barrier = "__syncthreads();" float4 = "make_float4" gep_arr_threshold = 8 code_for_workitem = {"g": lambda x: f"blockIdx.{chr(120+int(x))}", "l": lambda x: f"threadIdx.{chr(120+int(x))}"} code_for_op = { **CStyleLanguage.code_for_op, Ops.TRUNC: lambda x,dtype: f"htrunc({x})" if dtype in (dtypes.half, dtypes.bfloat16) else f"trunc({x})", Ops.SIN: lambda x,dtype: f"hsin({x})" if dtype in (dtypes.half, dtypes.bfloat16) else f"sin({x})", Ops.LOG2: lambda x,dtype: f"hlog2({x})" if dtype in (dtypes.half, dtypes.bfloat16) else f"log2({x})", Ops.EXP2: lambda x,dtype: f"hexp2({x})" if dtype in (dtypes.half, dtypes.bfloat16) else f"exp2({x})", Ops.SQRT: lambda x,dtype: f"hsqrt({x})" if dtype in (dtypes.half, dtypes.bfloat16) else f"sqrt({x})", Ops.RECIPROCAL: lambda x,dtype: f"hrcp({x})" if dtype in (dtypes.half, dtypes.bfloat16) else f"(1/{x})" } type_map = {**CStyleLanguage.type_map, dtypes.u32: "uint", dtypes.bf16: "nv_bfloat16", dtypes.fp8e4m3: "__nv_fp8_e4m3", dtypes.fp8e5m2: "__nv_fp8_e5m2"} extra_matcher = create_non_native_float_pats(dtypes.fp8s, casting=False) + PatternMatcher([ (UPat(Ops.CAST, dtypes.fp8s, UPat.var("x", dtypes.fp8s), name='y'), lambda x,y: x.cast(dtypes.float).cast(y.dtype) if x.dtype!=y.dtype else None), ]) string_rewrite = PatternMatcher([ (UPat(Ops.BITCAST, name="x"), lambda ctx,x: f"tg_bitcast<{ctx.render_dtype(x.dtype)}>(({ctx.render_dtype(x.src[0].dtype)})({ctx[x.src[0]]}))" if x.addrspace not in (AddrSpace.GLOBAL, AddrSpace.LOCAL) else None), ]) + base_rewrite def render_vector_prefix(self, dt:DType, count:int) -> str: vec, scal = self._render_dtype(dt, count, AddrSpace.REG), self.render_dtype(dt) elems, header = ', '.join(_nms[:count]), ', '.join([f"{scal} {x}" for x in _nms[:count]]) return f"struct __align__({dt.itemsize * count}) {vec} {{ {scal} {elems}; }}; " \ f"__device__ {vec} make_{vec}({header}) {{ {vec} r={{{elems}}}; return r; }}" def render_kernel(self, function_name, kernel, bufs, uops, prefix=None): prefix = ["typedef unsigned int uint;", "#define INFINITY (__int_as_float(0x7f800000))", "#define NAN (__int_as_float(0x7fffffff))", "template __device__ __forceinline__ T tg_bitcast(F v) { union U { F f; T t; }; U u; u.f = v; return u.t; }"] used_dtypes = uops_to_dtypes(uops) if any(dt in dtypes.fp8s for dt, _ in used_dtypes): prefix.append("#include ") if any(dt == dtypes.half for dt, _ in used_dtypes): prefix.append("#include ") if any(dt == dtypes.bfloat16 for dt, _ in used_dtypes): prefix.append("#include ") prefix += [self.render_vector_prefix(dt, count) for dt, count in used_dtypes if (count in (4,8) and dt in {dtypes.half, dtypes.bfloat16}) or (count in (2,4,8,16) and dt in dtypes.fp8s)] dt_map_in = { dtypes.f32: "tf32", dtypes.f16: "f16", dtypes.bf16: "bf16", dtypes.fp8e4m3: "e4m3", dtypes.fp8e5m2: "e5m2" } dt_map_out = { dtypes.f32: "f32", dtypes.f16: "f16" } for name, (N, M, K), dtype_in, dtype_out, upcast_sizes in wmma_args(uops): wmma_dtypes = [self._render_dtype(dtype, size, AddrSpace.REG) for dtype, size in zip([dtype_in, dtype_in, dtype_out], upcast_sizes)] n_operands = [size*dtype.itemsize//4 for dtype, size in zip([dtype_in, dtype_in, dtype_out], upcast_sizes)] # 4 => CUDA reg size in bytes operands = [f"%{i}" for i in range(sum(n_operands))] # mma operands => {c}, {a}, {b}, {c} prefix.append(f"""__device__ {wmma_dtypes[2]} __{name}({wmma_dtypes[0]} a, {wmma_dtypes[1]} b, {wmma_dtypes[2]} c){{ int *a_pk = (int *)(&a), *b_pk = (int *)(&b), *c_pk = (int *)(&c); asm("mma.sync.aligned.m{M}n{N}k{K}.row.col.{dt_map_out[dtype_out]}.{dt_map_in[dtype_in]}.{dt_map_in[dtype_in]}.{dt_map_out[dtype_out]}" "{{{", ".join(operands[:n_operands[2]])}}}, {{{", ".join(operands[n_operands[2]:n_operands[2]+n_operands[0]])}}}," "{{{", ".join(operands[-n_operands[1]:])}}}, {{{", ".join(operands[:n_operands[2]])}}};" : {", ".join([f'"+r"(c_pk[{i}])' for i in range(n_operands[2])])} : {", ".join([f'"r"(a_pk[{i}])' for i in range(n_operands[0])])}, {", ".join([f'"r"(b_pk[{i}])' for i in range(n_operands[1])])}); return c;\n}}""") return super().render_kernel(function_name, kernel, bufs, uops, prefix=prefix) def supported_dtypes(self): ver = int(self.target.arch[3:]) return {d for d in super().supported_dtypes() if (d != dtypes.half or ver >= 53) and (d != dtypes.bfloat16 or ver >= 80) and (d not in dtypes.fp8_ocp or ver >= 89) and d not in dtypes.fp8_fnuz} class NVCCRenderer(CUDARenderer): def __init__(self, target:Target): super().__init__(target, use_nvcc=True) def fp8_index(dtype: DType): return dtypes.fp8s.index(dtype) % 2 def amd_fp8s(arch:str): return {"gfx942": dtypes.fp8_fnuz, "gfx950": dtypes.fp8_ocp, "gfx1200": dtypes.fp8_ocp, "gfx1201": dtypes.fp8_ocp}.get(arch, ()) def _ocml(op): return lambda x,dtype: f"__ocml_{op}_f{ {dtypes.half:16, dtypes.double:64}.get(dtype, 32)}({x})" class HIPRenderer(CStyleLanguage): shared_max = 65536 # NOTE: this is only really needed on gfx12, even though gfx11 reports the same limitation global_max = (2147483647, 65535, 65535) global_prod_max = (0xFFFFFFFF, 0xFFFFFFFF, 0xFFFFFFFF) @staticmethod def is_cdna(arch): return arch in {"gfx942", "gfx950"} @staticmethod def is_cdna4(arch): return arch == "gfx950" def __init__(self, target:Target, use_hipcc=False): # gfx942 => MI300, gfx1100 => RX 7900, gfx1201 => RX 9700 super().__init__(target) from tinygrad.runtime.support.compiler_amd import HIPCompiler, HIPCCCompiler self.compiler, self.tensor_cores = (HIPCCCompiler if use_hipcc else HIPCompiler)(target.arch), tc.get_amd(target.arch) if not self.is_cdna4(target.arch): self.extra_matcher += pm_manual_bf16_cast if self.is_cdna(target.arch): self.string_rewrite = PatternMatcher([ (UPat(Ops.WMMA, name="x"), lambda ctx,x: f"__{_wmma_name(x)}({ctx[x.src[0]]}, {ctx[x.src[1]]}, {ctx[x.src[2]]}," f" {fp8_index(x.src[0].dtype)}, {fp8_index(x.src[0].dtype)}, 0, 0, 0, 0)" if x.arg[0][2] == 128 else None), (UPat(Ops.WMMA, name="x"), lambda ctx,x: f"__{_wmma_name(x)}({ctx[x.src[0]]}, {ctx[x.src[1]]}, {ctx[x.src[2]]}, 0, 0, 0)"), ]) + self.string_rewrite if amd_fp8s(target.arch): self.extra_matcher += tc.pm_wmma_fp8(dtypes.uint64 if self.is_cdna(target.arch) else dtypes.uint32) self.string_rewrite = PatternMatcher([ (UPat.cvar("c").cast(dtypes.fp8s, name="x"), lambda ctx,x,c: f"f32_to_fp8({ctx.nan if math.isnan(v:=c.val) else ctx.infinity if v == math.inf else f'-{ctx.infinity}' if v == -math.inf else f'{v}f'}," f" {fp8_index(x.dtype)})"), (UPat(Ops.CAST, dtypes.fp8s, (UPat(dtype=dtypes.float),), name="x",), lambda ctx,x: f"f32_to_fp8({ctx[x.src[0]]}, {fp8_index(x.dtype)})"), (UPat(Ops.CAST, dtypes.float, (UPat.var("y", dtypes.fp8s),), name="x",), lambda ctx,x,y: f"__builtin_amdgcn_cvt_f32_{('fp8', 'bf8')[fp8_index(y.dtype)]}((unsigned int){ctx[x.src[0]]}, 0)"), ]) + self.string_rewrite # a LOAD flagged nontemporal renders as the cache-bypassing builtin (only used on global loads) self.string_rewrite = PatternMatcher([(UPat(Ops.LOAD, arg="nontemporal", src=(UPat.var("bidx"),)), lambda ctx,bidx: f"__builtin_nontemporal_load({ctx.render_ptr(bidx)})")]) + self.string_rewrite if not self.is_cdna4(target.arch): self.string_rewrite = pm_bf16_ushort_const + self.string_rewrite # https://clang.llvm.org/docs/AttributeReference.html#amdgpu-flat-work-group-size # NOTE: this makes hlb_cifar10 twice as fast, there may be more gains in tweaking these parameters kernel_typedef = 'extern "C" __attribute__((global)) void __attribute__((amdgpu_flat_work_group_size(1, {launch_bounds})))' code_for_workitem = {"g": lambda x: f"__ockl_get_group_id({x})", "l": lambda x: f"__ockl_get_local_id({x})"} code_for_op = {**CStyleLanguage.code_for_op, Ops.TRUNC: _ocml("trunc"), Ops.SIN: _ocml("sin"), Ops.LOG2: _ocml("log2"), Ops.EXP2: _ocml("exp2"), Ops.SQRT: _ocml("sqrt")} smem_prefix = "__attribute__((shared, aligned(16)))" smem_prefix_for_cast: bool = False barrier = '__builtin_amdgcn_fence(__ATOMIC_RELEASE, "workgroup");' + '__builtin_amdgcn_s_barrier();' + \ '__builtin_amdgcn_fence(__ATOMIC_ACQUIRE, "workgroup");' float4 = "make_float4" type_map = {**CStyleLanguage.type_map, dtypes.bf16: "hip_bfloat16", **{d: ("hip_fp8", "hip_bf8")[fp8_index(d)] for d in dtypes.fp8s}} extra_matcher = create_non_native_float_pats((dtypes.bfloat16, *dtypes.fp8s)) def asm(self, prg:UOp, lin:UOp) -> bytes: from tinygrad.renderer.amd.elf import assemble_linear return assemble_linear(prg, lin, self.target.arch) def render_vector_prefix(self, dtype:DType, count:int) -> str: vec, scal = self._render_dtype(dtype, count, AddrSpace.REG), self.render_dtype(dtype) return f"typedef {scal} {vec} __attribute__((ext_vector_type({count})));\nstatic inline __attribute__((device)) "+ \ f"{vec} make_{vec}({', '.join([f'{scal} {x}' for x in _nms[:count]])}) {{ return {{ {', '.join(_nms[:count])} }}; }}" def render_kernel(self, function_name, kernel, bufs, uops, prefix=None) -> str: prefix, ockl = [], [] type_map = {dtypes.bf16: "bf16", dtypes.f32: "f32", dtypes.f16: "f16", **{d: ("_" if self.is_cdna(self.target.arch) else "") + ("fp8_fp8", "bf8_bf8")[fp8_index(d)] for d in dtypes.fp8s}} used_dtypes = uops_to_dtypes(uops) if any(u.op is Ops.CAST and u.src[0].op is Ops.CONST and not math.isfinite(u.src[0].val) for u in uops): prefix += ["#define INFINITY (__builtin_inff())", "#define NAN (__builtin_nanf(\"\"))"] if any(u.op is Ops.SPECIAL for u in uops): prefix.append("typedef long unsigned int size_t;") ockl = [(f"__ockl_get_{name}", "unsigned int", "size_t", "const") for name in ["local_id", "group_id", "local_size"]] ocml_ops = {Ops.EXP2: ("exp2", "pure"), Ops.LOG2: ("log2", "pure"), Ops.SQRT: ("sqrt", "const"), Ops.SIN: ("sin", ""), Ops.TRUNC: ("trunc", "")} ocml = [(f"__ocml_{ocml_ops[op][0]}_f{dt.bitsize}", self.render_dtype(dt), self.render_dtype(dt), ocml_ops[op][1]) for op, dt in dedup((u.op, u.dtype) for u in uops) if op in ocml_ops and dt in (dtypes.half, dtypes.float, dtypes.double)] if any(dt == dtypes.bfloat16 for dt, _ in used_dtypes): prefix.append(f"typedef {'__bf16' if self.is_cdna4(self.target.arch) else 'unsigned short'} hip_bfloat16;") if any(dt == dtypes.half for dt, _ in used_dtypes): prefix.append("#define half _Float16") if any(dt in dtypes.fp8s for dt, _ in used_dtypes): prefix += ["typedef unsigned char hip_bf8;", "typedef unsigned char hip_fp8;"] if any((u.op is Ops.CAST and u.dtype in dtypes.fp8s and u.src[0].dtype == dtypes.float) or (u.op is Ops.CAST and u.src[0].op is Ops.CONST and u.dtype in dtypes.fp8s) for u in uops): fp8_max = amd_fp8s(self.target.arch)[0].max prefix.append(f"""static inline __attribute__((device)) unsigned char f32_to_fp8(float v, int is_bf8) {{ v = (((*(unsigned*)&v)&0x7F800000)!=0x7F800000)?__builtin_amdgcn_fmed3f(v,is_bf8?57344.0f:{fp8_max}f,is_bf8?-57344.0f:-{fp8_max}f) : v; return (unsigned char)(is_bf8?__builtin_amdgcn_cvt_pk_bf8_f32(v,v,0,false):__builtin_amdgcn_cvt_pk_fp8_f32(v,v,0,false));\n}}""") prefix += [f'extern "C" __attribute__((device{f", {atr}" if atr else ""})) {dto} {meth}({dti});' for meth,dti,dto,atr in ockl+ocml] prefix += [self.render_vector_prefix(dt, count) for dt, count in used_dtypes if count > 1] for name, (N, M, K), dtype_in, dtype_out, _ in wmma_args(uops): # TODO: handle TCs f32_bf16 and bf16_bf16 w/ wrapper if self.is_cdna(self.target.arch): if (N, M, K) == (16, 16, 16): type_map[dtypes.bf16] = 'bf16_1k' elif (N, M, K) == (16, 16, 32): type_map = {**type_map, dtypes.bf16: "_bf16", dtypes.f16: "_f16"} elif (N, M, K) == (16, 16, 128): type_map = {**type_map, dtypes.fp8e4m3: "_f8f6f4", dtypes.fp8e5m2: "_f8f6f4"} prefix.append(f"#define __{name} __builtin_amdgcn_mfma_{'scale_' if K == 128 else ''}f32_{N}x{M}x{K}{type_map[dtype_in]}") elif dtype_out == dtypes.int32: # RDNA4 uses 8 int8 values per lane in 2 VGPRs, RDNA3 uses 16 in 4 VGPRs num_regs = 2 if (is_rdna4:=self.tensor_cores == tc.amd_rdna4) else 4 prefix.append(f"typedef int wmma_int{num_regs} __attribute__((ext_vector_type({num_regs})));\n"+ f"static inline __attribute__((device)) int8 __{name}"+f"""(signed_char{num_regs*4} a, signed_char{num_regs*4} b, int8 c) {{ return __builtin_amdgcn_wmma_i32_16x16x16_iu8_w32{'_gfx12' if is_rdna4 else ''}(true, __builtin_bit_cast(wmma_int{num_regs}, a), true, __builtin_bit_cast(wmma_int{num_regs}, b), c, false);\n}}""") # #define __WMMA_16_16_16_f16_f16 __builtin_amdgcn_wmma_f16_16x16x16_f16_w32_gfx12 elif self.tensor_cores == tc.amd_rdna4: prefix.append(f"#define __{name} __builtin_amdgcn_wmma_{type_map[dtype_out]}_16x16x16_{type_map[dtype_in]}_w32_gfx12") elif dtype_out == dtypes.float: prefix.append(f"#define __{name} __builtin_amdgcn_wmma_f32_16x16x16_{'f16' if dtype_in == dtypes.half else 'bf16'}_w32") else: prefix.append(f"static inline __attribute__((device)) half8 __{name}"+"""(half16 a, half16 b, half8 c) { half16 c_frag = {}; half8 d; for (int n = 0; n < 8; n++) { c_frag[n*2] = c[n]; } c_frag = __builtin_amdgcn_wmma_f16_16x16x16_f16_w32(a, b, c_frag, false); for (int n = 0; n < 8; n++) { d[n] = c_frag[n*2]; } return d;\n}""") return super().render_kernel(function_name, kernel, bufs, uops, prefix) def supported_dtypes(self): return {d for d in super().supported_dtypes() if d not in dtypes.fp8s or d in amd_fp8s(self.target.arch)} class HIPCCRenderer(HIPRenderer): def __init__(self, target:Target): super().__init__(target, use_hipcc=True) class QCOMCLRenderer(OpenCLRenderer): def __init__(self, target:Target): super().__init__(target) from tinygrad.runtime.support.compiler_qcom import QCOMCompiler self.compiler = QCOMCompiler(target.arch) # QCOM compiler is flaky with half def supported_dtypes(self): return {d for d in Renderer.supported_dtypes(self) if (d != dtypes.float16 or (bool(IMAGE) and bool(FLOAT16))) and d not in dtypes.fp8s+(dtypes.bfloat16,dtypes.double)} # QCOM's load vectorizer emits invalid IR for vectorized bool loads ("Range types must match load type"), type bool buffers as uchar def _render_dtype(self, dtype:DType, sz:int=1, addrspace=AddrSpace.ALU, mutable=True, override_ptr=False, shape=None): if dtype == dtypes.bool and addrspace == AddrSpace.GLOBAL: dtype = dtypes.uint8 return super()._render_dtype(dtype, sz, addrspace, mutable, override_ptr, shape)