[pull] master from tinygrad:master

docs/env_vars.md

@@ -57,6 +57,8 @@ AMD:LLVM      | use the AMD device with the LLVM renderer
 NV:CUDA:sm_70 | use the NV device with the CUDA renderer targetting sm_70
 AMD::gfx950   | use the AMD device targetting gfx950
 USB+AMD       | use the AMD device over the USB interface
+CPU:LLVM      | use the CPU device with the LLVM renderer
+CPU:LLVM:x86_64,znver2,avx2,-avx512f | use the CPU device with the LLVM renderer, with [additional arch flags](runtime.md#cpu-arch)
 
 ### Debug breakdown
 

docs/runtime.md

@@ -10,7 +10,7 @@ tinygrad supports various runtimes, enabling your code to scale across a wide ra
 | [METAL](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_metal.py) | Utilizes Metal for acceleration on Apple devices | - | M1+ Macs; Metal 3.0+ for `bfloat` support |
 | [CUDA](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_cuda.py) | Utilizes CUDA for acceleration on NVIDIA GPUs | nvrtc (default)<br> PTX (`DEV=CUDA:PTX`) | NVIDIA GPU with CUDA support |
 | [CL](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_cl.py) | Accelerates computations using OpenCL on GPUs | - | OpenCL 2.0 compatible device |
-| [CPU](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_cpu.py) | Runs on CPU using the clang or llvm compiler | Clang JIT (default)<br>LLVM IR (`DEV=CPU:LLVM`) | `clang` compiler in system `PATH` |
+| [CPU](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_cpu.py) | Runs on CPU using the clang or llvm compiler | Clang JIT (default)<br>LLVM IR (`DEV=CPU:LLVM`) | `clang` compiler in system `PATH`<br>You can specify additional arch parameters via [the `DEV` variable](env_vars.md#dev-variable). See [CPU arch](#cpu-arch) for details. |
 | [WEBGPU](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_webgpu.py) | Runs on GPU using the Dawn WebGPU engine (used in Google Chrome) | - | Dawn library installed and discoverable. Binaries: [pydawn v0.3.0](https://github.com/wpmed92/pydawn/releases/tag/v0.3.0) |
 
 
@@ -79,3 +79,13 @@ NV backend supports several interfaces for communicating with devices:
 
 * `NVK`: uses the nvidia driver
 * `PCI`: uses the [NV driver](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/support/nv/nvdev.py)
+
+## CPU Arch
+The CPU renderers may be additionally configured using the arch component of [the `DEV` environment variable](env_vars.md#dev-variable).
+CPU arch should be specified as a comma-separated list of parameters, and must contain at least two values: the architecture family (ie. x86_64, arm64, or riscv64) and the cpu type (as accepted by `clang`'s `-march`).
+If native is specified as the cpu type, tinygrad (or delegate compiler) will query the host cpu type. Additional comma-separated values may be specified as follows:
+
+* `AMX`: emit Apple silicon AMX instructions
+
+All other additional values are interpreted as cpu feature flags. When a value is preceded by a `-` character, the corresponding feature flag will be disabled, otherwise the flag will be enabled.
+Note that enabled feature flags should not be preceded by a `+`.

extra/optimization/extract_dataset.py

@@ -10,4 +10,4 @@ def extract_ast(*args) -> None:
   return None
 
 if __name__ == "__main__":
-  _pmap({"get_program":extract_ast})
+  _pmap({"do_to_program":extract_ast})

test/external/process_replay/process_replay.py

@@ -8,8 +8,8 @@
 if not int(os.getenv("ASSERT_PROCESS_REPLAY", "1")): ASSERT_DIFF = 0
 
 try:
-  from tinygrad.renderer import Renderer, ProgramSpec
-  from tinygrad.engine.realize import get_program
+  from tinygrad.renderer import Renderer
+  from tinygrad.codegen import to_program
   from tinygrad.uop.ops import UOp, Ops
   from tinygrad.helpers import VERSION, Context, ContextVar, colored, db_connection, getenv, tqdm
 except ImportError as e:
@@ -41,23 +41,25 @@ class ProcessReplayWarning(Warning): pass
 
 # *** replay the function and convert return values to string
 
-def replay_get_program(p:ProgramSpec, ast:UOp, renderer:Renderer) -> tuple[str, str, tuple[Any, ...]]:
+def replay_to_program(p:UOp, ast:UOp, renderer:Renderer) -> tuple[str, str, tuple[Any, ...]]:
   if ast.op is Ops.PROGRAM: input_ast = ast
   else:
     sink_arg = ast.arg
-    if sink_arg.beam: sink_arg = replace(sink_arg, opts_to_apply=p.applied_opts)
-    input_ast = ast.replace(arg=replace(sink_arg, name=p.name))
-  p2 = get_program(input_ast, renderer=renderer)
-  def to_str(ret:ProgramSpec) -> str:
+    if sink_arg.beam: sink_arg = replace(sink_arg, opts_to_apply=p.src[0].arg.applied_opts)
+    input_ast = ast.replace(arg=replace(sink_arg, name=p.src[0].arg.name))
+  p2 = to_program(input_ast, renderer=renderer)
+  device = p.src[1].arg
+  def to_str(ret:UOp) -> str:
+    src = ret.src[3].arg
     # PYTHON renderer pickles UOps, first unpickle and decode here
-    if p.device.startswith("PYTHON"): return "\n".join([str(x) for x in pickle.loads(base64.b64decode(ret.src))])
-    return ret.src
+    if device.startswith("PYTHON"): return "\n".join([str(x) for x in pickle.loads(base64.b64decode(src))])
+    return src
   # properly color the name arg
   ast_repr = codecs.decode(str(input_ast), "unicode_escape")
   return to_str(p2), to_str(p), (ast_repr, renderer)
 
 replayers: dict[str, Callable[..., tuple[str, str, tuple[Any, ...]]]] = {}
-replayers["get_program"] = replay_get_program
+replayers["do_to_program"] = replay_to_program
 
 # *** run replayers on captured rows and print diffs
 

test/null/test_device.py

@@ -132,6 +132,7 @@ def test_parse(self):
     for d, t in [("AMD", Target(device="AMD", renderer="")), ("AMD:LLVM", Target(device="AMD", renderer="LLVM")),
                  (":LLVM", Target(device="", renderer="LLVM")), ("AMD::gfx1100", Target(device="AMD", arch="gfx1100")),
                  ("AMD:LLVM:gfx1100", Target(device="AMD", renderer="LLVM", arch="gfx1100")), ("::gfx1100", Target(arch="gfx1100")),
+                 ("CPU:LLVM:arm64,native,AMX", Target(device="CPU", renderer="LLVM", arch="arm64,native,AMX")),
                  ("USB+", Target(interface="USB")), ("USB+AMD", Target(device="AMD", interface="USB")),
                  ("PCI:0+AMD", Target(device="AMD", interface="PCI", indices="0")), (":0+AMD", Target(device="AMD", indices="0")),
                  ("PCI:0,1+AMD", Target(device="AMD", interface="PCI", indices="0,1")),

test/null/test_elf.py

@@ -23,7 +23,7 @@ def test_clang_jit_compiler_external_raise(self):
       }
     '''
     with self.assertRaisesRegex(RuntimeError, 'evil_external_function'):
-      ClangJITCompiler({'AMD64':'x86_64', 'aarch64':'arm64'}.get(m:=platform.machine(), m)+",native").compile(src)
+      ClangJITCompiler([{'AMD64':'x86_64', 'aarch64':'arm64'}.get(m:=platform.machine(), m), "native"]).compile(src)
   def test_link(self):
     src = '''
       float powf(float, float); // from libm

test/null/test_process_replay.py

@@ -1,8 +1,8 @@
 import unittest
 from tinygrad import Tensor, Device, Context
-from tinygrad.engine.realize import get_program
+from tinygrad.codegen import do_to_program
 from tinygrad.codegen.opt import Opt, OptOps
-from test.external.process_replay.process_replay import replay_get_program
+from test.external.process_replay.process_replay import replay_to_program
 from test.helpers import replace_opts
 
 N = 16
@@ -14,30 +14,30 @@ def setUpClass(cls):
 
   def test_replay_no_opts(self):
     # opts=None means use default heuristic path
-    p = get_program(self.ast, self.renderer)
-    good, compare, _ = replay_get_program(p, self.ast, self.renderer)
+    p = do_to_program(self.ast, self.renderer)
+    good, compare, _ = replay_to_program(p, self.ast, self.renderer)
     self.assertEqual(good, compare)
 
   def test_replay_empty_opts(self):
     # opts=[] means explicitly apply zero opts (unoptimized)
     ast = replace_opts(self.ast, [])
-    p = get_program(ast, self.renderer)
-    good, compare, _ = replay_get_program(p, ast, self.renderer)
+    p = do_to_program(ast, self.renderer)
+    good, compare, _ = replay_to_program(p, ast, self.renderer)
     self.assertEqual(good, compare)
 
   def test_replay_with_opt(self):
     # opts=[Opt(...)] means apply a specific opt
     opts = [Opt(OptOps.UPCAST, 0, 4)]
     ast = replace_opts(self.ast, opts)
-    p = get_program(ast, self.renderer)
-    good, compare, _ = replay_get_program(p, ast, self.renderer)
+    p = do_to_program(ast, self.renderer)
+    good, compare, _ = replay_to_program(p, ast, self.renderer)
     self.assertEqual(good, compare)
 
   def test_beam(self):
     with Context(BEAM=1):
       si = (Tensor.empty(N, N) @ Tensor.empty(N, N)).schedule()[-1]
-    p = get_program(si.ast, self.renderer)
-    good, compare, _ = replay_get_program(p, si.ast, self.renderer)
+    p = do_to_program(si.ast, self.renderer)
+    good, compare, _ = replay_to_program(p, si.ast, self.renderer)
     self.assertEqual(good, compare)
 
 if __name__ == '__main__':

test/null/test_tensor_uop_mixin.py

@@ -12,6 +12,17 @@ def _t(*shape):
 def _check(tc: unittest.TestCase, t: Tensor, fn):
   tc.assertIs(fn(t).uop, fn(t.uop), f"\ntensor.uop = {fn(t).uop}\nuop = {fn(t.uop)}")
 
+class TestTensorUOpBinop(unittest.TestCase):
+  # Tensor's binop upcasts mixed dtypes via least_upper_dtype + explicit CAST; UOp should match.
+  def test_mul_float_int(self):
+    t = _t(3).float()
+    self.assertIs(_strip_unique((t * Tensor.arange(3)).uop), _strip_unique(t.uop * UOp.arange(3)))
+  def test_mul_bool_int(self):
+    t = _t(3)
+    self.assertIs(_strip_unique((t.eq(1) * Tensor.arange(3)).uop), _strip_unique(t.uop.eq(1) * UOp.arange(3)))
+  # Tensor's ufix picks float dtype when scalar is float and self is int; UOp should match.
+  def test_add_scalar_float_on_int(self): _check(self, _t(3), lambda x: x + 1.5)
+
 class TestTensorUOpGetitem(unittest.TestCase):
   # ---- pure slice patterns ----
   def test_slice_full(self):           _check(self, _t(4), lambda x: x[slice(None)])

test/null/test_uop_symbolic.py

@@ -6,7 +6,7 @@
 from tinygrad.helpers import Context
 from test.helpers import get_uops
 from tinygrad.uop.ops import UOp, Ops, graph_rewrite, sym_infer
-from tinygrad.uop.symbolic import sym, commutative, pm_simplify_valid
+from tinygrad.uop.symbolic import sym, commutative, pm_simplify_valid, pm_move_where_on_load
 from tinygrad.uop.validate import uops_to_z3
 
 def check_uop_against_string(self, v:UOp, s:str):
@@ -851,6 +851,19 @@ def test_simplex_lt(self):
     self.helper_test_variable((a+b+c*2<1).ne(True), 0, 1, "((((a+b)+c)<1)!=True)")
     self.helper_test_variable((a+b*2+c*4<1).ne(True), 0, 1, "((((a+b)+c)<1)!=True)")
 
+  def test_cast_bool_to_int_ne_const(self):
+    cond = Variable("a", 0, 3) < 2
+    # CAST(bool -> int) != 0  ->  cond
+    self.helper_test_variable(cond.cast(dtypes.int).ne(0), 0, 1, "(a<2)")
+    # CAST(bool -> int) != 1  ->  !cond
+    self.helper_test_variable(cond.cast(dtypes.int).ne(1), 0, 1, "((a<2)!=True)")
+    # CAST(bool -> int) != c (c not in {0,1})  ->  always True (CAST is 0 or 1)
+    self.helper_test_variable(cond.cast(dtypes.int).ne(2), 1, 1, "True")
+    self.helper_test_variable(cond.cast(dtypes.int).ne(-1), 1, 1, "True")
+    # CAST(bool -> weakint) folds too
+    self.helper_test_variable(cond.cast(dtypes.weakint).ne(0), 0, 1, "(a<2)")
+    self.helper_test_variable(cond.cast(dtypes.weakint).ne(1), 0, 1, "((a<2)!=True)")
+
   def test_where_removal(self):
     cond = Variable("a", 0, 3) < 2
     u1, u0 = cond.const_like(True), cond.const_like(False)
@@ -1234,6 +1247,23 @@ def test_store_load_folding(self):
     # Negative: store(idx, load(idx) + 1) should NOT fold
     self.assertEqual(graph_rewrite(index.store(index.load() + UOp.const(dtypes.int, 1)), sym).op, Ops.STORE)
 
+class TestMoveWhereOnLoad(unittest.TestCase):
+  def test_bool_index_preserves_dtype(self):
+    buf = UOp.param(0, dtypes.bool.ptr(8))
+    a = Variable("a", 0, 7)
+    r = UOp.range(8, 0)
+    # cond has a range that the rewrite can move into the valid: gate (a<4) goes into load valid
+    cond = (a < 4) & (r < 2)
+    valid = (a < 2)  # pre-existing valid on the load (to pass can_move check for the r-only clause)
+    idx = buf.index(a.valid(valid), ptr=True)
+    expr = cond.where(idx, 0)
+    out = graph_rewrite(expr, pm_move_where_on_load)
+    # any WHERE in the rewritten graph must have matched-dtype branches
+    for u in out.toposort():
+      if u.op is Ops.WHERE:
+        self.assertEqual(u.dtype, u.src[1].dtype, f"WHERE branch 1 dtype mismatch: {u}")
+        self.assertEqual(u.dtype, u.src[2].dtype, f"WHERE branch 2 dtype mismatch: {u}")
+
 class TestSymbolicRealWorld(unittest.TestCase):
   def test_resnet_half(self):
     gidx0 = Variable("gidx0", 0, 3)

test/null/test_viz.py

@@ -12,6 +12,7 @@
 
 from tinygrad.uop.ops import tracked_keys, tracked_ctxs, uop_fields, active_rewrites, active_group, _name_cnt, RewriteTrace
 from tinygrad.viz.serve import load_rewrites, get_full_rewrite, uop_to_json, VizData
+from tinygrad.codegen import to_program_cache
 
 @track_rewrites(name=True)
 def exec_rewrite(sink:UOp, pm_lst:list[PatternMatcher], names:None|list[str]=None) -> UOp:
@@ -39,6 +40,7 @@ def get_details(self, rewrite_idx:int, step:int) -> Generator[dict, None, None]:
 @contextlib.contextmanager
 def save_viz():
   for lst in [tracked_keys, tracked_ctxs, active_rewrites, active_group, _name_cnt]: lst.clear()
+  to_program_cache.clear()
   Buffer.profile_events.clear()
   cpu_events.clear()
   viz = VizTrace()
@@ -225,6 +227,10 @@ def test_const_node_visibility(self):
     self.assertEqual(list(graphs[0]), [id(a), id(alu)])
     self.assertEqual(list(graphs[1]), [id(z)])
 
+  # TODO: DEFINE_VAR (shape ()) now gets wrapped in RESHAPE+EXPAND when broadcast against a shaped operand
+  # (due to shared OpMixin._binop using _broadcasted). Either extend viz to fold RESHAPE/EXPAND around
+  # DEFINE_VAR/RANGE/SPECIAL the way it does for CONST, or redesign scalar-compiler-op broadcasting.
+  @unittest.expectedFailure
   def test_const_reshape_expand_folded(self):
     # CONST->RESHAPE->EXPAND should be folded into the ALU node, not shown as separate RESHAPE/EXPAND nodes
     c = UOp.const(dtypes.float, 1.0, device="CPU", shape=(3,4))  # creates CONST->RESHAPE->EXPAND chain

test/opt/test_gen_float4.py

@@ -3,9 +3,11 @@
 from tinygrad.uop.ops import UOp, Ops
 from tinygrad.codegen.opt import Opt, OptOps
 from tinygrad.engine.realize import get_program
-from tinygrad.helpers import AMX
+from tinygrad.helpers import DEV
 from test.helpers import replace_opts
 
+AMX = "AMX" in DEV.arch
+
 @unittest.skipUnless(Device[Device.DEFAULT].renderer.supports_float4, "need backends that support float4")
 class TestFloat4(unittest.TestCase):
   @staticmethod

test/opt/test_tensor_cores.py

@@ -7,7 +7,7 @@
 from tinygrad.uop.ops import Ops
 from tinygrad.dtype import DType
 from tinygrad.device import is_dtype_supported
-from tinygrad.helpers import AMX, DEV, Context
+from tinygrad.helpers import DEV, Context
 from test.helpers import slow, replace_opts
 from tinygrad.engine.realize import CompiledRunner, get_program
 from tinygrad.codegen.opt import Opt, OptOps, KernelOptError
@@ -18,6 +18,8 @@
 
 # NOTE: get_program always passes in Device[Device.DEFAULT].renderer explicitly for process_replay!!!
 
+AMX = "AMX" in DEV.arch
+
 def helper_tc_ensure_uops_and_opts_count(N: int, M:int, K:int, dtype_in:DType, dtype_out:DType, axis:int=0, tc_select:int=-1, tc_opt:int=0,
                                          ensure_triggered:bool=True):
   a, b = Tensor.rand(M, K, dtype=dtype_in), Tensor.rand(K, N, dtype=dtype_in)

tinygrad/codegen/__init__.py

@@ -1,7 +1,8 @@
 from typing import cast
 from dataclasses import replace
-import itertools
-from tinygrad.helpers import DISABLE_FAST_IDIV, DEVECTORIZE, TRANSCENDENTAL, SPEC, DEBUG, VIZ, IMAGE, TracingKey, Context, Target, panic
+import itertools, weakref
+from tinygrad.helpers import DISABLE_FAST_IDIV, DEVECTORIZE, TRANSCENDENTAL, SPEC, DEBUG, VIZ, IMAGE, NOOPT, EMULATED_DTYPES
+from tinygrad.helpers import TracingKey, Context, Target, panic
 from tinygrad.uop.ops import PatternMatcher, graph_rewrite, UOp, pm_lower_index_dtype, Ops, UPat, track_rewrites, KernelInfo, pyrender
 from tinygrad.uop.spec import type_verify, program_spec, kernel_spec
 from tinygrad.renderer import Renderer, ProgramSpec, Estimates
@@ -152,31 +153,34 @@ def do_compile(ctx:Renderer, prg:UOp, source:UOp) -> UOp|None:
   (UPat(Ops.PROGRAM, src=(UPat(), UPat(Ops.DEVICE), UPat(Ops.LINEAR), UPat(Ops.SOURCE, name="source")), name="prg"), do_compile),
 ])
 
+@track_rewrites(name=lambda ast,renderer,ret,**kwargs: TracingKey(ret.src[0].arg.name,(ret.src[0].arg.function_name, ast), ret=renderer), replay=True)
 @Context(ALLOW_DEVICE_USAGE=0)
-@track_rewrites(name=lambda ast,renderer,ret,**kwargs: TracingKey(ret.name, (ret.function_name, ast), ret=renderer), replay=True)
-def get_program(ast:UOp, renderer:Renderer) -> ProgramSpec:
+def do_to_program(ast:UOp, renderer:Renderer) -> UOp:
   """
-  Transform an AST into a ProgramSpec. May trigger BEAM search.
+  Transform an AST into a compiled PROGRAM. May trigger BEAM search.
 
   Args:
-    ast: The Ops.SINK rooted AST
+    ast: The Ops.SINK/Ops.PROGRAM rooted AST
     renderer: The renderer used to generate the code
 
   Returns:
-    The ProgramSpec of the program.
+    The Ops.PROGRAM with SINK/DEVICE/LINEAR/SOURCE/BINARY.
   """
-
   if ast.op is Ops.PROGRAM: prg = ast
   elif ast.op is Ops.SINK:
-    # rewrite to prg
-    assert isinstance(ast.arg, KernelInfo), "requires KernelInfo on arg to get_program"
+    assert isinstance(ast.arg, KernelInfo), "requires KernelInfo on arg to to_program"
     full_sink = full_rewrite_to_sink(ast, renderer, optimize=ast.tag is None, beam=ast.arg.beam)
     prg = UOp(Ops.PROGRAM, src=(full_sink, UOp(Ops.DEVICE, arg=renderer.target.device)))
-  else:
-    raise RuntimeError(f"can't call get_program on {ast.op}")
-
+  else: raise RuntimeError(f"can't call to_program on {ast.op}")
   prg = graph_rewrite(prg, pm_to_program, ctx=renderer, name="linearize/render")
   if VIZ: graph_rewrite(prg, PatternMatcher([]), name="View Program")
+  return prg
+
+to_program_cache: weakref.WeakValueDictionary[tuple, UOp] = weakref.WeakValueDictionary()
+def to_program(ast:UOp, renderer:Renderer) -> UOp:
+  if ast.op is Ops.PROGRAM and len(ast.src) >= 5 and ast.src[4].op is Ops.BINARY: return ast
+  key = (ast.key, type(renderer), renderer.target, NOOPT.value, DEVECTORIZE.value, EMULATED_DTYPES.value)
+  if (prg:=to_program_cache.get(key)) is None: to_program_cache[key] = prg = do_to_program(ast, renderer)
+  return prg
 
-  # create the ProgramSpec
-  return ProgramSpec.from_uop(prg)
+def get_program(ast:UOp, renderer:Renderer) -> ProgramSpec: return ProgramSpec.from_uop(to_program(ast, renderer))