sirhcm/tinygrad
1
0
Fork 0
forked from tinygrad/tinygrad
Code Pull Requests 2 Actions 8 Activity

casted CONST migration for nir [pr] (#17588)

Browse Source
This commit is contained in:
chenyu
2026-08-18 23:04:07 -04:00
committed by GitHub
parent df3b114fbc
commit ca86a42703
1 changed files with 6 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tinygrad/renderer/nir.py
+6 -4
View File
@@ -116,6 +116,7 @@ def nidx(b:mesa.nir_builder, buf, off, space, itemsize, gate=None) -> mesa.nir_d
class NIRRenderer(Renderer):
suffix = "NIR"
casted_consts = True
nir_options: bytes
global_max, local_max, shared_max = CUDARenderer.global_max, CUDARenderer.local_max, CUDARenderer.shared_max
code_for_op = {**{k:lambda:None for k in u_aop.keys()}, **{k:lambda:None for k in s_aop.keys()}, **{k:lambda:None for k in f_aop.keys()}}
@@ -145,7 +146,7 @@ class NIRRenderer(Renderer):
])
def_rewrite = PatternMatcher([
(UPat(Ops.CONST, name="x"), lambda ctx,x: nimm(ctx.b, x.val, x.dtype)),
(UPat.cvar("c").cast(name="x"), lambda ctx,x,c: nimm(ctx.b, c.val, x.dtype)),
(UPat(Ops.PARAM, name="x"), lambda ctx,x: ctx.param(ctx.b, x, x.dtype.itemsize if x.addrspace is AddrSpace.ALU else 8)),
(UPat(Ops.SPECIAL, name="x"), lambda ctx,x: nchannel(ctx.b, {'g':ngid, 'l':nlid, 'i': nid}[x.arg[0]](ctx.b), int(x.arg[-1]))),
(UPat(Ops.STORE, src=(UPat((Ops.INDEX, Ops.SHRINK), src=(UPat.var("buf"),UPat.var("off")), allow_any_len=True), UPat.var("val"))),
@@ -186,16 +187,17 @@ class NIRRenderer(Renderer):
def render(self, uops:list[UOp]):
self.prerender(uops)
for u in [u for u in uops if u.op is Ops.SPECIAL and u.arg[0] == "l"]: self.b.shader.contents.info.workgroup_size[int(u.arg[-1])] = u.src[0].val
for u in [u for u in uops if u.op is Ops.SPECIAL and u.arg[0] == "l"]:
self.b.shader.contents.info.workgroup_size[int(u.arg[-1])] = u.src[0].src[0].val
self.r: dict[UOp, Any] = {}
self.param_idx = 0
ranges: list[mesa.nir_def|None] = []
for u in uops:
if u.op in {Ops.NOOP, Ops.GROUP} or (u.op is Ops.STACK and len(u.src) == 0): pass
if u.op in {Ops.NOOP, Ops.GROUP, Ops.CONST} or (u.op is Ops.STACK and len(u.src) == 0): pass
elif u.op in {Ops.INDEX, Ops.SHRINK}:
# INDEX on a register value picks the element, memory INDEX is handled in the LOAD/STORE patterns
if u.src[0].op not in {Ops.PARAM, Ops.BUFFER, Ops.AFTER}: self.r[u] = nchannel(self.b, self.r[u.src[0]], u.src[1].val)
if u.src[0].op not in {Ops.PARAM, Ops.BUFFER, Ops.AFTER}: self.r[u] = nchannel(self.b, self.r[u.src[0]], u.src[1].src[0].val)
elif u.op is Ops.AFTER:
self.r[u] = self.r[u.src[0]]
elif u.op == Ops.SINK: