stack is fine

test/backend/test_linearizer.py

@@ -251,9 +251,9 @@ class TestLinearizer(unittest.TestCase):
     for u in uops:
       if u.op is Ops.STORE and u.src[0].addrspace is AddrSpace.REG:
         if uops.index(u) < begin_range:
-          assert u.src[1].op is Ops.CONST
+          assert u.src[1].op is Ops.CONST or (u.src[1].op is Ops.STACK and all(x.op is Ops.CONST for x in u.src[1].src))
         else:
-          assert u.src[1].op in GroupOp.ALU
+          assert u.src[1].op is Ops.CONST or (u.src[1].op is Ops.STACK and all(x.op in GroupOp.ALU for x in u.src[1].src))
           assert begin_range < uops.index(u) < end_range
       # children of END are placed after ENDRANGE
       if any(x.op is Ops.END and x.src[1].op in GroupOp.ALU for x in u.src):

tinygrad/codegen/late/coalese.py

@@ -1,7 +1,7 @@
 from typing import Any
 import itertools, functools
 from collections import defaultdict
-from tinygrad.dtype import dtypes, AddrSpace, Invalid, ImageDType, DType
+from tinygrad.dtype import dtypes, Invalid, ImageDType, DType
 from tinygrad.uop.ops import UOp, Ops, PatternMatcher, UPat, GroupOp
 from tinygrad.uop.symbolic import uop_given_valid, parse_valid, invalid_gate
 from tinygrad.helpers import getenv, IMAGE, OSX, ceildiv
@@ -109,7 +109,6 @@ def memory_coalesing(sink:UOp, ctx:Renderer) -> UOp:
       assert len(u.src) == (2 if u.op is Ops.STORE else 1), "memory coalesing does not support gated loads/stores"
       assert u.src[0].op is Ops.INDEX, f"memory coalesing should be on INDEX, not {u.src[0].op}"
       buf, idx_u = u.src[0].src
-      if buf.addrspace == AddrSpace.REG: continue
       idx: Any = idx_u.src[1] if idx_u.op is Ops.WHERE and idx_u.src[2].arg is Invalid else idx_u
       valid: Any = idx_u.src[0] if idx_u.op is Ops.WHERE and idx_u.src[2].arg is Invalid else None
       if idx.op is Ops.ADD and idx.src[1].op is Ops.CONST: root_src, arg = idx.src[0], idx.src[1].arg
@@ -130,8 +129,6 @@ def memory_coalesing(sink:UOp, ctx:Renderer) -> UOp:
       must_divide = False
     elif buf.dtype not in (dtypes.float, dtypes.half, *dtypes.fp8s) and not isinstance(buf.dtype, ImageDType):
       pass
-    elif buf.addrspace == AddrSpace.REG:
-      pass
     elif isinstance(buf.dtype, ImageDType):
       lengths = [4]
     elif ctx is not None and ctx.supports_float4: