update for ANON

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"

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 ")")