diff --git a/tinygrad/renderer/nir.py b/tinygrad/renderer/nir.py index a2153394f6..24e0b616bc 100644 --- a/tinygrad/renderer/nir.py +++ b/tinygrad/renderer/nir.py @@ -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: