From 3dbaa526fae8171e7b861b63e112c9079c8d3f47 Mon Sep 17 00:00:00 2001 From: George Hotz Date: Mon, 1 Jun 2026 17:49:12 -0700 Subject: [PATCH] update for ANON --- tinygrad/renderer/cstyle.py | 5 ++--- tinygrad/renderer/llvmir.py | 4 ++-- 2 files changed, 4 insertions(+), 5 deletions(-) diff --git a/tinygrad/renderer/cstyle.py b/tinygrad/renderer/cstyle.py index bd0d91b48e..df2dce7203 100644 --- a/tinygrad/renderer/cstyle.py +++ b/tinygrad/renderer/cstyle.py @@ -11,8 +11,8 @@ from tinygrad.codegen.late.devectorizer import no_vectorized_alu def render_index(ctx,buf,idx): base = buf while base.op is Ops.AFTER: base = base.src[0] - if base.addrspace == AddrSpace.REG or base.op not in {Ops.PARAM, Ops.DEFINE_LOCAL}: - assert idx.op is Ops.CONST + if base.addrspace == AddrSpace.ANON: + assert idx.op is Ops.CONST, f"{idx.op} must be CONST" return f"{ctx[buf]}[{idx.arg}]" else: return f"({ctx[buf]}+{strip_parens(ctx[idx]) if idx.arg == Ops.ADD else ctx[idx]})" @@ -161,7 +161,6 @@ class CStyleLanguage(Renderer): def render_cast(self, dt:DType, val: str) -> str: return f"({self.render_dtype(dt)})({val})" def render_dtype_with_shape(self, u:UOp) -> DType: return dtype_with_shape(u.dtype, u.shape) def render_access(self, bidx:UOp, dtype:DType) -> str: - if bidx.addrspace == AddrSpace.REG: return self[bidx] return f"(*(({self.render_dtype(dtype.ptr(addrspace=bidx.addrspace))})({self[bidx]})))" if dtype.count > 1 else f"(*{self[bidx]})" def render_dtype(self, dt:DType, mutable=True) -> str: if isinstance(dt, ImageDType): return f"{'write_only' if mutable else 'read_only'} image2d_t" diff --git a/tinygrad/renderer/llvmir.py b/tinygrad/renderer/llvmir.py index 52022f1262..b480ac6be1 100644 --- a/tinygrad/renderer/llvmir.py +++ b/tinygrad/renderer/llvmir.py @@ -56,12 +56,12 @@ def render_wmma_amd(ctx, wmma: UOp, cdna=False) -> str: N,M,K = wmma.arg[1] if cdna: if K == 32: dt_map.update({dtypes.half: ".f16", dtypes.bfloat16: ".bf16"}) - return f" {ctx[wmma]} = call {ldt(wmma.dtype)} @llvm.amdgcn.mfma.{dt_map[wmma.src[-1].dtype.scalar()]}" + \ + return f" {ctx[wmma]} = call {ldt(wmma.dtype, count=wmma.max_numel())} @llvm.amdgcn.mfma.{dt_map[wmma.src[-1].dtype.scalar()]}" + \ f".{N}x{M}x{K}{dt_map[wmma.arg[2]]}(" + ", ".join([f"{ldt(w.dtype, count=w.max_numel())} {ctx[w]}" for w in wmma.src]) + \ ", i32 0, i32 0, i32 0)" # https://github.com/llvm/llvm-project/blob/main/llvm/test/CodeGen/AMDGPU/GlobalISel/llvm.amdgcn.wmma_32.ll # example: %wmma0 = call <8 x float> @llvm.amdgcn.wmma.f32.16x16x16.f16(<16 x half> %v99,<16 x half> %v100,<8 x float> %v101) - return f" {ctx[wmma]} = call {ldt(wmma.dtype)} @llvm.amdgcn.wmma.{dt_map[wmma.src[-1].dtype.scalar()]}.16x16x16." + \ + return f" {ctx[wmma]} = call {ldt(wmma.dtype, count=wmma.max_numel())} @llvm.amdgcn.wmma.{dt_map[wmma.src[-1].dtype.scalar()]}.16x16x16." + \ f"{dt_map[wmma.src[0].dtype.scalar()]}(" + ", ".join([f"{ldt(w.dtype, count=w.max_numel())} {ctx[w]}" for w in wmma.src]) + (", i1 false)" \ if wmma.dtype.scalar() != dtypes.float else ")")