mirror of
https://github.com/tinygrad/tinygrad.git
synced 2026-08-29 11:56:08 +00:00
cstyle renderer is new style [pr] (#16484)
* cstyle new style * switch cstyle renderer to new style * fix hip * fixes * fix webgpu * correct webgpu is_packed * fix dsp * fixes * fix Ops.RANGE must be CONST * old style render access * this is correct * fix cstyle to good * dl/dr * as array * fix spec * remove define_local/define_reg * buffer in shrink * fix test_tiny * all tests fix * param args aren't realized * wgsl fix * work * new gate * fix opencl qcom * process replay * sort order * fix render index
This commit is contained in:
George Hotz
@@ -356,7 +356,7 @@ jobs:
|
||||
- name: Train MNIST
|
||||
run: time PYTHONPATH=. DEV=NV TARGET_EVAL_ACC_PCT=96.0 python3 examples/beautiful_mnist.py
|
||||
- name: Run 10 CIFAR training steps
|
||||
run: BENCHMARK_LOG=cifar_10steps ASSERT_MIN_STEP_TIME=120 DEV=NV STEPS=10 python3 examples/hlb_cifar10.py
|
||||
run: BENCHMARK_LOG=cifar_10steps ASSERT_MIN_STEP_TIME=130 DEV=NV STEPS=10 python3 examples/hlb_cifar10.py
|
||||
- name: Run 10 CIFAR training steps w HALF
|
||||
run: BENCHMARK_LOG=cifar_10steps_half ASSERT_MIN_STEP_TIME=120 DEV=NV STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py
|
||||
- name: Run 10 CIFAR training steps w BF16
|
||||
|
||||
@@ -41,7 +41,7 @@ pm_remove_vec_dtypes = PatternMatcher([
|
||||
# no LOADs on register dtypes
|
||||
(UPat(Ops.LOAD, name="x"), lambda x: x.src[0] if x.src[0].addrspace == AddrSpace.REG else None),
|
||||
# remove all vec dtypes
|
||||
(UPat(GroupOp.All-{Ops.PARAM, Ops.DEFINE_LOCAL, Ops.DEFINE_REG}, name="x"),
|
||||
(UPat(GroupOp.All-{Ops.PARAM, Ops.BUFFER, Ops.DEFINE_LOCAL, Ops.DEFINE_REG}, name="x"),
|
||||
lambda x: x.replace(dtype=x.dtype.base.scalar().base)),
|
||||
# replace DEFINE_LOCAL/DEFINE_REG with BUFFER
|
||||
(UPat((Ops.DEFINE_LOCAL, Ops.DEFINE_REG), name="x"), lambda x:
|
||||
|
||||
+73
-32
@@ -10,9 +10,8 @@ from tinygrad.codegen.late.devectorizer import no_vectorized_alu
|
||||
|
||||
|
||||
base_rewrite = PatternMatcher([
|
||||
# defines
|
||||
(UPat(Ops.DEFINE_REG, name="x"), lambda ctx,x: f"{ctx.render_dtype(x.dtype.base)} {ctx[x]}[{x.max_numel()}];"),
|
||||
(UPat(Ops.DEFINE_LOCAL, name="x"), lambda ctx,x: f"{ctx.smem_align}{ctx.smem_prefix}{ctx.render_dtype(x.dtype.base)} {ctx[x]}[{x.max_numel()}];"),
|
||||
# local/reg buffers
|
||||
(UPat(Ops.BUFFER, name="x"), lambda ctx,x: ctx.render_buffer(x)),
|
||||
|
||||
# range/if/endif
|
||||
(UPat(Ops.RANGE, name="x"),
|
||||
@@ -21,11 +20,10 @@ base_rewrite = PatternMatcher([
|
||||
(UPat((Ops.ENDIF, Ops.END)), lambda ctx: "}"),
|
||||
|
||||
# casting
|
||||
(UPat(Ops.CAST, name="x"), lambda ctx,x: f"__builtin_convertvector({ctx[x.src[0]]}, {ctx.render_dtype(x.dtype)})" \
|
||||
if x.max_numel() > 1 and not isinstance(x.dtype, PtrDType) else None),
|
||||
(UPat(Ops.CAST, name="x"), lambda ctx,x: f"__builtin_convertvector({ctx[x.src[0]]}, {ctx.render_type(x)})" \
|
||||
if x.max_numel() > 1 and x.addrspace is AddrSpace.REG else None),
|
||||
(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"__builtin_bit_cast({ctx.render_dtype(x.dtype)}, ({ctx.render_dtype(x.src[0].dtype)})({ctx[x.src[0]]}))"),
|
||||
(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),
|
||||
@@ -47,18 +45,18 @@ base_rewrite = PatternMatcher([
|
||||
# default const render
|
||||
(UPat(Ops.CONST, name="x"), lambda ctx,x: str(x.arg)),
|
||||
|
||||
# movement ops
|
||||
(UPat.var("buf").index(UPat.var('idx')), lambda ctx,buf,idx: f"({ctx[buf]}+{strip_parens(ctx[idx]) if idx.arg == Ops.ADD else ctx[idx]})"),
|
||||
# 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()), name="x"), lambda ctx,**kwargs: ctx.render_index(**kwargs)),
|
||||
(UPat(Ops.STACK, name="x"),
|
||||
lambda ctx,x: f"{ctx.float4.replace('float4', ctx.render_dtype(x.dtype))}" + \
|
||||
f"{ctx.float4_style[0]}{','.join([ctx[y] for y in x.src])}{ctx.float4_style[1]}"),
|
||||
(UPat(Ops.GEP, name="x"), lambda ctx,x: ctx[x.src[0]] + \
|
||||
(f"[{x.arg[0]}]" if x.src[0].max_numel() > ctx.gep_arr_threshold else f".{'xyzwabcd'[x.arg[0]]}")),
|
||||
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[bidx]})"),
|
||||
(UPat(Ops.LOAD, src=(UPat.var("bidx"), UPat.var("var"), UPat.var("gate"))), lambda ctx,bidx,var,gate: f"({ctx[gate]}?*{ctx[bidx]}:{ctx[var]})"),
|
||||
(UPat(Ops.STORE, src=(UPat.var('bidx'), UPat.var("var"))), lambda ctx,bidx,var: f"*{ctx[bidx]} = {ctx[var]};"),
|
||||
(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"__{x.arg[0]}({ctx[x.src[0]]}, {ctx[x.src[1]]}, {ctx[x.src[2]]})"),
|
||||
@@ -106,13 +104,22 @@ pm_manual_bf16_cast = PatternMatcher([
|
||||
(UPat(Ops.CAST, dtype=dtypes.bfloat16, src=(UPat.var("x", dtype=dtypes.float),)), cast_float_to_bf16),
|
||||
])
|
||||
|
||||
def uops_to_dtypes(uops:list[UOp]) -> list[DType]: return dedup(u.dtype for u in uops if not isinstance(u.dtype, (ImageDType, PtrDType)))
|
||||
def uops_to_dtypes(uops:list[UOp]) -> list[DType]:
|
||||
ret = []
|
||||
seen = set()
|
||||
for u in uops:
|
||||
if u.addrspace in (AddrSpace.REG, None) and u.dtype != dtypes.void and u._shape is not None and (key:=(u.dtype, u.max_numel())) not in seen:
|
||||
# TODO: this eventually needs to be removed
|
||||
ret.append(u.dtype.vec(u.max_numel()))
|
||||
seen.add(key)
|
||||
return ret
|
||||
|
||||
# (name, dims, dtype_in, dtype_out, device, threads, upcast_axes, reduce_axes)
|
||||
def wmma_args(uops:list[UOp]):
|
||||
return dedup((uop.arg[0], uop.arg[1], uop.arg[2], uop.dtype.scalar(), *(uop.arg[4:8])) for uop in uops if uop.op is Ops.WMMA)
|
||||
|
||||
class CStyleLanguage(Renderer):
|
||||
new_style = True
|
||||
kernel_typedef: str = "void"
|
||||
buffer_prefix: str = ""
|
||||
buffer_suffix: str = ""
|
||||
@@ -146,8 +153,8 @@ class CStyleLanguage(Renderer):
|
||||
tmp = ""
|
||||
if any(isinstance(u.dtype, ImageDType) for _,(u,_) in bufs):
|
||||
tmp = "const sampler_t smp = CLK_NORMALIZED_COORDS_FALSE | CLK_ADDRESS_CLAMP | CLK_FILTER_NEAREST;\n"
|
||||
buftypes = [(name, self.render_dtype(u.dtype, mutable)+self.buffer_suffix if isinstance(u.dtype, (ImageDType, PtrDType)) else
|
||||
self.arg_int_prefix if u.dtype == dtypes.int else None) for name,(u,mutable) in bufs]
|
||||
buftypes = [(name, self._render_dtype(u.dtype, sz=1, addrspace=u.addrspace, mutable=mutable)+self.buffer_suffix \
|
||||
if u.addrspace == AddrSpace.GLOBAL else self.arg_int_prefix if u.dtype == dtypes.int else None) 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}(",] +
|
||||
@@ -155,16 +162,49 @@ class CStyleLanguage(Renderer):
|
||||
[") {\n" + tmp] + ['\n'.join(kernel), "\n}"])
|
||||
return prg if prefix is None else "\n".join(prefix)+f"\n{prg}"
|
||||
|
||||
def render_cast(self, u:UOp, val:str) -> str: return f"({self.render_dtype(u.dtype)})({val})"
|
||||
def render_index(self, x:UOp, buf:UOp, idx:UOp):
|
||||
if buf.addrspace == AddrSpace.REG and buf.op not in {Ops.AFTER, Ops.BUFFER}:
|
||||
# this is lane access in C
|
||||
assert idx.op is Ops.CONST, f"{idx.op} must be CONST"
|
||||
return self[buf]+(f"[{idx.arg}]" if buf.max_numel() > self.gep_arr_threshold else f".{'xyzwabcd'[idx.arg]}")
|
||||
ptr = f"({self[buf]}+{strip_parens(self[idx]) if idx.arg == Ops.ADD else self[idx]})"
|
||||
if buf.addrspace != AddrSpace.REG: return ptr
|
||||
# REG buffers have no LOAD, so the access is rendered at the INDEX. the cast handles vector access, same as render_access
|
||||
return f"(*(({self.render_type(x)}*)({ptr})))" if x.max_numel() > 1 else f"(*{ptr})"
|
||||
|
||||
def render_buffer(self, x:UOp):
|
||||
shp = x.src[0].as_shape
|
||||
lanes = 1
|
||||
prefix = f"{self.smem_align}{self.smem_prefix}" if x.addrspace == AddrSpace.LOCAL else ""
|
||||
suffix = f"[{shp[0]}]" if len(shp) else ""
|
||||
if len(shp) > 1:
|
||||
# for DEFINE_REG, if it's a 2-D shape it's the number of lanes
|
||||
assert isinstance(shp[1], int)
|
||||
lanes = shp[1]
|
||||
return f"{prefix}{self._render_dtype(x.dtype, sz=lanes)} {self[x]}{suffix};"
|
||||
|
||||
def _render_dtype(self, dtype:DType, sz:int=1, addrspace=AddrSpace.REG, mutable=True):
|
||||
if isinstance(dtype, ImageDType): 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
|
||||
suffix = "*"
|
||||
if sz > 1:
|
||||
return prefix + self.type_map.get(scalar:=dtype.scalar(), scalar.name).replace(" ", "_") + str(sz) + suffix
|
||||
return prefix + self.type_map.get(scalar:=dtype.scalar(), scalar.name) + suffix
|
||||
|
||||
def render_type(self, u:UOp): return self._render_dtype(u.dtype, u.max_numel(), u.addrspace)
|
||||
def render_access(self, u:UOp):
|
||||
if u.addrspace in (AddrSpace.GLOBAL, AddrSpace.LOCAL):
|
||||
if u.max_numel() > 1: return f"*(({self.render_type(u)})({self[u]}))"
|
||||
else: return f"*{self[u]}"
|
||||
return self[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:
|
||||
if isinstance(dt, ImageDType): return f"{'write_only' if mutable else 'read_only'} image2d_t"
|
||||
if isinstance(dt, PtrDType):
|
||||
prefix = ""
|
||||
if dt.addrspace == AddrSpace.LOCAL and self.smem_prefix_for_cast: prefix = self.smem_prefix
|
||||
if dt.addrspace == AddrSpace.GLOBAL: prefix = self.buffer_prefix
|
||||
return prefix + self.render_dtype(dt.base) + "*"
|
||||
if dt.count > 1: return self.type_map.get(scalar:=dt.scalar(), scalar.name).replace(" ", "_") + str(dt.count)
|
||||
return self.type_map.get(scalar:=dt.scalar(), scalar.name)
|
||||
return self._render_dtype(dt, dt.count, dt.addrspace if isinstance(dt, PtrDType) else 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]]]]:
|
||||
@@ -181,6 +221,7 @@ class CStyleLanguage(Renderer):
|
||||
name = "test"
|
||||
for u in uops:
|
||||
if u.op in {Ops.NOOP, Ops.GROUP}: continue
|
||||
if u.op == Ops.STACK and len(u.src) == 0: continue
|
||||
if u.op is Ops.AFTER:
|
||||
r[u] = r[u.src[0]]
|
||||
continue
|
||||
@@ -199,7 +240,7 @@ class CStyleLanguage(Renderer):
|
||||
if u.op is Ops.SPECIAL: r[u] = u.arg
|
||||
elif u.op is Ops.RANGE: r[u] = f"{axis_letters[u.arg[-1]]}idx"+range_str(u)
|
||||
else:
|
||||
prefix = {Ops.WMMA: "wmma", Ops.DEFINE_LOCAL: "temp", Ops.CONST: "const",
|
||||
prefix = {Ops.WMMA: "wmma", Ops.DEFINE_LOCAL: "temp", Ops.CONST: "const", Ops.BUFFER: "buf",
|
||||
Ops.CAST: "cast", Ops.BITCAST: "cast", Ops.GEP: "gep", Ops.STACK: "cast",
|
||||
Ops.INDEX: "bidx", Ops.DEFINE_REG: "acc", Ops.LOAD: "val"}.get(u.op, "alu")
|
||||
r[u] = f"{prefix}{c[prefix]}"
|
||||
@@ -208,14 +249,14 @@ class CStyleLanguage(Renderer):
|
||||
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}: depth -= 1
|
||||
if (u.op is not Ops.CAST or u.dtype.vcount == 1) and (u.op in {Ops.CONST, Ops.GEP, Ops.INDEX, Ops.CUSTOMI} or \
|
||||
if (u.op is not Ops.CAST or u.dtype.vcount == 1) and (u.op in {Ops.CONST, Ops.GEP, Ops.INDEX, Ops.SHRINK, Ops.CUSTOMI} or \
|
||||
(u.op is Ops.LOAD and u.src[0].addrspace == AddrSpace.REG) or \
|
||||
(u.op is Ops.CAST 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.DEFINE_LOCAL, Ops.STORE, Ops.DEFINE_REG} and u.dtype != dtypes.void:
|
||||
l = f"{self.render_dtype(u.dtype)} {r[u]} = {l}" + (";" if u.op is not Ops.SPECIAL else "")
|
||||
if u.op not in {Ops.RANGE, Ops.DEFINE_LOCAL, Ops.STORE, Ops.DEFINE_REG, Ops.BUFFER} and u.dtype != dtypes.void:
|
||||
l = f"{self.render_type(u)} {r[u]} = {l}" + (";" if u.op is not Ops.SPECIAL else "")
|
||||
kernel.append(" "*depth + l)
|
||||
if prefix: c[prefix] += 1 # if it was used, increment
|
||||
if u.op in {Ops.IF, Ops.RANGE}: depth += 1
|
||||
|
||||
@@ -269,6 +269,7 @@ _nload_img = nir_instr(intrins=lambda dtype:{'IMAGE_DIM':mesa.GLSL_SAMPLER_DIM_2
|
||||
lambda b,img,idx_y,idx_x,dtype: mesa.nir_intrinsic_instr_create(b.shader, g("nir_intrinsic_image_load")))
|
||||
|
||||
class IR3Renderer(NIRRenderer, OpenCLRenderer):
|
||||
new_style = False
|
||||
has_aux = True
|
||||
|
||||
def nload_img(ctx,img,idx_y,idx_x):
|
||||
|
||||
+21
-23
@@ -1,4 +1,4 @@
|
||||
from tinygrad.dtype import DType, PtrDType, dtypes, truncate, AddrSpace
|
||||
from tinygrad.dtype import DType, dtypes, truncate, AddrSpace
|
||||
from tinygrad.uop.ops import UOp, Ops, PatternMatcher, UPat
|
||||
from tinygrad.renderer.cstyle import CStyleLanguage, base_rewrite, extra_pm
|
||||
from tinygrad.helpers import strip_parens
|
||||
@@ -27,13 +27,12 @@ def packed_load(root:UOp, bidx:UOp, dtype:DType, var:UOp|None=None, gate:UOp|Non
|
||||
val = (load.cast(dtypes.uint32) >> shift_am) & mask
|
||||
return sign_extend(val, 8*dtype.itemsize).cast(dtype) if dtype in [dtypes.char, dtypes.short] else val.cast(dtype)
|
||||
|
||||
def is_packed(dt:DType, odt:DType|None = None) -> bool:
|
||||
if odt is None: odt = dt
|
||||
# registers aren't packed
|
||||
if isinstance(odt, PtrDType) and odt.addrspace == AddrSpace.REG: return False
|
||||
return dt.itemsize < 4 and dt.base != dtypes.half
|
||||
def _packed_size(dt:PtrDType): return dt.size // (4//dt.itemsize) if is_packed(dt) else dt.size
|
||||
|
||||
def is_packed(x:UOp):
|
||||
if x.op is Ops.LOAD: dt, addrspace = x.dtype, x.src[0].addrspace
|
||||
elif x.op is Ops.STORE: dt, addrspace = x.src[1].dtype, x.src[0].addrspace
|
||||
else: dt, addrspace = x.dtype, x.addrspace
|
||||
return dt.itemsize < 4 and dt != dtypes.half and addrspace != AddrSpace.REG
|
||||
def _packed_size(u:UOp): return u.max_numel() // (4//u.dtype.itemsize) if is_packed(u) else u.max_numel()
|
||||
def is_nan(a):
|
||||
bs, (exp, mant) = a.dtype.bitsize, dtypes.finfo(a.dtype)
|
||||
return (a.bitcast(getattr(dtypes, f"uint{bs}")) & ((1 << (bs - 1)) - 1)) > (((1 << exp) - 1) << mant)
|
||||
@@ -43,12 +42,11 @@ wgsl_matcher = PatternMatcher([
|
||||
lambda a,b,c: a.cast(dtypes.int).alu(c.op, b.cast(dtypes.int)).cast(dtypes.bool)),
|
||||
# TODO: load alt value doesnt have to be a const
|
||||
(UPat.load(UPat.var("b"), UPat.cvar("c"), UPat.var("gate"), name="l"),
|
||||
lambda l,b,c,gate: packed_load(l,b,l.dtype,c.cast(dtypes.uint32),gate) if is_packed(l.dtype, b.dtype) else None),
|
||||
(UPat.load(UPat.var("b"), name='l'), lambda l,b: packed_load(l, b, l.dtype) if is_packed(l.dtype, b.dtype) else None),
|
||||
(UPat.store(UPat.var("bidx"), UPat.var("var"), UPat.var("gate")),
|
||||
lambda bidx,var,gate: packed_store(bidx,var,gate) if is_packed(var.dtype, bidx.dtype) else None),
|
||||
(UPat.store(UPat.var("bidx"), UPat.var("var")),
|
||||
lambda bidx,var: packed_store(bidx,var) if is_packed(var.dtype, bidx.dtype) else None),
|
||||
lambda l,b,c,gate: packed_load(l,b,l.dtype,c.cast(dtypes.uint32),gate) if is_packed(l) else None),
|
||||
(UPat.load(UPat.var("b"), name='l'), lambda l,b: packed_load(l,b,l.dtype) if is_packed(l) else None),
|
||||
(UPat.store(UPat.var("b"), UPat.var("var"), UPat.var("gate"), name="s"),
|
||||
lambda b,var,gate,s: packed_store(b,var,gate) if is_packed(s) else None),
|
||||
(UPat.store(UPat.var("b"), UPat.var("var"), name="s"), lambda b,var,s: packed_store(b,var) if is_packed(s) else None),
|
||||
(UPat.var("a") << UPat.var("b"),lambda a,b:(a.bitcast(dtypes.uint32)<<b.cast(dtypes.uint32)).bitcast(a.dtype) if b.dtype!=dtypes.uint32 else None),
|
||||
(UPat.var("x") >> UPat.var("y"), lambda x,y: UOp(Ops.SHR, x.dtype, (x,y.cast(dtypes.uint))) if y.dtype != dtypes.uint else None),
|
||||
# fix nan check: 'a != a -> is_nan()'
|
||||
@@ -73,8 +71,8 @@ class WGSLRenderer(CStyleLanguage):
|
||||
(UPat(Ops.CONST, dtype=(dtypes.uchar, dtypes.ushort, dtypes.uint32), name="x"),
|
||||
lambda x: f"bitcast<u32>({x.arg})" if x.arg < 0 else f"{x.arg&0xFFFFFFFF}u"),
|
||||
(UPat(Ops.CONST, dtype=dtypes.int32, name="x"), lambda ctx,x: f"{truncate[x.dtype](x.arg)}"),
|
||||
(UPat(Ops.DEFINE_LOCAL, name="x"), lambda ctx,x: f"var<workgroup> {ctx[x]}: array<{ctx.buf_map(x.dtype.base)},{_packed_size(x.dtype)}>;"),
|
||||
(UPat(Ops.DEFINE_REG, name="x"), lambda ctx,x: f"var {ctx[x]}: array<{ctx.buf_map(x.dtype)},{_packed_size(x.dtype)}>;"),
|
||||
(UPat(Ops.BUFFER, name="x"), lambda ctx,x:
|
||||
f"var{'<workgroup>' if x.addrspace == AddrSpace.LOCAL else ''} {ctx[x]}: array<{ctx.buf_map(x)},{_packed_size(x)}>;"),
|
||||
(UPat(Ops.BITCAST, dtype=dtypes.half, name="x", src=(UPat(dtype=(dtypes.short, dtypes.ushort, dtypes.uint32),),)),
|
||||
lambda ctx,x: f"bitcast<vec2<f16>>({ctx[x.src[0]]})[0]"),
|
||||
(UPat(Ops.BITCAST, dtype=dtypes.uchar, name="x"), lambda ctx,x: f"bitcast<u32>({ctx[x.src[0]]}&0xFF)"),
|
||||
@@ -86,20 +84,20 @@ class WGSLRenderer(CStyleLanguage):
|
||||
(UPat(Ops.BITCAST, name="x"), lambda ctx,x: f"bitcast<{ctx.type_map[x.dtype]}>({ctx[x.src[0]]})"),
|
||||
# TODO: load alt value doesnt have to be a const
|
||||
(UPat.load(UPat.var("b"), UPat.cvar("v"), UPat.var("gate")),
|
||||
lambda ctx,b,v,gate: f"select({ctx[v]}, {ctx.render_load(ctx[b],b.src[0].dtype)}, {ctx[gate]})"),
|
||||
(UPat.load(UPat.var("b")), lambda ctx, b: ctx.render_load(ctx[b], b.dtype)),
|
||||
lambda ctx,b,v,gate: f"select({ctx[v]}, {ctx.render_load(ctx[b], b.src[0])}, {ctx[gate]})"),
|
||||
(UPat.load(UPat.var("b")), lambda ctx, b: ctx.render_load(ctx[b], b)),
|
||||
(UPat.store(UPat.var("b"), UPat.var("v")), lambda ctx,b,v:\
|
||||
# (load & mask) | var -> mask = v.src[0].src[1], var = v.src[1]
|
||||
f"atomicAnd(&{ctx[b]},{ctx[v.src[0].src[1]]});\n atomicAdd(&{ctx[b]},{ctx[v.src[1]]});" if is_packed(b.src[0].dtype) \
|
||||
f"atomicAnd(&{ctx[b]},{ctx[v.src[0].src[1]]});\n atomicAdd(&{ctx[b]},{ctx[v.src[1]]});" if is_packed(b) \
|
||||
else f"{ctx[b]} = {ctx[v]};"),
|
||||
(UPat(Ops.INDEX, src=(UPat.var("b"), UPat.var("idx"))),
|
||||
lambda ctx,b,idx: f"{ctx[b]}[{strip_parens(ctx[idx]) if idx.arg is Ops.ADD else ctx[idx]}]"),
|
||||
]) + base_rewrite
|
||||
|
||||
def render_cast(self, u:UOp, val: str) -> str: return f"{self.type_map[u.dtype]}({val})"
|
||||
def render_dtype(self, dt:DType, mutable=True) -> str: return "var"
|
||||
def render_load(self, x:str, dt:DType) -> str: return f"atomicLoad(&{x})" if is_packed(dt) else x
|
||||
def buf_map(self, dt:DType) -> str: return "atomic<u32>" if is_packed(dt) else self.type_map[dt.base]
|
||||
def _render_dtype(self, dtype:DType, sz:int=1, addrspace=AddrSpace.REG, mutable=True): return "var"
|
||||
def render_load(self, x:str, u:UOp) -> str: return f"atomicLoad(&{x})" if is_packed(u) else x
|
||||
def buf_map(self, u:UOp) -> str: return "atomic<u32>" if is_packed(u) else self.type_map[u.dtype.base]
|
||||
def render_kernel(self, function_name:str, kernel:list[str], bufs:list[tuple[str,tuple[UOp,bool]]], uops:list[UOp], prefix=None) -> str:
|
||||
local_size = [u.src[0].ssimplify() for u in sorted([u for u in uops if u.op is Ops.SPECIAL and u.arg[0] == 'l'], key=lambda u: u.arg)]
|
||||
if not local_size: local_size = [1]
|
||||
@@ -111,7 +109,7 @@ class WGSLRenderer(CStyleLanguage):
|
||||
prg += "@group(0) @binding(0)\nvar<uniform> INFINITY : f32;\n"
|
||||
prg += "\n".join((external_local_bufs or [])+[f"@group(0) @binding({next(bind_it)+1})" +
|
||||
f"{'var<storage,read_write>' if u.addrspace == AddrSpace.GLOBAL else 'var<uniform>'}" +
|
||||
f"{name}:{f'array<{self.buf_map(u.dtype.base)}>' if u.addrspace == AddrSpace.GLOBAL else self.buf_map(u.dtype)};" for name,(u,_) in bufs])
|
||||
f"{name}:{f'array<{self.buf_map(u)}>' if u.addrspace == AddrSpace.GLOBAL else self.buf_map(u)};" for name,(u,_) in bufs])
|
||||
prg += f"\n@compute @workgroup_size({','.join([str(x) for x in local_size])}) fn {function_name}(@builtin(workgroup_id) gindex: vec3<u32>,"
|
||||
return prg + "@builtin(local_invocation_id) lindex: vec3<u32>) {\n" + "\n".join(kernel) + "\n}"
|
||||
|
||||
|
||||
@@ -22,6 +22,9 @@ class Ops(FastEnum):
|
||||
# define LOCAL/REG allocate things
|
||||
DEFINE_LOCAL = auto(); DEFINE_REG = auto()
|
||||
|
||||
# BUFFER is the new LOCAL/REG
|
||||
BUFFER = auto()
|
||||
|
||||
# ** 2 -- non op uops **
|
||||
|
||||
# uops that aren't rendered
|
||||
@@ -49,7 +52,7 @@ class Ops(FastEnum):
|
||||
# ** 3 -- load/store **
|
||||
|
||||
# INDEX is a BinaryOp similar to ADD, but it operates on pointers
|
||||
INDEX = auto()
|
||||
INDEX = auto(); SHRINK = auto()
|
||||
|
||||
# load/store before math
|
||||
LOAD = auto(); STORE = auto()
|
||||
@@ -99,10 +102,10 @@ class Ops(FastEnum):
|
||||
CONTIGUOUS = auto(); CONTIGUOUS_BACKWARD = auto(); DETACH = auto()
|
||||
|
||||
# buffer ops
|
||||
STAGE = auto(); COPY = auto(); BUFFER = auto(); SLICE = auto(); MSELECT = auto(); MSTACK = auto(); CUSTOM_FUNCTION = auto()
|
||||
STAGE = auto(); COPY = auto(); SLICE = auto(); MSELECT = auto(); MSTACK = auto(); CUSTOM_FUNCTION = auto()
|
||||
|
||||
# the core 6 movement ops! these only exist in the tensor graph
|
||||
RESHAPE = auto(); PERMUTE = auto(); EXPAND = auto(); PAD = auto(); SHRINK = auto(); FLIP = auto()
|
||||
RESHAPE = auto(); PERMUTE = auto(); EXPAND = auto(); PAD = auto(); FLIP = auto()
|
||||
MULTI = auto() # MULTI is really a movement op
|
||||
|
||||
# reduce
|
||||
|
||||
@@ -51,8 +51,8 @@ z3_renderer = PatternMatcher([
|
||||
])
|
||||
|
||||
def uops_to_z3(solver:z3.Solver, *uops: UOp) -> list[z3.ExprRef]:
|
||||
# gate on any upstream INDEX as a replacement for PtrDType
|
||||
lst = list(UOp.sink(*uops).toposort(gate=lambda x: x.op is not Ops.INDEX and \
|
||||
# gate on upstream AFTER/BUFFER as a replacement for PtrDType, but keep INDEX as an unknown LOAD
|
||||
lst = list(UOp.sink(*uops).toposort(gate=lambda x: x.op not in {Ops.AFTER, Ops.BUFFER, Ops.PARAM} and \
|
||||
(x.dtype.scalar() in dtypes.ints+(dtypes.bool, dtypes.weakint) or x.op is Ops.SINK)))[:-1]
|
||||
z3map: dict[UOp, z3.ExprRef] = {}
|
||||
for u in lst:
|
||||
|
||||
Reference in New Issue
Block a user