From 417245ab26c53b6500236f581ddf16a42fe86eab Mon Sep 17 00:00:00 2001 From: chenyu Date: Thu, 30 Jul 2026 12:46:24 -0400 Subject: [PATCH 01/44] rewrite writes new dtype with dtype_from_uop [pr] (#17302) some Ops currently depends on explicitly set dtype and not dtype_from_uop --- tinygrad/uop/ops.py | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/tinygrad/uop/ops.py b/tinygrad/uop/ops.py index a917297e06..448ab15e2b 100644 --- a/tinygrad/uop/ops.py +++ b/tinygrad/uop/ops.py @@ -1659,7 +1659,7 @@ class RewriteContext: else: # rebuild node with rewritten srcs new_src = tuple(self.replace.get(x, x) for x in n.src) - new_n = UOp(n.op, n.dtype, new_src, n.arg, n.tag) if new_src != n.src else n + new_n = UOp(n.op, _rebuild_dtype(n, new_src), new_src, n.arg, n.tag) if new_src != n.src else n # top-down: try pm on rebuilt node, use result as-is (no re-traversal) if self.pm is not None and (rewritten:=self.pm_rewrite(new_n)) is not None: new_n = rewritten self.replace[n] = new_n @@ -1718,7 +1718,7 @@ class RewriteContext: continue else: # if srcs changed from rewrites, construct a new UOp with the new srcs - new_src_n = UOp(new_n.op, new_n.dtype, new_src, new_n.arg, new_n.tag) + new_src_n = UOp(new_n.op, _rebuild_dtype(new_n, new_src), new_src, new_n.arg, new_n.tag) # trigger a rewrite of new_src_n, then after that rewrite is done, link it back to n stack.append((n, 2, new_src_n)) stack.append((new_src_n, 0, new_src_n)) @@ -1738,6 +1738,13 @@ def graph_rewrite(sink:UOp, pm:PatternMatcher, ctx=None, bottom_up=False, name=N rewrite_ctx = RewriteContext(pm if not bottom_up else None, pm if bottom_up else bpm, ctx, enter_calls) return rewrite_ctx.walk_rewrite(sink) if walk else rewrite_ctx.unified_rewrite(sink) +def _rebuild_dtype(n:UOp, new_src:tuple[UOp,...]) -> DType: + # TODO: delete this once the dtype field is removed, every rebuild will re-derive + # TODO: these ops keep their stored dtype until dtype_from_uop works + if n.op in {Ops.INS, Ops.INDEX, Ops.CUSTOM, Ops.CUSTOMI, Ops.PYLITERAL} or \ + all(a.dtype is b.dtype or b.base.arg is Invalid for a,b in zip(n.src, new_src)): return n.dtype + return dtype_from_uop(n.op, new_src, n.arg) or n.dtype + def sint_to_uop(x:sint, dtype=dtypes.weakint) -> UOp: return UOp.const(dtype, x) def to_max_shape(shape:tuple[sint, ...]) -> tuple[int, ...]: return tuple(int(x.vmax) if isinstance(x, UOp) else x for x in shape) From d05a3e6c0ba21a0aca29569eb58e94c958db5a89 Mon Sep 17 00:00:00 2001 From: chenyu Date: Thu, 30 Jul 2026 13:08:26 -0400 Subject: [PATCH 02/44] Revert "rewrite writes new dtype with dtype_from_uop [pr] (#17302)" (#17303) This reverts commit 417245ab26c53b6500236f581ddf16a42fe86eab. --- tinygrad/uop/ops.py | 11 ++--------- 1 file changed, 2 insertions(+), 9 deletions(-) diff --git a/tinygrad/uop/ops.py b/tinygrad/uop/ops.py index 448ab15e2b..a917297e06 100644 --- a/tinygrad/uop/ops.py +++ b/tinygrad/uop/ops.py @@ -1659,7 +1659,7 @@ class RewriteContext: else: # rebuild node with rewritten srcs new_src = tuple(self.replace.get(x, x) for x in n.src) - new_n = UOp(n.op, _rebuild_dtype(n, new_src), new_src, n.arg, n.tag) if new_src != n.src else n + new_n = UOp(n.op, n.dtype, new_src, n.arg, n.tag) if new_src != n.src else n # top-down: try pm on rebuilt node, use result as-is (no re-traversal) if self.pm is not None and (rewritten:=self.pm_rewrite(new_n)) is not None: new_n = rewritten self.replace[n] = new_n @@ -1718,7 +1718,7 @@ class RewriteContext: continue else: # if srcs changed from rewrites, construct a new UOp with the new srcs - new_src_n = UOp(new_n.op, _rebuild_dtype(new_n, new_src), new_src, new_n.arg, new_n.tag) + new_src_n = UOp(new_n.op, new_n.dtype, new_src, new_n.arg, new_n.tag) # trigger a rewrite of new_src_n, then after that rewrite is done, link it back to n stack.append((n, 2, new_src_n)) stack.append((new_src_n, 0, new_src_n)) @@ -1738,13 +1738,6 @@ def graph_rewrite(sink:UOp, pm:PatternMatcher, ctx=None, bottom_up=False, name=N rewrite_ctx = RewriteContext(pm if not bottom_up else None, pm if bottom_up else bpm, ctx, enter_calls) return rewrite_ctx.walk_rewrite(sink) if walk else rewrite_ctx.unified_rewrite(sink) -def _rebuild_dtype(n:UOp, new_src:tuple[UOp,...]) -> DType: - # TODO: delete this once the dtype field is removed, every rebuild will re-derive - # TODO: these ops keep their stored dtype until dtype_from_uop works - if n.op in {Ops.INS, Ops.INDEX, Ops.CUSTOM, Ops.CUSTOMI, Ops.PYLITERAL} or \ - all(a.dtype is b.dtype or b.base.arg is Invalid for a,b in zip(n.src, new_src)): return n.dtype - return dtype_from_uop(n.op, new_src, n.arg) or n.dtype - def sint_to_uop(x:sint, dtype=dtypes.weakint) -> UOp: return UOp.const(dtype, x) def to_max_shape(shape:tuple[sint, ...]) -> tuple[int, ...]: return tuple(int(x.vmax) if isinstance(x, UOp) else x for x in shape) From 341c4ed4f52c428b00fbc1f73aa013c7ad8289c8 Mon Sep 17 00:00:00 2001 From: chenyu Date: Thu, 30 Jul 2026 13:47:21 -0400 Subject: [PATCH 03/44] rewrite writes new dtype with dtype_from_uop try 2 [pr] (#17305) with NIR fix which is INDEX related some Ops currently depends on explicitly set dtype and not dtype_from_uop --- tinygrad/renderer/nir.py | 6 +++--- tinygrad/uop/ops.py | 11 +++++++++-- 2 files changed, 12 insertions(+), 5 deletions(-) diff --git a/tinygrad/renderer/nir.py b/tinygrad/renderer/nir.py index 15352a1dda..542e579870 100644 --- a/tinygrad/renderer/nir.py +++ b/tinygrad/renderer/nir.py @@ -138,9 +138,9 @@ class NIRRenderer(Renderer): # load/store use pointer arithmetic, and the cast does nothing. NOTE: this doesn't apply to image indexing cause it's 1-D (UPat((Ops.INDEX, Ops.SHRINK), src=(UPat.var("buf"), UPat.var("off")), allow_any_len=True, name="x"), lambda x,buf,off: x.replace( src=(buf,off.cast(dtypes.long))+x.src[2:]) if buf.addrspace != AddrSpace.REG and not is_image_shape(buf._shape) else None), - # images need index to be int for nir - (UPat.var("buf").index(UPat.var("idx_y"), UPat.var("idx_x")), - lambda buf,idx_y,idx_x: buf.index(idx_y.cast(dtypes.int), idx_x.cast(dtypes.int))), + # images need index to be int for nir (coordinates only: the INDEX keeps its access dtype) + (UPat.var("buf").index(UPat.var("idx_y"), UPat.var("idx_x"), name="x"), + lambda x,buf,idx_y,idx_x: x.replace(src=(buf, idx_y.cast(dtypes.int), idx_x.cast(dtypes.int)))), ]) def_rewrite = PatternMatcher([ diff --git a/tinygrad/uop/ops.py b/tinygrad/uop/ops.py index a917297e06..448ab15e2b 100644 --- a/tinygrad/uop/ops.py +++ b/tinygrad/uop/ops.py @@ -1659,7 +1659,7 @@ class RewriteContext: else: # rebuild node with rewritten srcs new_src = tuple(self.replace.get(x, x) for x in n.src) - new_n = UOp(n.op, n.dtype, new_src, n.arg, n.tag) if new_src != n.src else n + new_n = UOp(n.op, _rebuild_dtype(n, new_src), new_src, n.arg, n.tag) if new_src != n.src else n # top-down: try pm on rebuilt node, use result as-is (no re-traversal) if self.pm is not None and (rewritten:=self.pm_rewrite(new_n)) is not None: new_n = rewritten self.replace[n] = new_n @@ -1718,7 +1718,7 @@ class RewriteContext: continue else: # if srcs changed from rewrites, construct a new UOp with the new srcs - new_src_n = UOp(new_n.op, new_n.dtype, new_src, new_n.arg, new_n.tag) + new_src_n = UOp(new_n.op, _rebuild_dtype(new_n, new_src), new_src, new_n.arg, new_n.tag) # trigger a rewrite of new_src_n, then after that rewrite is done, link it back to n stack.append((n, 2, new_src_n)) stack.append((new_src_n, 0, new_src_n)) @@ -1738,6 +1738,13 @@ def graph_rewrite(sink:UOp, pm:PatternMatcher, ctx=None, bottom_up=False, name=N rewrite_ctx = RewriteContext(pm if not bottom_up else None, pm if bottom_up else bpm, ctx, enter_calls) return rewrite_ctx.walk_rewrite(sink) if walk else rewrite_ctx.unified_rewrite(sink) +def _rebuild_dtype(n:UOp, new_src:tuple[UOp,...]) -> DType: + # TODO: delete this once the dtype field is removed, every rebuild will re-derive + # TODO: these ops keep their stored dtype until dtype_from_uop works + if n.op in {Ops.INS, Ops.INDEX, Ops.CUSTOM, Ops.CUSTOMI, Ops.PYLITERAL} or \ + all(a.dtype is b.dtype or b.base.arg is Invalid for a,b in zip(n.src, new_src)): return n.dtype + return dtype_from_uop(n.op, new_src, n.arg) or n.dtype + def sint_to_uop(x:sint, dtype=dtypes.weakint) -> UOp: return UOp.const(dtype, x) def to_max_shape(shape:tuple[sint, ...]) -> tuple[int, ...]: return tuple(int(x.vmax) if isinstance(x, UOp) else x for x in shape) From b488cc7df2705f4377015f8e38732dccbc9e01c1 Mon Sep 17 00:00:00 2001 From: chenyu Date: Thu, 30 Jul 2026 13:57:58 -0400 Subject: [PATCH 04/44] update dtype_from_uop for SHR/SHL (#17307) certain backend cast the distance, and we should not re-broadcast it --- test/unit/test_dtype_weak.py | 4 ++-- tinygrad/uop/ops.py | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/test/unit/test_dtype_weak.py b/test/unit/test_dtype_weak.py index 41bee04a99..d605345ecd 100644 --- a/test/unit/test_dtype_weak.py +++ b/test/unit/test_dtype_weak.py @@ -79,8 +79,8 @@ class TestWeakPromotion(unittest.TestCase): def test_weak_int_binop(self): v = UOp.variable("i", 0, 10, dtypes.weakint) self.assertEqual((v << 1).dtype, dtypes.weakint) - self.assertEqual(dtype_from_uop(Ops.SHL, (UOp.const(dtypes.int8, 1), UOp.const(dtypes.uint32, 1)), None), dtypes.int64) - self.assertEqual(UOp.const(dtypes.weakint, 1).alu(Ops.SHL, UOp.const(dtypes.uint8, 1)).dtype, dtypes.uint8) + self.assertEqual(dtype_from_uop(Ops.SHL, (UOp.const(dtypes.int8, 1), UOp.const(dtypes.uint32, 1)), None), dtypes.int8) + self.assertEqual(UOp.const(dtypes.weakint, 1).alu(Ops.SHL, UOp.const(dtypes.uint, 1)).dtype, dtypes.weakint) self.assertEqual((v & 3).dtype, dtypes.weakint) with self.assertRaises(RuntimeError): Tensor.const(dtypes.weakfloat, 1.0) << Tensor.const(dtypes.weakfloat, 1.0) with self.assertRaises(RuntimeError): UOp.const(dtypes.int32, 1).alu(Ops.SHL, UOp.const(dtypes.float64, 1)) diff --git a/tinygrad/uop/ops.py b/tinygrad/uop/ops.py index 448ab15e2b..bbb96a19a5 100644 --- a/tinygrad/uop/ops.py +++ b/tinygrad/uop/ops.py @@ -158,7 +158,7 @@ def dtype_from_uop(op:Ops, src:tuple[UOp,...], arg:Any) -> DType|None: return dtypes.uint64 case Ops.SHL | Ops.SHR: if not all(dtypes.is_int(x.dtype) for x in src): raise RuntimeError(f"shift operands must be int, got {[x.dtype for x in src]}") - return promo_dtype(src) + return src[0].dtype case Ops.BUFFER | Ops.PARAM: assert isinstance(arg, ParamArg), "BUFFER/PARAM must have ParamArg" return arg.dtype @@ -191,7 +191,7 @@ class UOpMetaClass(type): if dtype is None: dtype = dtype_from_uop(op, src, arg) or dtypes.void # CONST derives its dtype by value only when the constructor omits one # TODO: delete this once the dtype field is removed, for now it just re-implements spec.py - if SPEC == 2 and op is not Ops.CONST and not (op in (Ops.SHL, Ops.SHR) and src[1].dtype == dtypes.uint and dtype == src[0].dtype) and \ + if SPEC == 2 and op is not Ops.CONST and \ not any(s.base.arg is Invalid for s in src) and (expected_dtype:=dtype_from_uop(op, src, arg)) is not None and expected_dtype != dtype: raise RuntimeError(f"bad dtype {dtype}, expected {expected_dtype} on {op}") if (wret:=UOpMetaClass.ucache.get(key:=(op, dtype, src, arg, tag), None)) is not None and (ret:=wret()) is not None: return ret From da15c43e51db335c6c9c522711b656db37fcb8e8 Mon Sep 17 00:00:00 2001 From: Christopher Milan Date: Thu, 30 Jul 2026 14:22:20 -0400 Subject: [PATCH 05/44] update openpilot benchmarks (#17304) --- .github/workflows/benchmark.yml | 58 ++++++++++++++++----------------- 1 file changed, 29 insertions(+), 29 deletions(-) diff --git a/.github/workflows/benchmark.yml b/.github/workflows/benchmark.yml index 96dbaf90b1..578917726b 100644 --- a/.github/workflows/benchmark.yml +++ b/.github/workflows/benchmark.yml @@ -402,6 +402,35 @@ jobs: run: PYTHONPATH=. DEV=PCI+NV:NAK python3.11 test/test_tiny.py testcommalatest: + name: comma Benchmark (0.11.2) + runs-on: [self-hosted, Linux, comma] + timeout-minutes: 12 + defaults: + run: + shell: bash -e -o pipefail {0} + if: github.repository_owner == 'tinygrad' + steps: + - name: Checkout Code + uses: actions/checkout@v6 + - name: setup staging db + if: github.ref == 'refs/heads/update_benchmark_staging' + run: | + echo "CACHEDB=/tmp/staging.db" >> $GITHUB_ENV + rm -f /tmp/staging.db /tmp/staging.db-shm /tmp/staging.db-wal + - name: reset process replay + run: test/external/process_replay/reset.py + - name: openpilot compile3 0.11.2 supercombo + run: BENCHMARK_LOG=openpilot_0_11_2_supercombo PYTHONPATH="." ASSERT_MIN_STEP_TIME=26 DEV=QCOM FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/433f85f956837606ad1f1cbee4aa7e2158ad23c768dea914b20436c97232741b + - name: openpilot compile3 0.11.2 supercombo (from pickle) + run: BENCHMARK_LOG=openpilot_0_11_2_supercombo_run_pickle RUN_PICKLE=1 PYTHONPATH="." ASSERT_MIN_STEP_TIME=26 DEV=QCOM taskset -c 4-7 python3 examples/openpilot/compile3.py + - name: IR3 openpilot compile3 0.11.2 supercombo + run: BENCHMARK_LOG=ir3_openpilot_0_11_2_supercombo PYTHONPATH="." ASSERT_MIN_STEP_TIME=41 DEV=QCOM:IR3 FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/433f85f956837606ad1f1cbee4aa7e2158ad23c768dea914b20436c97232741b + - name: openpilot compile3 0.11.2 dmonitoring + run: BENCHMARK_LOG=openpilot_0_11_2_dmonitoring PYTHONPATH="." ASSERT_MIN_STEP_TIME=11 DEV=QCOM FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/3e7b31dfbc0a5234f1baf196513b77fc6af12204b8a8ffe8ee0417e48352f316 + - name: Run process replay tests + uses: ./.github/actions/process-replay + + testcommaold: name: comma Benchmark (0.11.0) runs-on: [self-hosted, Linux, comma] timeout-minutes: 12 @@ -432,35 +461,6 @@ jobs: - name: Run process replay tests uses: ./.github/actions/process-replay - testcommaold: - name: comma Benchmark (0.10.1) - runs-on: [self-hosted, Linux, comma] - timeout-minutes: 12 - defaults: - run: - shell: bash -e -o pipefail {0} - if: github.repository_owner == 'tinygrad' - steps: - - name: Checkout Code - uses: actions/checkout@v6 - - name: setup staging db - if: github.ref == 'refs/heads/update_benchmark_staging' - run: | - echo "CACHEDB=/tmp/staging.db" >> $GITHUB_ENV - rm -f /tmp/staging.db /tmp/staging.db-shm /tmp/staging.db-wal - - name: reset process replay - run: test/external/process_replay/reset.py - - name: DEBUG=2 openpilot compile3 0.10.1 driving_vision - run: PYTHONPATH="." DEBUG=2 DEV=QCOM FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_vision.onnx - - name: openpilot compile3 0.10.1 driving_vision - run: BENCHMARK_LOG=openpilot_0_10_1_vision PYTHONPATH="." ASSERT_MIN_STEP_TIME=17 DEV=QCOM FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_vision.onnx - - name: openpilot compile3 0.10.1 driving_policy - run: BENCHMARK_LOG=openpilot_0_10_1_policy PYTHONPATH="." ASSERT_MIN_STEP_TIME=3.2 DEV=QCOM FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_policy.onnx - - name: openpilot compile3 0.10.1 dmonitoring - run: BENCHMARK_LOG=openpilot_0_10_1_dmonitoring PYTHONPATH="." ASSERT_MIN_STEP_TIME=11 DEV=QCOM FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/dmonitoring_model.onnx - - name: Run process replay tests - uses: ./.github/actions/process-replay - testqualcommdsp: name: DSP Benchmark runs-on: [self-hosted, Linux, comma4] From b2903721212409c3002fed21d75f0998005fe8e2 Mon Sep 17 00:00:00 2001 From: Christopher Milan Date: Thu, 30 Jul 2026 14:28:45 -0400 Subject: [PATCH 06/44] ftdi reset chestnut before running comma benchmark (#17306) --- .github/workflows/benchmark.yml | 2 ++ 1 file changed, 2 insertions(+) diff --git a/.github/workflows/benchmark.yml b/.github/workflows/benchmark.yml index 578917726b..04d314e837 100644 --- a/.github/workflows/benchmark.yml +++ b/.github/workflows/benchmark.yml @@ -515,6 +515,8 @@ jobs: run: | echo "CACHEDB=/tmp/staging.db" >> $GITHUB_ENV rm -f /tmp/staging.db /tmp/staging.db-shm /tmp/staging.db-wal + - name: reset chestnut + run: python3 extra/usbgpu/debug.py -rn - name: openpilot compile3 0.10.1 driving_vision run: BENCHMARK_LOG=usbgpu_openpilot_0_10_1_vision PYTHONPATH="." GMMU=0 DEV=USB+AMD:LLVM ASSERT_MIN_STEP_TIME=50 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_vision.onnx - name: openpilot load_pickle 0.10.1 driving_vision From ce500c19469a3ac327623de7caf14c61a7e21b78 Mon Sep 17 00:00:00 2001 From: chenyu Date: Thu, 30 Jul 2026 18:51:59 -0400 Subject: [PATCH 07/44] broadcast doesn't cast CONST [pr] (#17308) * broadcast doesn't cast CONST [pr] it might change arg from python int to python float, but not cast the CONST, so we don't need or_casted in symbolic. we lower the weak dtype right before decomp * fix --- test/null/test_uop_symbolic.py | 2 +- test/null/test_viz.py | 6 ++-- test/unit/test_dtype_weak.py | 51 +++++++++++++++++++++++++++++++--- tinygrad/codegen/__init__.py | 6 ++-- tinygrad/dtype.py | 2 ++ tinygrad/mixin/elementwise.py | 11 ++++++-- tinygrad/uop/ops.py | 24 +++++++++++++--- tinygrad/uop/render.py | 4 ++- tinygrad/uop/symbolic.py | 1 + 9 files changed, 89 insertions(+), 18 deletions(-) diff --git a/test/null/test_uop_symbolic.py b/test/null/test_uop_symbolic.py index 9334440753..3a2912f96e 100644 --- a/test/null/test_uop_symbolic.py +++ b/test/null/test_uop_symbolic.py @@ -1022,7 +1022,7 @@ class TestSymbolic(unittest.TestCase): # the vars are now scalar PARAMs pvar = {u.expr: u for u in rewritten_uop.toposort() if u.op is Ops.PARAM} - self.assertEqual(rewritten_uop, (pvar['s']<2).where(pvar['a'].cast(dtypes.half), pvar['b'].cast(dtypes.half))) + self.assertEqual(rewritten_uop, (pvar['s'] UOp: # floordiv+mod / dtype decomp (early) supported_ops = tuple(ren.code_for_op.keys()) - pm_decomp = symbolic_simple+get_simplifying_rewrite_patterns(supported_ops) + pm_decomp = symbolic_simple+pm_commit_weak+get_simplifying_rewrite_patterns(supported_ops) sink = graph_rewrite(sink, pm_decomp, name="early decompositions") # late decomps + move gates from unrenderable INVALID where - sink = graph_rewrite(sink, pm_dtype_decomps, ctx=(set(), ren), name="decomp dtypes") + sink = graph_rewrite(sink, pm_dtype_decomps+pm_commit_weak, ctx=(set(), ren), name="decomp dtypes") pm_decomp = pm_decomp+\ get_late_rewrite_patterns(supported_ops, bool(DISABLE_FAST_IDIV))+\ get_transcendental_patterns(supported_ops, TRANSCENDENTAL>=2) diff --git a/tinygrad/dtype.py b/tinygrad/dtype.py index 2e97e9a8de..af8f7c7267 100644 --- a/tinygrad/dtype.py +++ b/tinygrad/dtype.py @@ -165,6 +165,8 @@ assert dtypes.is_float(dtypes.default_float), f"{DEFAULT_FLOAT.value} is not a f assert dtypes.is_int(dtypes.default_int), f"{DEFAULT_INT.value} is not an int dtype" def strong_dtype(dtype:DType) -> DType: return {dtypes.weakint: dtypes.default_int, dtypes.weakfloat: dtypes.default_float}.get(dtype, dtype) +def weak_dtype(dtype:DType) -> DType: + return dtypes.weakfloat if dtypes.is_float(dtype) else dtypes.weakint if dtypes.is_int(dtype) else dtype # https://jax.readthedocs.io/en/latest/jep/9407-type-promotion.html # we don't support complex type diff --git a/tinygrad/mixin/elementwise.py b/tinygrad/mixin/elementwise.py index 7c9c2b3047..31a5d7b25a 100644 --- a/tinygrad/mixin/elementwise.py +++ b/tinygrad/mixin/elementwise.py @@ -1,7 +1,7 @@ import math, functools, operator from typing import TYPE_CHECKING, Literal, Self from tinygrad.uop import Ops -from tinygrad.dtype import dtypes, ConstType, PyConst, least_upper_dtype, least_upper_float +from tinygrad.dtype import dtypes, ConstType, PyConst, least_upper_dtype, least_upper_float, weak_dtype from tinygrad.helpers import argfix, polyN from tinygrad.mixin.creation import CreationMixin @@ -21,7 +21,12 @@ class ElementwiseMixin(CreationMixin): def _broadcasted(self, y: 'Self|ConstType|UOp', reverse: bool = False) -> tuple[Self, Self]: y = self.ufix(y) x, y = (self, y) if not reverse else (y, self) - return x.cast(out_dtype := least_upper_dtype(x.dtype, y.dtype)), y.cast(out_dtype) + out_dtype = least_upper_dtype(x.dtype, y.dtype) + # keep weak CONST weak, might lift weakint -> weakfloat + def promote(t): + if t.dtype in dtypes.weaks and t._uop.base.op is Ops.CONST: return t._wrap_uop(t._uop.const(weak_dtype(out_dtype), t._uop.base.arg, t.shape)) + return t.cast(out_dtype) + return promote(x), promote(y) def _binop(self, op: Ops, x: Self | ConstType, reverse: bool) -> Self: lhs, rhs = self._broadcasted(x, reverse) @@ -390,6 +395,8 @@ class ElementwiseMixin(CreationMixin): ``` """ t, x = self._broadcasted(x) + # ~ is width-dependent: min(a,b) == ~max(~a,~b) only holds at a common width, so a weak operand commits at its sibling's + t, x = t.cast(dt:=least_upper_dtype(t.dtype, x.dtype)), x.cast(dt) return t._inverse().maximum(x._inverse())._inverse() def copysign(self, other: Self | ConstType) -> Self: diff --git a/tinygrad/uop/ops.py b/tinygrad/uop/ops.py index bbb96a19a5..5b2f6ec8f3 100644 --- a/tinygrad/uop/ops.py +++ b/tinygrad/uop/ops.py @@ -5,7 +5,7 @@ from dataclasses import dataclass, replace from enum import Enum, auto from tinygrad.uop import Ops, GroupOp from tinygrad.dtype import ConstType, dtypes, DType, DTypeLike, truncate, least_upper_dtype, least_upper_float, Invalid, AddrSpace, strong_dtype -from tinygrad.dtype import ConstFloat, PyConst, InvalidType, storage_fmt_for_dtype, to_storage_scalar, from_storage_scalar +from tinygrad.dtype import ConstFloat, PyConst, InvalidType, storage_fmt_for_dtype, to_storage_scalar, from_storage_scalar, weak_dtype from tinygrad.device import Buffer, MultiBuffer, canonicalize_device, TinyELF from tinygrad.helpers import ContextVar, all_int, prod, getenv, all_same, Context, partition, temp, unwrap, T, argfix, Metadata, flatten, TRACEMETA from tinygrad.helpers import PROFILE, dedup, cdiv, cmod, floordiv, floormod, diskcache_put, to_function_name, cpu_profile, TracingKey @@ -191,8 +191,10 @@ class UOpMetaClass(type): if dtype is None: dtype = dtype_from_uop(op, src, arg) or dtypes.void # CONST derives its dtype by value only when the constructor omits one # TODO: delete this once the dtype field is removed, for now it just re-implements spec.py + # an INDEX presents its access dtype, which a still-weak source matches up to weakness if SPEC == 2 and op is not Ops.CONST and \ - not any(s.base.arg is Invalid for s in src) and (expected_dtype:=dtype_from_uop(op, src, arg)) is not None and expected_dtype != dtype: + not any(s.base.arg is Invalid for s in src) and (expected_dtype:=dtype_from_uop(op, src, arg)) is not None and expected_dtype != dtype and \ + not (op is Ops.INDEX and weak_dtype(expected_dtype) == weak_dtype(dtype)): raise RuntimeError(f"bad dtype {dtype}, expected {expected_dtype} on {op}") if (wret:=UOpMetaClass.ucache.get(key:=(op, dtype, src, arg, tag), None)) is not None and (ret:=wret()) is not None: return ret UOpMetaClass.ucache[key] = weakref.ref(created:=super().__call__(*key)) @@ -586,7 +588,7 @@ class UOp(RandMixin, metaclass=UOpMetaClass): return UOp.const(dtype or self.dtype, b).broadcast(self.max_numel()) def ufix(self, x): if isinstance(x, UOp): return x - return UOp.const(least_upper_dtype(self.dtype, dtypes.from_py(x)), x) + return UOp.const(None, x) def broadcast(self, count:int): if count == 1: return self return UOp(Ops.STACK, src=(self,)*count) @@ -1780,7 +1782,21 @@ def lower_weak_srcs(ctx:dict[UOp, UOp]|None, u:UOp) -> UOp|None: # a comparison demands a common operand width: lower it whole so the Binary rule unifies its operands ret = lower(u) if u.op in GroupOp.Comparison else u.replace(src=tuple(lower(s) if s.dtype in dtypes.weaks else s for s in u.src)) return None if ret is u else ret -pm_lower_index_dtype = PatternMatcher([ + +def commit_weak_srcs(u:UOp) -> UOp|None: + if (dt:=least_upper_dtype(*(s.dtype for s in u.src))) in dtypes.weaks: return None + # the root re-derives: a shift's dtype is its lhs's, so committing the lhs commits the node too + return u.replace(dtype=None, src=tuple(s.cast(dt) if s.dtype in dtypes.weaks else s for s in u.src)) + +# runs in index lowering and in the decomps: a rule that mints a weak const commits it in the same rewrite, so none reaches the renderer +pm_commit_weak = PatternMatcher([ + (UPat(GroupOp.Broadcastable, name="u"), commit_weak_srcs), + # demand from the destination: a STORE's weak value commits at the destination's dtype + (UPat(Ops.STORE, src=(UPat(), UPat(dtype=dtypes.weaks)), allow_any_len=True, name="u"), + lambda u: u.replace(src=(u.src[0], u.src[1].cast(u.src[0].dtype), *u.src[2:]))), +]) + +pm_lower_index_dtype = pm_commit_weak+PatternMatcher([ (UPat(GroupOp.All, name="u"), lambda ctx,u: lower_weak_srcs(ctx, u) if u.dtype not in dtypes.weaks and any(s.dtype in dtypes.weaks for s in u.src) else None), # a valid index into an n-element buffer lives in [0,n): a gated long index narrows when n-1 fits int32 (out-of-gate wraps, discarded) diff --git a/tinygrad/uop/render.py b/tinygrad/uop/render.py index 84d2a6bf24..bf8b32a4f1 100644 --- a/tinygrad/uop/render.py +++ b/tinygrad/uop/render.py @@ -75,7 +75,7 @@ def render_marg(ctx,x:UOp): return f"({','.join(pieces)})" if len(pieces) != 1 else f"({pieces[0]},)" sugar = {Ops.SINK, Ops.END, Ops.STORE, Ops.LOAD, Ops.SQRT, Ops.INDEX, Ops.REDUCE, Ops.AFTER, Ops.THREEFRY, - Ops.WHERE, Ops.RECIPROCAL, Ops.EXP2, Ops.LOG2, Ops.SIN, Ops.CONTIGUOUS, Ops.BARRIER, Ops.DETACH} + Ops.RECIPROCAL, Ops.EXP2, Ops.LOG2, Ops.SIN, Ops.CONTIGUOUS, Ops.BARRIER, Ops.DETACH} pm_pyrender_extra = PatternMatcher([ (UPat(Ops.CONST, src=(), name="x"), lambda x: f"UOp.const({x.dtype}, {x.arg})"), (UPat((Ops.CAST, Ops.BITCAST), name="x"), lambda ctx,x: f"{ctx[x.src[0]]}.{x.op.name.lower()}({x.dtype})"), @@ -100,6 +100,8 @@ pm_pyrender_extra = PatternMatcher([ # explicit trunc ops: `//` and `%` parse as FLOORDIV/FLOORMOD, so render CDIV/CMOD via .alu() (UPat(Ops.CDIV, name="x"), lambda ctx,x: f"{ctx[x.src[0]]}.alu(Ops.CDIV, {ctx[x.src[1]]})"), (UPat(Ops.CMOD, name="x"), lambda ctx,x: f"{ctx[x.src[0]]}.alu(Ops.CMOD, {ctx[x.src[1]]})"), + # `.where` re-promotes its operands, so render WHERE via .alu() too + (UPat(Ops.WHERE, name="x"), lambda ctx,x: f"{ctx[x.src[0]]}.alu(Ops.WHERE, {ctx[x.src[1]]}, {ctx[x.src[2]]})"), (UPat(set(syms.keys())-{Ops.SUB, Ops.CDIV, Ops.CMOD}, name="x"), lambda ctx,x: strip_binary_parens(x, ctx[x.src[0]], ctx[x.src[1]], lambda a,b: f"({a}{syms[x.op]}{b})")), (UPat(sugar, src=(), name="x"), lambda x: f"UOp.{x.op.name.lower()}("+', '.join(([f'arg={repr(x.arg)}'] if x.arg is not None else []))+")"), diff --git a/tinygrad/uop/symbolic.py b/tinygrad/uop/symbolic.py index 65da3f2f56..0b35760e85 100644 --- a/tinygrad/uop/symbolic.py +++ b/tinygrad/uop/symbolic.py @@ -155,6 +155,7 @@ symbolic_simple = pm_data_invalid + PatternMatcher([ (UPat.var("x") * 0, lambda x: x.const_like(float("nan") if x.op is Ops.CONST and isinstance(x.arg, float) and (math.isnan(x.arg) or math.isinf(x.arg)) else 0)), # *** cast/bitcast *** + # TODO: delete this once CONST has no dtype (UPat(Ops.CAST, name="root", src=(UPat.cvar("c"),)), lambda root, c: root.const_like(c.arg)), (UPat((Ops.CAST, Ops.BITCAST), name="root"), lambda root: root.src[0] if root.dtype == root.src[0].dtype else None), (UPat(Ops.BITCAST, name="root", src=(UPat.cvar("c"),)), fold_bitcast), From fe8ece7efaad0a29f965a7977d17eb06398277bf Mon Sep 17 00:00:00 2001 From: chenyu Date: Thu, 30 Jul 2026 19:47:47 -0400 Subject: [PATCH 08/44] fix weak for image gate fusion [pr] (#17310) --- test/null/test_simplify_valid_idx.py | 12 +++++++++++- test/unit/test_dtype_weak.py | 11 +++++++++-- tinygrad/uop/ops.py | 8 ++++++-- 3 files changed, 26 insertions(+), 5 deletions(-) diff --git a/test/null/test_simplify_valid_idx.py b/test/null/test_simplify_valid_idx.py index 7bf3dc000f..436dee3d96 100644 --- a/test/null/test_simplify_valid_idx.py +++ b/test/null/test_simplify_valid_idx.py @@ -2,7 +2,7 @@ import unittest, itertools from tinygrad.codegen.late.coalesce import indexing_simplify from tinygrad.dtype import dtypes -from tinygrad.uop.ops import UOp, Ops, graph_rewrite +from tinygrad.uop.ops import UOp, Ops, graph_rewrite, pm_lower_index_dtype from tinygrad.uop.symbolic import simplify_valid, sym, pm_move_where_on_load from tinygrad.helpers import Context from test.helpers import full_rewrite @@ -492,6 +492,16 @@ class TestImageSimplification(unittest.TestCase): load = get_load_image_uop((1, 48, 4), valid, idx) self.check(load, None, "(r12*4+(gidx0+3)%4+(gidx0+3)//4*24+-3888)", "0") + def test_drop_gate_committed_in_the_index_pass(self): + # the fused index pass runs without symbolic, so committing a weak src must not leave a CAST that + # symbolic later folds inside the index only: the gate's copy of the expression has to stay the same node + f = UOp.variable("f", 0.0, 9.0, dtypes.float) + idx_y = (f + UOp.const(None, 1.0)).cast(dtypes.int) + load = get_load_image_uop((10, 10, 4), (UOp.const(None, -1) < idx_y) & (idx_y < UOp.const(None, 10)), + (Special("gidx0", 10), idx_y)) + off = graph_rewrite(load.sink(), pm_lower_index_dtype+indexing_simplify, ctx={}).src[0].src[0] + self.assertEqual(off.src[1].get_valid(), UOp.const(dtypes.bool, True)) + class TestDropTrueGate(unittest.TestCase): def test_drop_true_gate_on_index(self): # test that INDEX with a constant True valid gets simplified to drop the valid diff --git a/test/unit/test_dtype_weak.py b/test/unit/test_dtype_weak.py index 179d37e24b..03b2fc8cda 100644 --- a/test/unit/test_dtype_weak.py +++ b/test/unit/test_dtype_weak.py @@ -3,7 +3,8 @@ import tempfile, unittest, math from tinygrad import Tensor, dtypes, TinyJit from tinygrad.helpers import Context from tinygrad.dtype import least_upper_float -from tinygrad.uop.ops import UOp, Ops, dtype_from_uop, graph_rewrite, pm_lower_index_dtype +from tinygrad.uop.ops import UOp, Ops, dtype_from_uop, graph_rewrite, pm_lower_index_dtype, pm_commit_weak +from tinygrad.uop.symbolic import symbolic_simple from tinygrad.uop.spec import spec_shared, type_verify from tinygrad.engine.jit import JitError @@ -82,7 +83,8 @@ class TestWeakPromotion(unittest.TestCase): dst = UOp.param(0, dtypes.bfloat16, (1,)).index(UOp.const(None, 0).cast(dtypes.int32)) gate = UOp.const(None, True) out = graph_rewrite(dst.store(UOp.const(None, 5.0), gate), pm_lower_index_dtype, ctx={}) - self.assertEqual((out.src[1].dtype, out.src[2]), (dtypes.bfloat16, gate)) + # a bare weak CONST commits directly: the pass runs without symbolic, so a CAST here would survive it + self.assertEqual((out.src[1], out.src[2]), (UOp.const(dtypes.bfloat16, 5.0), gate)) def test_weak_srcs_commit_only_at_a_concrete_lub(self): weak_lub = UOp(Ops.ADD, src=(UOp.const(None, 1), UOp.const(None, 1.0))) @@ -91,6 +93,11 @@ class TestWeakPromotion(unittest.TestCase): where = graph_rewrite(UOp(Ops.WHERE, src=(UOp.const(None, True), concrete, UOp.const(None, 1.0))), pm_lower_index_dtype, ctx={}) self.assertEqual(tuple(x.dtype for x in where.src), (dtypes.bool, dtypes.float16, dtypes.float16)) + def test_weak_shift_lhs_commits_the_node(self): + # a shift derives its lhs's dtype, so committing the lhs restates the root (WGSL's packed store writes `mask << shift_am`) + shl = graph_rewrite(UOp.const(None, 0xFFFF) << UOp.variable("x", 0, 16, dtypes.uint), symbolic_simple+pm_commit_weak) + self.assertEqual((shl.dtype, shl.src[0]), (dtypes.uint, UOp.const(dtypes.uint, 0xFFFF))) + @unittest.expectedFailure # TODO: a weak const defers to its consumer (JAX): these dtypes change once python scalars are weak consts def test_changed_rows(self): t_i8, t_f16, t_bf16 = Tensor([1], dtype=dtypes.int8), Tensor([1], dtype=dtypes.float16), Tensor([1], dtype=dtypes.bfloat16) diff --git a/tinygrad/uop/ops.py b/tinygrad/uop/ops.py index 5b2f6ec8f3..bc6cdfc71d 100644 --- a/tinygrad/uop/ops.py +++ b/tinygrad/uop/ops.py @@ -1783,17 +1783,21 @@ def lower_weak_srcs(ctx:dict[UOp, UOp]|None, u:UOp) -> UOp|None: ret = lower(u) if u.op in GroupOp.Comparison else u.replace(src=tuple(lower(s) if s.dtype in dtypes.weaks else s for s in u.src)) return None if ret is u else ret +def commit_weak(s:UOp, dt:DType) -> UOp: + # a bare weak CONST commits directly (its number must fit), a weak non-const src takes the demand cast + return UOp.const(dt, s.arg) if s.op is Ops.CONST else s.cast(dt) + def commit_weak_srcs(u:UOp) -> UOp|None: if (dt:=least_upper_dtype(*(s.dtype for s in u.src))) in dtypes.weaks: return None # the root re-derives: a shift's dtype is its lhs's, so committing the lhs commits the node too - return u.replace(dtype=None, src=tuple(s.cast(dt) if s.dtype in dtypes.weaks else s for s in u.src)) + return u.replace(dtype=None, src=tuple(commit_weak(s, dt) if s.dtype in dtypes.weaks else s for s in u.src)) # runs in index lowering and in the decomps: a rule that mints a weak const commits it in the same rewrite, so none reaches the renderer pm_commit_weak = PatternMatcher([ (UPat(GroupOp.Broadcastable, name="u"), commit_weak_srcs), # demand from the destination: a STORE's weak value commits at the destination's dtype (UPat(Ops.STORE, src=(UPat(), UPat(dtype=dtypes.weaks)), allow_any_len=True, name="u"), - lambda u: u.replace(src=(u.src[0], u.src[1].cast(u.src[0].dtype), *u.src[2:]))), + lambda u: u.replace(src=(u.src[0], commit_weak(u.src[1], u.src[0].dtype), *u.src[2:]))), ]) pm_lower_index_dtype = pm_commit_weak+PatternMatcher([ From f2c2f4456b31ebd7e23e6c3ada1c387d6536e442 Mon Sep 17 00:00:00 2001 From: Noah Schiro <59176275+NoahSchiro@users.noreply.github.com> Date: Thu, 30 Jul 2026 20:05:11 -0400 Subject: [PATCH 09/44] Add softmin (#17292) * Add softmin * Remove extra tests --- test/backend/test_ops.py | 2 ++ tinygrad/mixin/op.py | 22 ++++++++++++++++++++++ 2 files changed, 24 insertions(+) diff --git a/test/backend/test_ops.py b/test/backend/test_ops.py index d3ebe02eff..520ed9d793 100644 --- a/test/backend/test_ops.py +++ b/test/backend/test_ops.py @@ -1728,6 +1728,8 @@ class TestOps(unittest.TestCase): helper_test_op([(10,10,10)], lambda x: x.log_softmax(0), atol=1e-7, grad_atol=1e-7) helper_test_op([(10,10,10)], lambda x: x.log_softmax(1), atol=1e-7, grad_atol=1e-7) helper_test_op([(10,10,10)], lambda x: x.log_softmax(2), atol=1e-7, grad_atol=1e-7) + def test_softmin(self): + helper_test_op([(45,65)], torch.nn.Softmin(dim=1), Tensor.softmin, atol=1e-7, grad_atol=1e-7) def test_normalize(self): helper_test_op([(45,65)], lambda x: torch.nn.functional.normalize(x), lambda x: x.normalize(), atol=1e-7, grad_atol=1e-7) diff --git a/tinygrad/mixin/op.py b/tinygrad/mixin/op.py index 0c5e1cfd16..90f5d8f4fd 100644 --- a/tinygrad/mixin/op.py +++ b/tinygrad/mixin/op.py @@ -703,6 +703,28 @@ class OpMixin(ElementwiseMixin, ReduceMixin): m, _, ss = self._softmax(axis, dtype) return m - ss.log() + def softmin(self, axis=-1, dtype:DTypeLike|None=None) -> Self: + """ + Applies the softmin function to the tensor along the specified axis. + + Rescales the elements of the tensor such that they lie in the range [0, 1] and sum to 1. + + You can pass in the `axis` keyword argument to control the axis along which the softmin is computed. + + ```python exec="true" source="above" session="tensor" result="python" + Tensor.manual_seed(42) + t = Tensor.randn(2, 3) + print(t.numpy()) + ``` + ```python exec="true" source="above" session="tensor" result="python" + print(t.softmin().numpy()) + ``` + ```python exec="true" source="above" session="tensor" result="python" + print(t.softmin(axis=0).numpy()) + ``` + """ + return (-self).softmax(axis, dtype) + def cat(self, *args:Self, dim:int=0) -> Self: """ Concatenates self with other tensors in `args` along an axis specified by `dim`. From 6608b9d8e2aff2f1c841c2f8f720849647efd834 Mon Sep 17 00:00:00 2001 From: chenyu Date: Thu, 30 Jul 2026 23:36:12 -0400 Subject: [PATCH 10/44] adjust lower weak order [pr] (#17314) * adjust lower weak order [pr] * no symbolic_simple --- tinygrad/codegen/__init__.py | 6 +++--- tinygrad/uop/ops.py | 7 +++++++ 2 files changed, 10 insertions(+), 3 deletions(-) diff --git a/tinygrad/codegen/__init__.py b/tinygrad/codegen/__init__.py index 0364654a31..230897630a 100644 --- a/tinygrad/codegen/__init__.py +++ b/tinygrad/codegen/__init__.py @@ -3,7 +3,7 @@ import itertools, functools from tinygrad.helpers import DISABLE_FAST_IDIV, TRANSCENDENTAL, SPEC, DEBUG, VIZ, IMAGE, NOOPT, EMULATED_DTYPES, NOLOCALS, USE_TC from tinygrad.helpers import ALLOW_TF32, DEFAULT_FLOAT, DEFAULT_INT, TracingKey, Context, panic from tinygrad.uop.ops import PatternMatcher, graph_rewrite, UOp, pm_lower_index_dtype, Ops, UPat, track_rewrites, KernelInfo, ProgramInfo, GroupOp -from tinygrad.uop.ops import AxisType, pm_commit_weak +from tinygrad.uop.ops import AxisType, pm_commit_weak, pm_cast_weak from tinygrad.uop.render import pyrender from tinygrad.uop.spec import type_verify, spec_tensor, spec_program from tinygrad.renderer import Renderer, Estimates @@ -358,7 +358,7 @@ def full_rewrite_to_sink(ast:UOp, ren:Renderer, optimize:bool=True) -> UOp: # floordiv+mod / dtype decomp (early) supported_ops = tuple(ren.code_for_op.keys()) - pm_decomp = symbolic_simple+pm_commit_weak+get_simplifying_rewrite_patterns(supported_ops) + pm_decomp = symbolic_simple+get_simplifying_rewrite_patterns(supported_ops) sink = graph_rewrite(sink, pm_decomp, name="early decompositions") # late decomps + move gates from unrenderable INVALID where @@ -371,7 +371,7 @@ def full_rewrite_to_sink(ast:UOp, ren:Renderer, optimize:bool=True) -> UOp: # final rules for the renderer (without sym) extra_matcher = ren.extra_matcher if ren.extra_matcher is not None else PatternMatcher([]) - pm_final_rewrite = pm_decomp+extra_matcher+pm_split_ends + pm_final_rewrite = pm_commit_weak+pm_cast_weak+pm_decomp+extra_matcher+pm_split_ends sink = graph_rewrite(sink, pm_final_rewrite+pm_remove_invalid, ctx=ren, name="final rewrite") # add implicit barriers (stores/loads through LOCAL memory ordered by AFTER or across loop iterations need workgroup barriers) diff --git a/tinygrad/uop/ops.py b/tinygrad/uop/ops.py index bc6cdfc71d..0670a57b2d 100644 --- a/tinygrad/uop/ops.py +++ b/tinygrad/uop/ops.py @@ -1800,6 +1800,13 @@ pm_commit_weak = PatternMatcher([ lambda u: u.replace(src=(u.src[0], commit_weak(u.src[1], u.src[0].dtype), *u.src[2:]))), ]) +# push cast to weak src +pm_cast_weak = PatternMatcher([ + (UPat(Ops.CAST, name="c", src=(UPat(GroupOp.Broadcastable, dtype=dtypes.weaks, name="u"),)), + lambda c,u: u.replace(dtype=None, src=tuple(commit_weak(s, c.dtype) if s.dtype in dtypes.weaks else s for s in u.src)).cast(c.dtype) + if c.dtype not in dtypes.weaks else None), +]) + pm_lower_index_dtype = pm_commit_weak+PatternMatcher([ (UPat(GroupOp.All, name="u"), lambda ctx,u: lower_weak_srcs(ctx, u) if u.dtype not in dtypes.weaks and any(s.dtype in dtypes.weaks for s in u.src) else None), From d65ea465edfd561d69a5d5f465fd1c820f733acd Mon Sep 17 00:00:00 2001 From: George Hotz <72895+geohot@users.noreply.github.com> Date: Thu, 30 Jul 2026 20:53:28 -0700 Subject: [PATCH 11/44] cleanup gemm fragment + add store unshard (#17313) * cleanup gemm fragment + add store unshard * multi * fix --- extra/gemm_fragment.py | 156 +++++++++-------------------- test/backend/test_custom_kernel.py | 72 +++++++++++++ tinygrad/schedule/multi.py | 41 ++++++-- 3 files changed, 157 insertions(+), 112 deletions(-) diff --git a/extra/gemm_fragment.py b/extra/gemm_fragment.py index d9ad4c5d22..457439e3d1 100644 --- a/extra/gemm_fragment.py +++ b/extra/gemm_fragment.py @@ -1,55 +1,25 @@ """ tilelang-style matmul_relu written with tinygrad UOp APIs. -Demonstrates that tilelang's T.alloc_fragment is expressible with existing -tinygrad primitives: a per-thread REG buffer, wrapped in one Ops.UNSHARD per -sharded axis over the LOCAL thread-grid ranges to form the full logical tile. -Here the 64 threads are an 8x8 grid and each thread owns an 8x8 sub-tile -- -the 2-D fragment layout tilelang infers. The kernel is written against the -full-tile UNSHARD view, and multi_pm (the same pass that lowers multi-device -UNSHARDs) resolves it into per-thread shard code. - Reference tilelang kernel: - - @tilelang.jit - def matmul_relu(A, B, block_M=64, block_N=64, block_K=64, - dtype=T.float16, accum_dtype=T.float32): - M, N, K = T.const('M, N, K') - C = T.empty([M, N], dtype) - with T.Kernel(T.ceildiv(N, block_N), T.ceildiv(M, block_M), threads=128) as (bx, by): - A_shared = T.alloc_shared((block_M, block_K), dtype) - B_shared = T.alloc_shared((block_K, block_N), dtype) - C_local = T.alloc_fragment((block_M, block_N), accum_dtype) - T.clear(C_local) - for ko in T.Pipelined(T.ceildiv(K, block_K), num_stages=3): - T.copy(A[by * block_M, ko * block_K], A_shared) - T.copy(B[ko * block_K, bx * block_N], B_shared) - T.gemm(A_shared, B_shared, C_local) - for i, j in T.Parallel(block_M, block_N): - C_local[i, j] = T.max(C_local[i, j], 0) - T.copy(C_local, C[by * block_M, bx * block_N]) - return C - -API mapping (tilelang -> tinygrad UOps, idioms from test/backend/test_custom_kernel.py): - - T.Kernel(gx, gy, threads=T) -> AxisType.GLOBAL ranges (blocks) + AxisType.LOCAL ranges (thread grid) - T.alloc_shared(shape, dtype) -> UOp.placeholder(shape, dtype, slot, AddrSpace.LOCAL) - T.alloc_fragment(shape, dt) -> per-thread REG placeholder, wrapped in one Ops.UNSHARD per sharded axis over - the AxisType.LOCAL ranges: fragment.unshard((axis_y, axis_x), (ty, tx)). - The full logical tile is the shard with each sharded axis multiplied by its - range size, exactly like device sharding, but the sharding axes are thread - axes carried by the RANGE metadata instead of a device tuple. C_local[i, j] - with [i, j] in this thread's shard is INDEX on the UNSHARD, which multi_pm - resolves into INDEX on the per-thread REG shard, axis by axis. - T.copy(gmem_slice, smem) -> smem[thread_idx].set(gmem_slice[thread_idx], end=copy_rng). set returns the - smem tile AFTER the copy; the implicit-barrier pass turns the store->load - dependency of the loop that consumes it into a workgroup barrier - T.gemm (no WMMA) -> C_local[..].set(C_local.after(k)[..] + a_shared[..] * b_shared[..], end=k) - with k a loop-carried LOOP range (codegen builds the register accumulator - from this self-referential store automatically) - T.copy(fragment, gmem) -> gmem.index(gidx).store(C_local[..]).end(all_ranges) - UNSHARD lowering -> multi_pm in codegen (full_rewrite_to_sink): INDEX/AFTER/STORE ops on the - full-tile view become per-thread shard ops, no UNSHARD survives into the program. + @tilelang.jit + def matmul_relu(A, B, block_M=64, block_N=64, block_K=64, + dtype=T.float16, accum_dtype=T.float32): + M, N, K = T.const('M, N, K') + C = T.empty([M, N], dtype) + with T.Kernel(T.ceildiv(N, block_N), T.ceildiv(M, block_M), threads=128) as (bx, by): + A_shared = T.alloc_shared((block_M, block_K), dtype) + B_shared = T.alloc_shared((block_K, block_N), dtype) + C_local = T.alloc_fragment((block_M, block_N), accum_dtype) + T.clear(C_local) + for ko in T.Pipelined(T.ceildiv(K, block_K), num_stages=3): + T.copy(A[by * block_M, ko * block_K], A_shared) + T.copy(B[ko * block_K, bx * block_N], B_shared) + T.gemm(A_shared, B_shared, C_local) + for i, j in T.Parallel(block_M, block_N): + C_local[i, j] = T.max(C_local[i, j], 0) + T.copy(C_local, C[by * block_M, bx * block_N]) + return C """ from tinygrad.dtype import dtypes, AddrSpace, DType @@ -61,25 +31,17 @@ from tinygrad.tensor import Tensor # tilelang builtins, expressed with tinygrad UOp APIs # --------------------------------------------------------------------------- -def alloc_shared(shape:tuple[int, ...], dtype:DType) -> UOp: +def alloc_shared(shape:tuple[int, ...], dtype:DType, slot:int) -> UOp: """T.alloc_shared: one LOCAL buffer shared by all threads in the block.""" - return UOp.placeholder(tuple(shape), dtype, next(UOp.unique_num), AddrSpace.LOCAL) + return UOp.placeholder(tuple(shape), dtype, slot, AddrSpace.LOCAL) -def alloc_fragment(shape:tuple[int, ...], dtype:DType, axes:tuple[int, ...], rngs:tuple[UOp, ...]) -> UOp: - """T.alloc_fragment: per-thread REG fragment + UNSHARD over the LOCAL thread grid. - - Each thread privately owns shape[axis]//threads elements along every sharded - axis in a REG buffer. The UNSHARDs over the LOCAL thread ranges present the - full logical tile: full_shape = shard_shape with each sharded axis multiplied - by its range size. This is exactly how UNSHARD carries a DEVICE axis today, - except the sharding axes are thread axes carried by the RANGE metadata. - """ +def alloc_fragment(shape:tuple[int, ...], dtype:DType, slot:int, axes:tuple[int, ...], rngs:tuple[UOp, ...]) -> UOp: + """T.alloc_fragment: per-thread REG fragment + UNSHARD over the LOCAL thread grid.""" assert len(axes) == len(rngs) assert all(tnum.op is Ops.RANGE and tnum.arg[-1] is AxisType.LOCAL for tnum in rngs), "fragments shard over LOCAL ranges" - assert all(shape[a] % (int(rng.vmax)+1) == 0 for a, rng in zip(axes, rngs)) by_axis = dict(zip(axes, rngs)) shard_shape = tuple(s // (int(by_axis[i].vmax)+1) if i in by_axis else s for i, s in enumerate(shape)) - fragment = UOp.placeholder(shard_shape, dtype, next(UOp.unique_num), AddrSpace.REG) + fragment = UOp.placeholder(shard_shape, dtype, slot, AddrSpace.REG) return fragment.unshard(axes, rngs) # --------------------------------------------------------------------------- @@ -106,64 +68,45 @@ def matmul_relu_kernel(c:UOp, a:UOp, b:UOp) -> UOp: # with T.Kernel(T.ceildiv(N, BLOCK_N), T.ceildiv(M, BLOCK_M), threads=128) as (bx, by): bx = UOp.range(cdiv(N, BLOCK_N), 0, AxisType.GLOBAL) by = UOp.range(cdiv(M, BLOCK_M), 1, AxisType.GLOBAL) - # tx (N, 16) is the fast/inner LOCAL axis so a warp covers 16 cols x 2 rows -- - # matching tilelang's (tidx>>4, tidx&15) warp composition. This keeps the 8 A_shared - # reads in a warp on only 2 row-groups (broadcast across 16 cols) instead of 8 rows - # (8-way bank conflict), since A_shared[row*512 + ...] all map to the same bank when 8 - # distinct rows land in one warp. + + # 16*8 threads = 128 threads tx = UOp.range(TX, 2, AxisType.LOCAL) ty = UOp.range(TY, 3, AxisType.LOCAL) - # A_shared = T.alloc_shared((BLOCK_M, BLOCK_K), dtype) - # B_shared = T.alloc_shared((BLOCK_K, BLOCK_N), dtype) - A_shared = alloc_shared((BLOCK_M, BLOCK_K), a.dtype) - B_shared = alloc_shared((BLOCK_K, BLOCK_N), b.dtype) + # shared + fragment (regs) + A_shared = alloc_shared((BLOCK_M, BLOCK_K), a.dtype, 0) + B_shared = alloc_shared((BLOCK_K, BLOCK_N), b.dtype, 1) + C_local = alloc_fragment((TM, TY, TX, TN), dtypes.float32, 0, (1, 2), (ty, tx)) - # C_local = T.alloc_fragment((BLOCK_M, BLOCK_N), accum_dtype) -- an 8x4 REG tile per thread of the 8x16 grid - C_local = alloc_fragment((BLOCK_M, BLOCK_N), dtypes.float32, (0, 1), (ty, tx)) - - # T.clear(C_local) -- each thread zeroes its own fragment sub-tile - ic, jc = UOp.range(TM, 4, AxisType.LOOP), UOp.range(TN, 5, AxisType.UPCAST) - C_loc = C_local[ic*TM + ty, tx*TN + jc].set(0.0, end=(ic, jc)) + # zero out the regs to start. this is expanded by the devectorizer + C_local = C_local.after(C_local.store(0.0)) # for ko in T.Pipelined(T.ceildiv(K, BLOCK_K), num_stages=3): - # (num_stages pipelining is async copy + multi-buffering; this is the synchronous single-buffer version) ko = UOp.range(cdiv(K, BLOCK_K), 6, AxisType.LOOP) - # T.copy(A[by * BLOCK_M, ko * BLOCK_K], A_shared) -- each thread copies its own 8x4 sub-tile. - # Row index is iar*TM + ty (strided by TM across ty), matching tilelang's layout: thread ty owns - # rows {ty, ty+8, ..., ty+56} not {ty*8, ..., ty*8+7}. - iar, ka = UOp.range(TM, 7, AxisType.LOOP), UOp.range(TN, 8, AxisType.UPCAST) - A_store = A_shared[iar*TM + ty, tx*TN + ka].store(a[by*BLOCK_M + iar*TM + ty, ko*BLOCK_K + tx*TN + ka]).end(iar, ka) + # index the outer matrices + a = a.rearrange("(m bm) (k bk) -> m k bm bk", bm=BLOCK_M, bk=BLOCK_K)[by, ko] + b = b.rearrange("(k bk) (n bn) -> k n bk bn", bk=BLOCK_K, bn=BLOCK_N)[ko, bx] + c = c.rearrange("(m bm) (n bn) -> m n bm bn", bm=BLOCK_M, bn=BLOCK_N)[by, bx] - # T.copy(B[ko * BLOCK_K, bx * BLOCK_N], B_shared) - kb, ibr = UOp.range(TM, 9, AxisType.LOOP), UOp.range(TN, 10, AxisType.UPCAST) - B_store = B_shared[kb*TM + ty, tx*TN + ibr].store(b[ko*BLOCK_K + kb*TM + ty, bx*BLOCK_N + tx*TN + ibr]).end(kb, ibr) + # T.copy: A_shared <- a, B_shared <- b + def with_threads(x:UOp): return x.rearrange("(tm ty) (tx tn) -> ty tx tm tn", tm=TM, tn=TN)[ty, tx] + A_shared = A_shared.after(with_threads(A_shared).store(with_threads(a))) + B_shared = B_shared.after(with_threads(B_shared).store(with_threads(b))) - # get the shared after the stores (single barrier) - A_shared = A_shared.after(A_store, B_store) - B_shared = B_shared.after(A_store, B_store) - - # T.gemm(A_shared, B_shared, C_local), no WMMA -- per-thread accumulate over its fragment sub-tile. - # identical to custom_gemm: a self-referential store over the loop-carried kk range, - # which codegen turns into a register accumulator - # kk is the outer compute loop (axis 11) so that for each kk we read all 8 A rows and reuse - # the B[kk] read across them -- matching tilelang's ko > kk > row > col access order exactly. - kk, ir = UOp.range(BLOCK_K, 11, AxisType.LOOP), UOp.range(TM, 12, AxisType.LOOP) + # T.gemm(A_shared, B_shared, C_local), no WMMA + kk = UOp.range(BLOCK_K, 11, AxisType.LOOP) + ir = UOp.range(TM, 12, AxisType.LOOP) jj = UOp.range(TN, 13, AxisType.UPCAST) - acc = C_loc.after(kk)[ir*TM + ty, tx*TN + jj] + A_shared[ir*TM + ty, kk].cast(dtypes.float32) * B_shared[kk, tx*TN + jj].cast(dtypes.float32) + acc = C_local.after(kk)[ir, ty, tx, jj] + A_shared[ir*TM + ty, kk].cast(dtypes.float32) * B_shared[kk, tx*TN + jj].cast(dtypes.float32) # closing the ko loop here too; codegen adds the barrier so no thread overwrites the tiles while others still read them - C_loc = C_loc[ir*TM + ty, tx*TN + jj].set(acc, end=(kk, ir, jj, ko)) + C_local = C_local[ir, ty, tx, jj].set(acc, end=(kk, ir, jj, ko)) - # for i, j in T.Parallel(BLOCK_M, BLOCK_N): C_local[i, j] = T.max(C_local[i, j], 0) - # T.copy(C_local, C[by * BLOCK_M, bx * BLOCK_N]) -- per-thread store of the fragment shard (relu fused into it) - # LOOP: these loops are the per-thread output layout; convert_loop_to_global must not globalize them - ie, je = UOp.range(TM, 14, AxisType.LOOP), UOp.range(TN, 15, AxisType.UPCAST) - c_st = c[by*BLOCK_M + ie*TM + ty, bx*BLOCK_N + tx*TN + je].store(C_loc[ie*TM + ty, tx*TN + je].relu().cast(c.dtype)) + # c <- C_local (with relu and cast): every thread stores its shard's sub-view of the output tile + c_st = c.reshape(C_local.shape).store(C_local.relu().cast(c.dtype)) - # all open ranges are closed at the final store (ko was closed above). - # the fragment UNSHARDs go to codegen as is: multi_pm there resolves the full-tile view into per-thread shard code - return c_st.end(je, ie, tx, ty, bx, by).sink(arg=KernelInfo(name="matmul_relu", opts_to_apply=())) + # close the locals and globals + return c_st.end(tx, ty, bx, by).sink(arg=KernelInfo(name="matmul_relu", opts_to_apply=())) # --------------------------------------------------------------------------- # python wrapper: same signature as the tilelang function @@ -188,7 +131,8 @@ if __name__ == "__main__": b = Tensor.randn(K, N, dtype=dtype_in).contiguous() ref = (a @ b).relu().realize() - out = matmul_relu(a, b).realize() + for _ in range(10): + out = matmul_relu(a, b).realize() import numpy as np np.testing.assert_allclose(out.numpy(), ref.numpy(), atol=1e-1, rtol=1e-2) diff --git a/test/backend/test_custom_kernel.py b/test/backend/test_custom_kernel.py index 182e619899..7e087da373 100644 --- a/test/backend/test_custom_kernel.py +++ b/test/backend/test_custom_kernel.py @@ -1,5 +1,6 @@ import unittest from tinygrad import Tensor, UOp, GlobalCounters, Context, Device +import numpy as np from tinygrad.dtype import AddrSpace, dtypes, Invalid from tinygrad.uop.ops import KernelInfo, AxisType, Ops from tinygrad.renderer.ptx import PTXRenderer @@ -527,6 +528,77 @@ class TestUnshardIndex(unittest.TestCase): with self.assertRaisesRegex(RuntimeError, "cannot shard index"): self._run(kernel, (64, 8)) +def _run_fragment_kernel(testcase, kernel, out_shape, inputs=()): + c = Tensor.empty(*out_shape) + out = Tensor.custom_kernel(c, *inputs, fxn=kernel)[0] + try: return out.numpy() + except RuntimeError as e: + if isinstance(Device[Device.DEFAULT].renderer, PTXRenderer) and "dynamic register indexing" in str(e): + testcase.skipTest("PTX does not support dynamic register indexing") + raise + +class TestUnshardAlu(unittest.TestCase): + """Tests for ALU on (fragment) UNSHARD values in schedule/multi.py's alu_multi. + + An ALU with UNSHARD srcs lowers to per-shard ops when every src is one of: + same sharding: peel the UNSHARD, keep the layout + scalar: broadcast to every shard + whole unsharded same-shape value: takes its per-shard sub-view (shard_subview) + """ + @unittest.skipIf(not Device[Device.DEFAULT].renderer.has_local, "fragment tests need LOCAL ranges") + def test_alu_scalar_broadcast(self): + # scalar srcs broadcast to every shard: frag*2.0 where frag is 1.5 per thread -> 3.0 everywhere + def kernel(C:UOp) -> UOp: + ty = UOp.range(8, 0, AxisType.LOCAL) + # 8 values per thread, 8 threads -> 64-value full view + frag = UOp.placeholder((8,), dtypes.float32, 0, AddrSpace.LOCAL).unshard((0,), (ty,)) + v = frag.after(frag.store(1.5)) * 2.0 + return C.store(v).end(ty).sink(arg=KernelInfo(name="alu_scalar", opts_to_apply=())) + out = _run_fragment_kernel(self, kernel, (64,)) + np.testing.assert_allclose(out, 3.0) + + @unittest.skipIf(not Device[Device.DEFAULT].renderer.has_local, "fragment tests need LOCAL ranges") + def test_alu_whole_value_subview(self): + # UNSHARD + whole unsharded same-shape value: each shard adds its own sub-view of A. + def kernel(C:UOp, A:UOp) -> UOp: + ty = UOp.range(8, 0, AxisType.LOCAL) + frag = UOp.placeholder((8,), dtypes.float32, 0, AddrSpace.LOCAL).unshard((0,), (ty,)) + v = frag.after(frag.store(0.0)) + A + return C.store(v).end(ty).sink(arg=KernelInfo(name="alu_subview", opts_to_apply=())) + a = Tensor(np.arange(64, dtype=np.float32)) + out = _run_fragment_kernel(self, kernel, (64,), inputs=(a,)) + np.testing.assert_allclose(out, a.numpy(), atol=1e-4) + +class TestUnshardStore(unittest.TestCase): + """Tests for STORE of a sharded value into an unsharded dest (store_value_multi in schedule/multi.py). + + Every shard stores its value into its own contiguous sub-view of the dest, one SHRINK per sharded axis. + """ + @unittest.skipIf(not Device[Device.DEFAULT].renderer.has_local, "fragment tests need LOCAL ranges") + def test_store_unshard_value(self): + # single-axis: 8 threads each own 8 values of the 64-value output tile + def kernel(C:UOp) -> UOp: + ty = UOp.range(8, 0, AxisType.LOCAL) + frag = UOp.placeholder((8,), dtypes.float32, 0, AddrSpace.LOCAL).unshard((0,), (ty,)) + v = frag.after(frag.store(0.0)) + 2.5 + return C.store(v).end(ty).sink(arg=KernelInfo(name="store_unshard", opts_to_apply=())) + out = _run_fragment_kernel(self, kernel, (64,)) + np.testing.assert_allclose(out, 2.5) + + @unittest.skipIf(not Device[Device.DEFAULT].renderer.has_local, "fragment tests need LOCAL ranges") + def test_store_unshard_value_2axis(self): + # two sharded axes (the gemm fragment layout): thread (ty, tx) owns the (2, 1, 1, 2) sub-view of the + # (2, 4, 2, 2) output tile; the store must SHRINK dest on both sharded axes + def kernel(C:UOp, A:UOp) -> UOp: + ty = UOp.range(4, 0, AxisType.LOCAL) + tx = UOp.range(2, 1, AxisType.LOCAL) + frag = UOp.placeholder((2, 1, 1, 2), dtypes.float32, 0, AddrSpace.REG).unshard((1, 2), (ty, tx)) + v = frag.after(frag.store(0.0)) + A + return C.store(v).end(tx, ty).sink(arg=KernelInfo(name="store_unshard_2axis", opts_to_apply=())) + a = Tensor(np.arange(32, dtype=np.float32).reshape(2, 4, 2, 2)) + out = _run_fragment_kernel(self, kernel, (2, 4, 2, 2), inputs=(a,)) + np.testing.assert_allclose(out, a.numpy(), atol=1e-4) + class TestUOpReduce(unittest.TestCase): def test_uop_sum(self): a = Tensor([1.0, 2, 3, 4, 5]) diff --git a/tinygrad/schedule/multi.py b/tinygrad/schedule/multi.py index 275cffc95d..440b1565f9 100644 --- a/tinygrad/schedule/multi.py +++ b/tinygrad/schedule/multi.py @@ -71,13 +71,29 @@ def shard_srcs(msrcs:tuple[UOp, ...], axis:int) -> list[UOp]: srcs.append(full if axis in broadcast_axes(mlb.shape, out_shape) else full._shard(src_axis, sharding_rng)) return srcs +def shard_subview(full:UOp, multi:UOp) -> UOp: + """the sub-view of an unsharded full-shape value (shape == multi.shape) that belongs to this shard: + _shard along every sharded axis (contiguous blocks, like the device path).""" + assert tuple(full.shape) == tuple(multi.shape), f"shard sub-view shape mismatch {full.shape} != {multi.shape}" + # an EXPAND of a scalar over the full shape is the same broadcast on every shard: re-expand over the shard shape + if full.op is Ops.EXPAND and full.src[0].shape == (): return full.src[0].expand(multi.src[0].shape) + for ax, rng in multi.sharding: full = full._shard(ax, rng) + return full + def alu_multi(root:UOp): multis = [m for m in root.src if m.op is Ops.UNSHARD] if not multis: return None sharding = multis[0].sharding - if len(multis) == len(root.src) and all(m.sharding == sharding for m in multis): - srcs = [m.src[0] for m in root.src] - return srcs[0].alu(root.op, *srcs[1:]).unshard(multis[0].arg, multis[0].src[1:]) + target = multis[0] + def can_handle(m:UOp) -> bool: + # same sharding (peel the UNSHARD), or a whole unsharded value of the full tile shape (takes its per-shard + # sub-view), or a broadcast scalar + if m.sharding: return m.sharding == sharding + return m.shape == () or tuple(m.shape) == tuple(target.shape) + if all(can_handle(m) for m in root.src): + # every src either has the target sharding or is whole on every shard: run the alu per-shard + srcs = [m.src[0] if m.op is Ops.UNSHARD else m if m.shape == () else shard_subview(m, target) for m in root.src] + return srcs[0].alu(root.op, *srcs[1:]).unshard(target.arg, target.src[1:]) # resharding: single-axis fallback via shard_srcs axis = root.axis assert axis is not None @@ -233,6 +249,18 @@ def copy_multi(multi:UOp, device:str | tuple[str, ...]): def store_after_multi(dest:UOp, src:UOp): return dest.after(dest.store(src.src[0])).unshard(src.arg, src.src[1:]) +def store_value_multi(dest:UOp, multi:UOp): + # storing a sharded value into an unsharded dest: every shard stores into its own sub-view of the dest + return shard_subview(dest, multi).store(multi.src[0]) + +def store_dest_multi(root:UOp, multi:UOp): + # STORE with a sharded dest: every shard stores into its own shard of the dest. + # the value is handled like in alu_multi: UNSHARD srcs peel, full-shape values take their per-shard sub-view + # (scalars arrive EXPANDed to the full shape by UOp.store's const_like, so they sub-view like everything else) + srcs = [multi.src[0]] + [x.src[0] if x.op is Ops.UNSHARD else shard_subview(x, multi) if tuple(x.shape) == tuple(multi.shape) else x + for x in root.src[1:]] + return UOp(root.op, root.dtype, tuple(srcs), root.arg) + def passthrough_multi(root:UOp, multi:UOp): new_src = (multi.src[0],)+tuple(x.src[0] if x.op is Ops.UNSHARD else x for x in root.src[1:]) return UOp(root.op, root.dtype, src=new_src, arg=root.arg).unshard(multi.arg, multi.src[1:]) @@ -284,7 +312,8 @@ multi_pm = PatternMatcher([ UOp(root.op, root.dtype, tuple(x.src[0] if x.op is Ops.UNSHARD else x for x in root.src), root.arg)), (UPat((Ops.CAST, Ops.BITCAST, Ops.CONTIGUOUS, Ops.DETACH, Ops.CONTIGUOUS_BACKWARD), src=(UPat(Ops.UNSHARD, name="multi"), ), name="root"), passthrough_multi), - # remove UNSHARD from STORE - (UPat(Ops.STORE, src=(UPat(Ops.UNSHARD, name="multi"), ), name="root", allow_any_len=True), - lambda root,multi: UOp(root.op, root.dtype, (multi.src[0],)+tuple(x.src[0] if x.op is Ops.UNSHARD else x for x in root.src[1:]), root.arg)), + # STORE of a sharded value into an unsharded dest (e.g. a fragment into a full output tile) + (UPat(Ops.STORE, src=(UPat.var("dest"), UPat(Ops.UNSHARD, name="multi"))), store_value_multi), + # STORE into a sharded dest (e.g. the fragment init): every shard stores into its own shard + (UPat(Ops.STORE, src=(UPat(Ops.UNSHARD, name="multi"), ), name="root", allow_any_len=True), store_dest_multi), ])+replace_allreduce From a8c1e89500f83e7de2b6c09098453cc2c1999de7 Mon Sep 17 00:00:00 2001 From: qazal <77887910+Qazalin@users.noreply.github.com> Date: Fri, 31 Jul 2026 13:25:23 +0800 Subject: [PATCH 12/44] fp4 asm gemm 6+ pflops (#17315) * fp4 gemm * better kernargs structure * move to .s files * work * work * p2 * style * use .py * move to dsl * cleanup * add MFMA_SCALE_X2_ENCODING * cleanup mfma * fma docs * gemm_mxfp4 * more cleanup * move * move to cdna_asm_gemm * change * rm * change * mx --- extra/gemm/cdna_asm_gemm.py | 39 +- extra/gemm/gemm_mxfp4.py | 4124 +++++++++++++++++++++++++++++++++ test/backend/test_asm_gemm.py | 13 + 3 files changed, 4171 insertions(+), 5 deletions(-) create mode 100644 extra/gemm/gemm_mxfp4.py diff --git a/extra/gemm/cdna_asm_gemm.py b/extra/gemm/cdna_asm_gemm.py index cc7add0c52..31c75c9513 100644 --- a/extra/gemm/cdna_asm_gemm.py +++ b/extra/gemm/cdna_asm_gemm.py @@ -1,8 +1,9 @@ import atexit, functools, pathlib from tinygrad import Tensor, Device, dtypes +from tinygrad.dtype import AddrSpace from tinygrad.uop.ops import UOp, Ops, KernelInfo, AxisType from tinygrad.renderer import Estimates -from tinygrad.helpers import getenv, all_same, DEBUG +from tinygrad.helpers import getenv, all_same, DEBUG, ceildiv from tinygrad.runtime.support.compiler_amd import HIPCCCompiler from examples.mlperf.models.flat_llama import FP8_DTYPE, quantize_fp8 @@ -107,6 +108,23 @@ def custom_hk_mxfp8_gemm(C:UOp, A:UOp, B:UOp, scale_A:UOp, scale_B:UOp, *extra:U return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=(*sink.src, sink)), UOp(Ops.SOURCE, arg=src), UOp(Ops.BINARY, arg=lib))) +# ** MXFP4 GEMM custom kernel + +@functools.cache +def custom_mxfp4_gemm(C:UOp, A:UOp, B:UOp, scale_a:UOp, scale_b:UOp, tile_m:int, tile_n:int) -> UOp: + from extra.gemm.gemm_mxfp4 import build_kernel + M, half_k = A.shape[0]*A.shape[1], A.shape[2] + N, half_k_b = B.shape + K = half_k * 2 + assert half_k == half_k_b and C.shape == (*A.shape[:-1], N) + threads = UOp.special(256, "lidx0") + groups_x, groups_y = UOp.special(ceildiv(N, tile_n), "gidx0"), UOp.special(ceildiv(M, tile_m), "gidx1") + lds = UOp.placeholder((163840,), dtypes.uint8, 0, AddrSpace.LOCAL) + sink = UOp.sink(C.base, A.base, B.base, scale_a.base, scale_b.base, lds, threads, groups_x, groups_y, + arg=KernelInfo(f"custom_mxfp4_{M}_{N}_{K}", estimates=Estimates(ops=2*M*N*K))) + insts = build_kernel(M, N, K, tile_m, tile_n) + return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=tuple(UOp(Ops.INS, arg=x) for x in insts)))) + def quantize_mxfp8(x:Tensor) -> tuple[Tensor, Tensor, Tensor]: # 1x32 block scaling along the last axis *batch, K = x.shape @@ -144,7 +162,9 @@ atexit.register(_asm_gemm_report) def can_use_asm_gemm(a:Tensor, b:Tensor) -> bool: if a.dtype != b.dtype: return todo(f"dtypes must match {a.dtype} != {b.dtype}") - if a.dtype not in {dtypes.bfloat16, dtypes.float16, FP8_DTYPE}: return todo(f"only bfloat16/float16/fp8, got {a.dtype}") + # fp4 encoded as packed uint8 + # TODO: add fp4 dtype? + if a.dtype not in {dtypes.bfloat16, dtypes.float16, FP8_DTYPE, dtypes.uint8}: return todo(f"only bfloat16/float16/fp8/fp4, got {a.dtype}") batch, M, K = (1, *a.shape) if a.ndim == 2 else a.shape N = b.shape[1] if isinstance(a.device, tuple): @@ -244,7 +264,6 @@ def hk_bf16_atb_gemm(a:Tensor, b:Tensor) -> Tensor: if reduce_out: out = out.sum(0) return out.squeeze(0) if out.ndim == 3 else out - # ** backward gemm, might use the asm gemm def custom_gemm_bw(gradient:UOp, kernel:UOp, n_scales:int=2, has_grad_amax:bool=False, has_w_post:bool=False): @@ -348,6 +367,12 @@ def asm_gemm(a:Tensor, b:Tensor, x_scale:Tensor|None=None, w_scale:Tensor|None=N w_post_scale:Tensor|None=None, mx:bool=False, mx_scales:tuple|None=None, mx_w_stored:bool=False, g_amax:Tensor|None=None, a_pretranspose:Tensor|None=None) -> Tensor: assert can_use_asm_gemm(a, b), f"{counters['todos'][-1]}" + if (mxfp4:=a.dtype == dtypes.uint8): + assert mx_scales is not None and len(mx_scales) == 2 + scale_a, scale_b = mx_scales + K = a.shape[-1] * 2 + assert scale_a.shape == (*a.shape[:-1], K // 32) and scale_b.shape == (b.shape[1], K // 32) + assert scale_a.dtype == scale_b.dtype == dtypes.uint8 and a.device == b.device == scale_a.device == scale_b.device counters["used"] += 1 unfold_batch = a.ndim == 3 and isinstance(a.device, tuple) and a.uop.axis == 2 and b.uop.axis == 0 if unfold_batch: @@ -355,7 +380,7 @@ def asm_gemm(a:Tensor, b:Tensor, x_scale:Tensor|None=None, w_scale:Tensor|None=N a = a.reshape(a.shape[0]*a.shape[1], a.shape[2]) squeeze = a.ndim == 2 if squeeze: a = a.unsqueeze(0) - out_dtype = dtypes.bfloat16 if a.dtype == FP8_DTYPE else a.dtype + out_dtype = dtypes.bfloat16 if a.dtype == FP8_DTYPE or mxfp4 else a.dtype batch, M, K = a.shape N = b.shape[1] @@ -378,7 +403,11 @@ def asm_gemm(a:Tensor, b:Tensor, x_scale:Tensor|None=None, w_scale:Tensor|None=N renderer = Device[dname:=(a.device[0] if is_multi else a.device)].renderer dname, arch = dname.split(":")[0], renderer.target.arch if arch.startswith("gfx950") and getenv("USE_ASM", 1): - if mx: + if mxfp4: + tile_m, tile_n = next((tm, tn) for tm, tn in ((256, 256), (192, 256), (128, 512)) if (batch*M) % tm == N % tn == 0) + fxn = functools.partial(custom_mxfp4_gemm, tile_m=tile_m, tile_n=tile_n) + out = Tensor.custom_kernel(out, a, b.T, scale_a, scale_b, fxn=fxn)[0] + elif mx: # mxfp8 1x32 block scaling if mx_scales is not None: a_si, a_e8, b_si, b_e8 = mx_scales diff --git a/extra/gemm/gemm_mxfp4.py b/extra/gemm/gemm_mxfp4.py new file mode 100644 index 0000000000..04d8195eaa --- /dev/null +++ b/extra/gemm/gemm_mxfp4.py @@ -0,0 +1,4124 @@ +# ruff: noqa: E501,F403,F405 +from tinygrad.runtime.autogen.amd.cdna.ins import * + +class Kernel: + def __init__(self): self.instructions, self.labels, self.pos = [], {}, 0 + def label(self, name): self.labels[name] = self.pos + def emit(self, inst, target=None): + self.instructions.append(inst) + inst._target, inst._pos = target, self.pos + self.pos += inst.size() + def finalize(self): + for inst in self.instructions: + if inst._target is not None: inst.simm16 = (self.labels[inst._target] - inst._pos - inst.size()) // 4 + return self.instructions + +def v_mfma_fp4(dst, a, b, opsel, opsel_hi, scale_a, scale_b): + # select fp4 for both inputs, 0xD3AC is the load scale encoding and write to acc vgprs + return v_mfma_scale_f32_16x16x128_f8f6f4(dst, a, b, dst, 0, 0, opsel, opsel_hi, 4, 1, 1, 0, 4, 0xD3AC, scale_a.offset, scale_b.offset) + +def build_kernel(M: int, N: int, K: int, tile_m: int, tile_n: int): + k = Kernel() + scale_k = K // 32 + if (tile_m, tile_n) == (128, 512): + k.emit(s_and_b32(s[1], s[1], LIT, 65535)) + k.emit(s_mov_b32(s[47], s[2])) + k.emit(s_mov_b32(s[48], s[3])) + k.emit(s_mov_b32(s[64], s[4])) + k.emit(s_load_dwordx2(s[4:5], s[0:1], s[0], 0, 0, 0, 0, 1)) + k.emit(s_mov_b32(s[8], 0)) + k.emit(s_mov_b32(s[9], 0)) + k.emit(s_load_dwordx2(s[12:13], s[0:1], s[0], 8, 0, 0, 0, 1)) + k.emit(s_load_dwordx2(s[16:17], s[0:1], s[0], 16, 0, 0, 0, 1)) + k.emit(s_mov_b32(s[41], 1.0)) + k.emit(s_mov_b32(s[42], 0)) + k.emit(s_mov_b32(s[36], N)) + k.emit(s_mov_b32(s[37], K)) + k.emit(s_mov_b32(s[38], K)) + k.emit(s_mov_b32(s[43], M)) + k.emit(s_mov_b32(s[44], N)) + k.emit(s_mov_b32(s[45], K)) + k.emit(s_load_dwordx2(s[20:21], s[0:1], s[0], 24, 0, 0, 0, 1)) + k.emit(s_load_dwordx2(s[24:25], s[0:1], s[0], 32, 0, 0, 0, 1)) + k.emit(s_mov_b32(s[39], scale_k)) + k.emit(s_mov_b32(s[40], scale_k)) + k.emit(s_mov_b32(s[65], 0)) + k.emit(v_lshrrev_b32_e32(v[1], 10)) + k.emit(v_lshrrev_b32_e32(v[2], 10, v[1])) + k.emit(v_and_b32_e32(v[2], LIT, v[2], 1023)) + k.emit(v_and_b32_e32(v[1], LIT, v[1], 1023)) + k.emit(v_and_b32_e32(v[0], LIT, v[0], 1023)) + k.emit(v_lshrrev_b32_e32(v[3], 6)) + k.emit(v_and_b32_e32(v[0], 63)) + k.emit(v_readfirstlane_b32_e32(v[46], v[3])) + k.emit(s_waitcnt(49279)) + for i in range(2): + k.emit(s_mov_b32(s[6 + i * 8], -16)) + k.emit(s_mov_b32(s[10 + i * 12], -16)) + k.emit(s_mov_b32(s[18 + i * 8], -16)) + for i in range(2): + k.emit(s_mov_b32(s[7 + i * 8], LIT, 131072)) + k.emit(s_mov_b32(s[11 + i * 12], LIT, 131072)) + k.emit(s_mov_b32(s[19 + i * 8], LIT, 131072)) + for i in range(2): + k.emit(s_and_b32(s[5 + i * 8], s[5 + i * 8], LIT, 65535)) + k.emit(s_and_b32(s[9 + i * 12], s[9 + i * 12], LIT, 65535)) + k.emit(s_and_b32(s[17 + i * 8], s[17 + i * 8], LIT, 65535)) + for i in range(2): + k.emit(s_or_b32(s[5 + i * 8], s[5 + i * 8], LIT, 262144)) + k.emit(s_or_b32(s[9 + i * 12], s[9 + i * 12], LIT, 262144)) + k.emit(s_or_b32(s[17 + i * 8], s[17 + i * 8], LIT, 262144)) + k.emit(s_cmp_gt_i32(s[65], 0)) + k.emit(s_cbranch_scc0(9), target='L0_0194') + k.emit(s_lshr_b32(s[66], s[45], s[65])) + k.emit(s_add_u32(s[66], s[66], LIT, 255)) + k.emit(s_lshr_b32(s[66], s[66], 8)) + k.emit(s_lshl_b32(s[66], s[66], 8)) + k.emit(s_mul_i32(s[63], s[66], s[64])) + k.emit(s_sub_i32(s[62], s[45], s[63])) + k.emit(s_cmp_lt_i32(s[62], s[66])) + k.emit(s_cselect_b32(s[45], s[62], s[66])) + k.label('L0_0194') + k.emit(s_lshr_b32(s[37], s[37], 1)) + k.emit(s_mul_i32(s[62], s[48], LIT, 128)) + k.emit(s_mul_hi_u32(s[63], s[37], s[62])) + k.emit(s_add_u32(s[13], s[13], s[63])) + k.emit(s_mul_i32(s[63], s[37], s[62])) + k.emit(s_add_u32(s[12], s[12], s[63])) + k.emit(s_addc_u32(s[13], s[13], 0)) + k.emit(s_sub_i32(s[63], s[43], s[62])) + k.emit(s_cmp_lt_u32(s[63], LIT, 128)) + k.emit(s_cselect_b32(s[62], s[63], LIT, 128)) + k.emit(s_mul_i32(s[63], s[37], s[62])) + k.emit(s_mov_b32(s[14], s[63])) + k.emit(s_mov_b32(s[15], LIT, 131072)) + k.emit(s_cmp_gt_i32(s[65], 0)) + k.emit(s_cbranch_scc0(5), target='L0_01F4') + k.emit(s_mul_i32(s[63], s[66], s[64])) + k.emit(s_lshr_b32(s[62], s[63], 1)) + k.emit(s_add_u32(s[12], s[12], s[62])) + k.emit(s_addc_u32(s[13], s[13], 0)) + k.emit(s_sub_u32(s[14], s[14], s[62])) + k.label('L0_01F4') + k.emit(v_lshrrev_b32_e32(v[4], 3)) + k.emit(v_lshrrev_b32_e32(v[5], 2, v[4])) + k.emit(v_lshlrev_b32_e32(v[5], 4, v[5])) + k.emit(v_and_b32_e32(v[4], 3, v[4])) + k.emit(v_lshrrev_b32_e32(v[6], 1, v[4])) + k.emit(v_lshlrev_b32_e32(v[6], 2, v[6])) + k.emit(v_add_u32_e32(v[5], v[5], v[6])) + k.emit(v_and_b32_e32(v[4], 1, v[4])) + k.emit(v_add_u32_e32(v[5], v[5], v[4])) + k.emit(v_mul_lo_u32(v[212], s[37], v[5])) + k.emit(v_and_b32_e32(v[4], 7)) + k.emit(v_lshlrev_b32_e32(v[4], 4, v[4])) + k.emit(v_add_u32_e32(v[212], v[4], v[212])) + k.emit(s_lshr_b32(s[62], s[46], 1)) + k.emit(s_mul_i32(s[62], s[62], 8)) + k.emit(s_and_b32(s[63], s[46], 1)) + k.emit(s_mul_i32(s[63], s[63], 2)) + k.emit(s_add_u32(s[62], s[62], s[63])) + k.emit(s_mul_i32(s[62], s[37], s[62])) + k.emit(v_add_u32_e32(v[212], s[62], v[212])) + k.emit(s_mul_i32(s[62], s[37], 32)) + k.emit(v_add_u32_e32(v[213], s[62], v[212])) + k.emit(v_add_u32_e32(v[214], s[62], v[213])) + k.emit(v_add_u32_e32(v[215], s[62], v[214])) + k.emit(s_mul_i32(s[67], LIT, s[46], 1056)) + k.emit(s_add_u32(s[67], LIT, s[67], 2048)) + k.emit(v_and_b32_e32(v[4], 15)) + k.emit(v_lshrrev_b32_e32(v[5], 3, v[4])) + k.emit(v_mul_i32_i24_e32(v[5], 2, v[5])) + k.emit(v_and_b32_e32(v[4], 3)) + k.emit(v_lshrrev_b32_e32(v[6], 1, v[4])) + k.emit(v_add_u32_e32(v[4], v[5], v[6])) + k.emit(v_mul_i32_i24_e32(v[216], LIT, v[4], 1056)) + k.emit(v_and_b32_e32(v[4], 7)) + k.emit(v_lshrrev_b32_e32(v[5], 2, v[4])) + k.emit(v_mul_i32_i24_e32(v[5], LIT, v[5], 256)) + k.emit(v_add_u32_e32(v[216], v[5], v[216])) + k.emit(v_and_b32_e32(v[4], 1)) + k.emit(v_mul_i32_i24_e32(v[6], LIT, v[4], 128)) + k.emit(v_add_u32_e32(v[216], v[6], v[216])) + k.emit(v_lshrrev_b32_e32(v[4], 4)) + k.emit(v_mul_i32_i24_e32(v[4], 16, v[4])) + k.emit(v_add_u32_e32(v[216], v[4], v[216])) + k.emit(s_mov_b32(s[62], LIT, 2048)) + k.emit(v_add_u32_e64(v[216], v[216], s[62])) + k.emit(v_add_u32_e32(v[217], LIT, v[216], 16896)) + k.emit(s_mul_i32(s[62], s[48], LIT, 128)) + k.emit(s_mul_hi_u32(s[63], s[39], s[62])) + k.emit(s_add_u32(s[21], s[21], s[63])) + k.emit(s_mul_i32(s[63], s[39], s[62])) + k.emit(s_add_u32(s[20], s[20], s[63])) + k.emit(s_addc_u32(s[21], s[21], 0)) + k.emit(s_add_u32(s[63], s[43], 31)) + k.emit(s_lshr_b32(s[63], s[63], 5)) + k.emit(s_lshl_b32(s[63], s[63], 5)) + k.emit(s_sub_i32(s[63], s[63], s[62])) + k.emit(s_cmp_lt_u32(s[63], LIT, 128)) + k.emit(s_cselect_b32(s[62], s[63], LIT, 128)) + k.emit(s_mul_i32(s[63], s[39], s[62])) + k.emit(s_mov_b32(s[22], s[63])) + k.emit(s_mov_b32(s[23], LIT, 131072)) + k.emit(s_cmp_gt_i32(s[65], 0)) + k.emit(s_cbranch_scc0(4), target='L0_0334') + k.emit(s_mul_i32(s[63], s[66], s[64])) + k.emit(s_add_u32(s[20], s[20], s[63])) + k.emit(s_addc_u32(s[21], s[21], 0)) + k.emit(s_sub_u32(s[22], s[22], s[63])) + k.label('L0_0334') + k.emit(v_lshlrev_b32_e32(v[218], 2)) + k.emit(s_mul_i32(s[63], s[46], 32)) + k.emit(s_mul_i32(s[63], s[63], s[39])) + k.emit(v_add_u32_e32(v[218], s[63], v[218])) + k.emit(s_mul_i32(s[68], s[46], LIT, 256)) + k.emit(s_add_i32(s[68], s[68], 0)) + k.emit(v_lshlrev_b32_e32(v[219], 2)) + k.emit(v_add_u32_e32(v[219], 0, v[219])) + k.emit(s_lshr_b32(s[38], s[38], 1)) + k.emit(s_mul_i32(s[62], s[47], LIT, 512)) + k.emit(s_mul_hi_u32(s[63], s[38], s[62])) + k.emit(s_add_u32(s[17], s[17], s[63])) + k.emit(s_mul_i32(s[63], s[38], s[62])) + k.emit(s_add_u32(s[16], s[16], s[63])) + k.emit(s_addc_u32(s[17], s[17], 0)) + k.emit(s_sub_i32(s[63], s[44], s[62])) + k.emit(s_cmp_lt_u32(s[63], LIT, 512)) + k.emit(s_cselect_b32(s[62], s[63], LIT, 512)) + k.emit(s_mul_i32(s[63], s[38], s[62])) + k.emit(s_mov_b32(s[18], s[63])) + k.emit(s_mov_b32(s[19], LIT, 131072)) + k.emit(s_cmp_gt_i32(s[65], 0)) + k.emit(s_cbranch_scc0(6), target='L0_03BC') + k.emit(s_mul_i32(s[63], s[66], s[64])) + k.emit(s_lshr_b32(s[62], s[63], 1)) + k.emit(s_mul_i32(s[62], s[62], 16)) + k.emit(s_add_u32(s[16], s[16], s[62])) + k.emit(s_addc_u32(s[17], s[17], 0)) + k.emit(s_sub_u32(s[18], s[18], s[62])) + k.label('L0_03BC') + k.emit(v_lshlrev_b32_e32(v[220], 4)) + k.emit(s_mul_i32(s[63], s[46], LIT, 128)) + k.emit(s_mul_i32(s[62], s[63], s[38])) + k.emit(v_add_u32_e32(v[220], s[62], v[220])) + k.emit(s_mul_i32(s[62], 16, s[38])) + k.emit(v_add_u32_e32(v[221], s[62], v[220])) + k.emit(v_add_u32_e32(v[222], s[62], v[221])) + k.emit(v_add_u32_e32(v[223], s[62], v[222])) + for i in range(4): + k.emit(v_add_u32_e32(v[224 + i * 1], LIT, v[220 + i * 1], 1024)) + k.emit(s_mul_i32(s[62], 64, s[38])) + for i in range(8): + k.emit(v_add_u32_e32(v[228 + i * 1], s[62], v[220 + i * 1])) + k.emit(s_mul_i32(s[62], s[47], LIT, 512)) + k.emit(s_mul_hi_u32(s[63], s[40], s[62])) + k.emit(s_add_u32(s[25], s[25], s[63])) + k.emit(s_mul_i32(s[63], s[40], s[62])) + k.emit(s_add_u32(s[24], s[24], s[63])) + k.emit(s_addc_u32(s[25], s[25], 0)) + k.emit(s_sub_i32(s[63], s[44], s[62])) + k.emit(s_cmp_lt_u32(s[63], LIT, 512)) + k.emit(s_cselect_b32(s[62], s[63], LIT, 512)) + k.emit(s_mul_i32(s[63], s[40], s[62])) + k.emit(s_mov_b32(s[26], s[63])) + k.emit(s_mov_b32(s[27], LIT, 131072)) + k.emit(s_cmp_gt_i32(s[65], 0)) + k.emit(s_cbranch_scc0(4), target='L0_047C') + k.emit(s_mul_i32(s[63], s[66], s[64])) + k.emit(s_add_u32(s[24], s[24], s[63])) + k.emit(s_addc_u32(s[25], s[25], 0)) + k.emit(s_sub_u32(s[26], s[26], s[63])) + k.label('L0_047C') + k.emit(v_lshlrev_b32_e32(v[236], 2)) + k.emit(s_mul_i32(s[63], s[46], LIT, 128)) + k.emit(s_mul_i32(s[63], s[63], s[40])) + k.emit(v_add_u32_e32(v[236], s[63], v[236])) + k.emit(s_mul_i32(s[62], 32, s[40])) + k.emit(v_add_u32_e32(v[237], s[62], v[236])) + k.emit(v_add_u32_e32(v[238], s[62], v[237])) + k.emit(v_add_u32_e32(v[239], s[62], v[238])) + k.emit(s_mov_b32(s[69], LIT, 128)) + k.emit(s_mov_b32(s[70], LIT, 2048)) + k.emit(s_mov_b32(s[71], LIT, 256)) + k.emit(s_mov_b32(s[72], LIT, 256)) + k.emit(s_mov_b32(s[60], 0)) + k.emit(s_mov_b32(s[61], s[45])) + k.emit(s_add_u32(NULL, 0, s[67])) + k.emit(buffer_load_dwordx4(v[0:3], v[212], s[12:15], 0, 0, 1, 0, 0, 0, 0, 1)) + for i in range(8): + k.emit(v_accvgpr_write(v[0 + i * 1], 0)) + k.emit(s_add_u32(NULL, LIT, s[67], 4224)) + k.emit(buffer_load_dwordx4(v[0:3], v[213], s[12:15], 0, 0, 1, 0, 0, 0, 0, 1)) + for i in range(8): + k.emit(v_accvgpr_write(v[8 + i * 1], 0)) + k.emit(s_add_u32(NULL, 0, s[68])) + k.emit(buffer_load_dword(v[0], v[218], s[20:23], 0, 0, 1, 0, 0, 0, 0, 1)) + for i in range(2): + for j0 in range(8): + k.emit(v_accvgpr_write(v[16 + j0 * 1 + i * 8], 0)) + k.emit(s_add_u32(NULL, LIT, s[67], 8448 + i * 4224)) + k.emit(buffer_load_dwordx4(v[0:3], v[214 + i * 1], s[12:15], 0, 0, 1, 0, 0, 0, 0, 1)) + for i in range(8): + k.emit(v_accvgpr_write(v[32 + i * 1], 0)) + k.emit(s_add_u32(s[62], LIT, s[60], 256)) + k.emit(s_cmp_lt_u32(s[62], s[61])) + k.emit(s_cselect_b32(s[69], s[69], 0)) + k.emit(s_cselect_b32(s[71], s[71], 0)) + for i in range(2): + k.emit(s_add_u32(s[12 + i * 8], s[12 + i * 8], s[69 + i * 2])) + k.emit(s_addc_u32(s[13 + i * 8], 0, s[13 + i * 8])) + k.emit(s_sub_u32(s[14 + i * 8], s[14 + i * 8], s[69 + i * 2])) + for i in range(2): + for j1 in range(8): + k.emit(buffer_load_dwordx4(v[72 + j1 * 4 + i * 32:75 + j1 * 4 + i * 32], v[220 + j1 * 1 + i * 8], s[16:19], 0, 0, 1)) + k.emit(v_accvgpr_write(v[40 + j1 * 8 + i * 80], 0)) + k.emit(v_accvgpr_write(v[41 + j1 * 8 + i * 80], 0)) + k.emit(v_accvgpr_write(v[42 + j1 * 8 + i * 80], 0)) + k.emit(v_accvgpr_write(v[43 + j1 * 8 + i * 80], 0)) + k.emit(v_accvgpr_write(v[44 + j1 * 8 + i * 80], 0)) + k.emit(v_accvgpr_write(v[45 + j1 * 8 + i * 80], 0)) + k.emit(v_accvgpr_write(v[46 + j1 * 8 + i * 80], 0)) + k.emit(v_accvgpr_write(v[47 + j1 * 8 + i * 80], 0)) + for j2 in range(2): + k.emit(buffer_load_dword(v[204 + j2 * 1 + i * 2], v[236 + j2 * 1 + i * 2], s[24:27], 0, 0, 1)) + k.emit(v_accvgpr_write(v[104 + j2 * 8 + i * 80], 0)) + k.emit(v_accvgpr_write(v[105 + j2 * 8 + i * 80], 0)) + k.emit(v_accvgpr_write(v[106 + j2 * 8 + i * 80], 0)) + k.emit(v_accvgpr_write(v[107 + j2 * 8 + i * 80], 0)) + k.emit(v_accvgpr_write(v[108 + j2 * 8 + i * 80], 0)) + k.emit(v_accvgpr_write(v[109 + j2 * 8 + i * 80], 0)) + k.emit(v_accvgpr_write(v[110 + j2 * 8 + i * 80], 0)) + k.emit(v_accvgpr_write(v[111 + j2 * 8 + i * 80], 0)) + k.emit(s_add_u32(s[63], LIT, s[60], 256)) + k.emit(s_cmp_lt_u32(s[63], s[61])) + k.emit(s_cselect_b32(s[70], s[70], 0)) + k.emit(s_cselect_b32(s[72], s[72], 0)) + for i in range(2): + k.emit(s_add_u32(s[16 + i * 8], s[16 + i * 8], s[70 + i * 2])) + k.emit(s_addc_u32(s[17 + i * 8], 0, s[17 + i * 8])) + k.emit(s_sub_u32(s[18 + i * 8], s[18 + i * 8], s[70 + i * 2])) + for i in range(2): + k.emit(s_add_u32(NULL, LIT, s[67], 16896 + i * 4224)) + k.emit(buffer_load_dwordx4(v[0:3], v[212 + i * 1], s[12:15], 0, 0, 1, 0, 0, 0, 0, 1)) + for j3 in range(8): + k.emit(v_accvgpr_write(v[200 + j3 * 1 + i * 8], 0)) + k.emit(s_add_u32(NULL, LIT, s[68], 1024)) + k.emit(buffer_load_dword(v[0], v[218], s[20:23], 0, 0, 1, 0, 0, 0, 0, 1)) + for i in range(2): + for j4 in range(8): + k.emit(v_accvgpr_write(v[216 + j4 * 1 + i * 8], 0)) + k.emit(s_add_u32(NULL, LIT, s[67], 25344 + i * 4224)) + k.emit(buffer_load_dwordx4(v[0:3], v[214 + i * 1], s[12:15], 0, 0, 1, 0, 0, 0, 0, 1)) + for i in range(24): + k.emit(v_accvgpr_write(v[232 + i * 1], 0)) + k.emit(s_waitcnt(20347)) + k.emit(s_barrier()) + k.emit(ds_read_b128(v[8:11], v[216])) + k.emit(ds_read_b128(v[24:27], v[216], v[0], v[0], 0, 64)) + for i in range(2): + k.emit(ds_read_b128(v[12 + i * 20:15 + i * 20], v[216], v[0], v[0], 0, 0 + i * 192, 2 + i * 14)) + k.emit(ds_read_b128(v[28 + i * -8:31 + i * -8], v[216], v[0], v[0], 0, 64 + i * 64, 2 + i * 16)) + k.emit(ds_read_b128(v[16 + i * 20:19 + i * 20], v[216], v[0], v[0], 0, 128 + i * 64, 16 + i * 2)) + k.emit(ds_read_b32(v[200], v[219])) + k.emit(ds_read_b32(v[201], v[219], v[0], v[0], 0, 0, 1)) + for i in range(5): + k.emit(s_nop()) + k.emit(s_lshl_b32(s[36], s[36], 1)) + k.emit(s_mul_i32(s[62], s[48], LIT, 128)) + k.emit(s_mul_hi_u32(s[63], s[36], s[62])) + k.emit(s_add_u32(s[5], s[5], s[63])) + k.emit(s_mul_i32(s[63], s[36], s[62])) + k.emit(s_add_u32(s[4], s[4], s[63])) + k.emit(s_addc_u32(s[5], s[5], 0)) + k.emit(s_mul_i32(s[63], s[47], LIT, 512)) + k.emit(s_lshl_b32(s[63], s[63], 1)) + k.emit(s_add_u32(s[4], s[4], s[63])) + k.emit(s_addc_u32(s[5], s[5], 0)) + k.emit(s_sub_i32(s[62], s[43], s[62])) + k.emit(s_cmp_lt_u32(s[62], LIT, 128)) + k.emit(s_cselect_b32(s[62], s[62], LIT, 128)) + k.emit(s_mul_i32(s[62], s[36], s[62])) + k.emit(s_sub_i32(s[62], s[62], s[63])) + k.emit(s_mov_b32(s[6], s[62])) + k.emit(s_mov_b32(s[7], LIT, 131072)) + k.emit(s_cmp_gt_i32(s[65], 0)) + k.emit(s_cbranch_scc0(11), target='L0_0F54') + k.emit(v_mul_i32_i24_e64(v[4], v[0], 4)) + k.emit(s_mul_i32(s[62], s[46], LIT, 256)) + k.emit(v_add_u32_e32(v[240], s[62], v[4])) + k.emit(v_add_u32_e32(v[241], LIT, v[240], 128)) + k.emit(s_mul_i32(s[62], s[36], 64)) + k.emit(v_add_u32_e32(v[242], s[62], v[240])) + k.emit(v_add_u32_e32(v[243], s[62], v[241])) + k.emit(s_branch(22), target='L0_0FAC') + k.label('L0_0F54') + k.emit(v_and_b32_e64(v[4], v[0], 15)) + k.emit(v_mul_lo_u32(v[240], s[36], v[4])) + k.emit(v_lshrrev_b32_e32(v[4], 5)) + k.emit(v_mul_i32_i24_e64(v[4], v[4], 16)) + k.emit(v_lshrrev_b32_e32(v[5], 4)) + k.emit(v_and_b32_e64(v[5], v[5], 1)) + k.emit(v_mul_i32_i24_e64(v[5], v[5], 32)) + k.emit(v_add_u32_e32(v[4], v[4], v[5])) + k.emit(v_add_u32_e32(v[240], v[4], v[240])) + k.emit(s_mul_i32(s[62], s[46], LIT, 256)) + k.emit(v_add_u32_e32(v[240], s[62], v[240])) + k.emit(v_add_u32_e32(v[241], LIT, v[240], 128)) + k.emit(s_mul_i32(s[62], s[36], 64)) + k.emit(v_add_u32_e32(v[242], s[62], v[240])) + k.emit(v_add_u32_e32(v[243], s[62], v[241])) + k.label('L0_0FAC') + k.emit(s_cmp_lt_i32(s[46], 2)) + k.emit(s_cbranch_scc0(1283), target='L0_23C0') + k.label('L0_0FB4') + k.emit(s_waitcnt(127)) + k.emit(v_mfma_fp4(v[0:3], v[72:75], v[8:11], 0, 0, v[204], v[200])) + k.emit(s_barrier()) + for i in range(2): + k.emit(v_mfma_fp4(v[4 + i * 32:7 + i * 32], v[72 + i * 8:75 + i * 8], v[12:15], 2, 0, v[204 + i * 1], v[200])) + k.emit(buffer_load_dwordx4(v[136 + i * 8:139 + i * 8], v[220 + i * 2], s[16:19], 0, 0, 1)) + for j5 in range(2): + k.emit(v_mfma_fp4(v[16 + j5 * -8 + i * 32:19 + j5 * -8 + i * 32], v[76 + j5 * -4 + i * 8:79 + j5 * -4 + i * 8], v[8 + j5 * 8:11 + j5 * 8], 1 + j5 * -1, 0, v[204 + i * 1], v[200 + j5 * 1])) + k.emit(ds_read_b128(v[40 + j5 * 16 + i * 8:43 + j5 * 16 + i * 8], v[216], v[0], v[0], 0, 0 + j5 * 64 + i * 128, 33 + i * 16)) + k.emit(v_mfma_fp4(v[20 + j5 * -8 + i * 32:23 + j5 * -8 + i * 32], v[76 + j5 * -4 + i * 8:79 + j5 * -4 + i * 8], v[12 + j5 * 8:15 + j5 * 8], 3 + j5 * -1, 0, v[204 + i * 1], v[200 + j5 * 1])) + k.emit(buffer_load_dwordx4(v[140 + i * 8:143 + i * 8], v[221 + i * 2], s[16:19], 0, 0, 1)) + k.emit(v_mfma_fp4(v[24 + i * 32:27 + i * 32], v[76 + i * 8:79 + i * 8], v[16:19], 1, 0, v[204 + i * 1], v[201])) + k.emit(ds_read_b128(v[44 + i * 8:47 + i * 8], v[216], v[0], v[0], 0, 0 + i * 128, 35 + i * 16)) + k.emit(v_mfma_fp4(v[28 + i * 32:31 + i * 32], v[76 + i * 8:79 + i * 8], v[20:23], 3, 0, v[204 + i * 1], v[201])) + k.emit(v_mfma_fp4(v[32 + i * -32:35 + i * -32], v[80 + i * 8:83 + i * 8], v[8 + i * 16:11 + i * 16], 0, 0 + i * 3, v[205 + i * -1], v[200])) + k.emit(ds_read_b128(v[60 + i * 8:63 + i * 8], v[216], v[0], v[0], 0, 64 + i * 128, 35 + i * 16)) + k.emit(v_mfma_fp4(v[4:7], v[88:91], v[28:31], 2, 3, v[204], v[200])) + k.emit(buffer_load_dwordx4(v[152:155], v[224], s[16:19], 0, 0, 1)) + for i in range(2): + k.emit(v_mfma_fp4(v[16 + i * -8:19 + i * -8], v[92 + i * -4:95 + i * -4], v[24 + i * 8:27 + i * 8], 1 + i * -1, 3, v[204], v[200 + i * 1])) + k.emit(ds_read_b32(v[202 + i * 1], v[219], v[0], v[0], 0, 0, 2 + i * 1)) + k.emit(v_mfma_fp4(v[20 + i * -8:23 + i * -8], v[92 + i * -4:95 + i * -4], v[28 + i * 8:31 + i * 8], 3 + i * -1, 3, v[204], v[200 + i * 1])) + for i in range(2): + k.emit(buffer_load_dwordx4(v[156 + i * 4:159 + i * 4], v[225 + i * 1], s[16:19], 0, 0, 1)) + k.emit(v_mfma_fp4(v[24 + i * 24:27 + i * 24], v[92 + i * 8:95 + i * 8], v[32 + i * -8:35 + i * -8], 1, 3, v[204 + i * 1], v[201 + i * -1])) + k.emit(v_mfma_fp4(v[28 + i * 24:31 + i * 24], v[92 + i * 8:95 + i * 8], v[36 + i * -8:39 + i * -8], 3, 3, v[204 + i * 1], v[201 + i * -1])) + k.emit(v_mfma_fp4(v[32 + i * 8:35 + i * 8], v[96:99], v[24 + i * 8:27 + i * 8], 0, 3, v[205], v[200 + i * 1])) + k.emit(v_mfma_fp4(v[36 + i * 8:39 + i * 8], v[96:99], v[28 + i * 8:31 + i * 8], 2, 3, v[205], v[200 + i * 1])) + k.emit(buffer_load_dwordx4(v[164:167], v[227], s[16:19], 0, 0, 1)) + k.emit(v_mfma_fp4(v[56:59], v[100:103], v[32:35], 1, 3, v[205], v[201])) + k.emit(v_mfma_fp4(v[60:63], v[100:103], v[36:39], 3, 3, v[205], v[201])) + k.emit(s_waitcnt(3965)) + k.emit(v_mfma_fp4(v[64:67], v[104:107], v[8:11], 0, 0, v[206], v[200])) + k.emit(s_add_u32(s[62], LIT, s[60], 512)) + k.emit(v_mfma_fp4(v[68:71], v[104:107], v[12:15], 2, 0, v[206], v[200])) + k.emit(buffer_load_dword(v[208], v[236], s[24:27], 0, 0, 1)) + k.emit(v_mfma_fp4(v[80:83], v[108:111], v[8:11], 1, 0, v[206], v[200])) + k.emit(s_cmp_lt_u32(s[62], s[61])) + k.emit(v_mfma_fp4(v[84:87], v[108:111], v[12:15], 3, 0, v[206], v[200])) + k.emit(v_mfma_fp4(v[72:75], v[104:107], v[16:19], 0, 0, v[206], v[201])) + k.emit(s_cselect_b32(s[69], s[69], 0)) + k.emit(v_mfma_fp4(v[76:79], v[104:107], v[20:23], 2, 0, v[206], v[201])) + k.emit(buffer_load_dword(v[209], v[237], s[24:27], 0, 0, 1)) + k.emit(v_mfma_fp4(v[88:91], v[108:111], v[16:19], 1, 0, v[206], v[201])) + k.emit(s_cselect_b32(s[71], s[71], 0)) + k.emit(v_mfma_fp4(v[92:95], v[108:111], v[20:23], 3, 0, v[206], v[201])) + k.emit(v_mfma_fp4(v[96:99], v[112:115], v[8:11], 0, 0, v[207], v[200])) + k.emit(s_add_u32(s[12], s[12], s[69])) + k.emit(v_mfma_fp4(v[100:103], v[112:115], v[12:15], 2, 0, v[207], v[200])) + k.emit(buffer_load_dwordx4(v[168:171], v[228], s[16:19], 0, 0, 1)) + k.emit(v_mfma_fp4(v[112:115], v[116:119], v[8:11], 1, 0, v[207], v[200])) + k.emit(s_addc_u32(s[13], 0, s[13])) + k.emit(v_mfma_fp4(v[116:119], v[116:119], v[12:15], 3, 0, v[207], v[200])) + k.emit(v_mfma_fp4(v[104:107], v[112:115], v[16:19], 0, 0, v[207], v[201])) + k.emit(s_sub_u32(s[14], s[14], s[69])) + k.emit(v_mfma_fp4(v[108:111], v[112:115], v[20:23], 2, 0, v[207], v[201])) + k.emit(buffer_load_dwordx4(v[172:175], v[229], s[16:19], 0, 0, 1)) + k.emit(v_mfma_fp4(v[120:123], v[116:119], v[16:19], 1, 0, v[207], v[201])) + k.emit(s_add_u32(s[20], s[20], s[71])) + k.emit(v_mfma_fp4(v[124:127], v[116:119], v[20:23], 3, 0, v[207], v[201])) + k.emit(v_mfma_fp4(v[64:67], v[120:123], v[24:27], 0, 3, v[206], v[200])) + k.emit(s_addc_u32(s[21], 0, s[21])) + k.emit(v_mfma_fp4(v[68:71], v[120:123], v[28:31], 2, 3, v[206], v[200])) + k.emit(buffer_load_dwordx4(v[176:179], v[230], s[16:19], 0, 0, 1)) + k.emit(v_mfma_fp4(v[80:83], v[124:127], v[24:27], 1, 3, v[206], v[200])) + k.emit(s_sub_u32(s[22], s[22], s[71])) + k.emit(v_mfma_fp4(v[84:87], v[124:127], v[28:31], 3, 3, v[206], v[200])) + k.emit(v_mfma_fp4(v[72:75], v[120:123], v[32:35], 0, 3, v[206], v[201])) + for i in range(2): + k.emit(v_mfma_fp4(v[76 + i * 24:79 + i * 24], v[120 + i * 8:123 + i * 8], v[36 + i * -8:39 + i * -8], 2, 3, v[206 + i * 1], v[201 + i * -1])) + k.emit(buffer_load_dwordx4(v[180 + i * 4:183 + i * 4], v[231 + i * 1], s[16:19], 0, 0, 1)) + k.emit(v_mfma_fp4(v[88 + i * 24:91 + i * 24], v[124 + i * 8:127 + i * 8], v[32 + i * -8:35 + i * -8], 1, 3, v[206 + i * 1], v[201 + i * -1])) + k.emit(v_mfma_fp4(v[92 + i * 24:95 + i * 24], v[124 + i * 8:127 + i * 8], v[36 + i * -8:39 + i * -8], 3, 3, v[206 + i * 1], v[201 + i * -1])) + k.emit(v_mfma_fp4(v[96 + i * 8:99 + i * 8], v[128:131], v[24 + i * 8:27 + i * 8], 0, 3, v[207], v[200 + i * 1])) + k.emit(v_mfma_fp4(v[108:111], v[128:131], v[36:39], 2, 3, v[207], v[201])) + k.emit(buffer_load_dwordx4(v[188:191], v[233], s[16:19], 0, 0, 1)) + k.emit(v_mfma_fp4(v[120:123], v[132:135], v[32:35], 1, 3, v[207], v[201])) + k.emit(v_mfma_fp4(v[124:127], v[132:135], v[36:39], 3, 3, v[207], v[201])) + k.emit(s_waitcnt(16498)) + k.emit(v_mfma_fp4(v[128:131], v[72:75], v[40:43], 0, 0, v[204], v[202])) + k.emit(s_barrier()) + k.emit(v_mfma_fp4(v[132:135], v[72:75], v[44:47], 2, 0, v[204], v[202])) + k.emit(buffer_load_dwordx4(v[192:195], v[234], s[16:19], 0, 0, 1)) + k.emit(v_mfma_fp4(v[144:147], v[76:79], v[40:43], 1, 0, v[204], v[202])) + k.emit(ds_read_b128(v[8:11], v[217])) + k.emit(v_mfma_fp4(v[148:151], v[76:79], v[44:47], 3, 0, v[204], v[202])) + k.emit(v_mfma_fp4(v[136:139], v[72:75], v[48:51], 0, 0, v[204], v[203])) + k.emit(ds_read_b128(v[24:27], v[217], v[0], v[0], 0, 64)) + k.emit(v_mfma_fp4(v[140:143], v[72:75], v[52:55], 2, 0, v[204], v[203])) + k.emit(buffer_load_dwordx4(v[196:199], v[235], s[16:19], 0, 0, 1)) + for i in range(2): + k.emit(v_mfma_fp4(v[152 + i * 24:155 + i * 24], v[76 + i * 8:79 + i * 8], v[48 + i * -8:51 + i * -8], 1, 0, v[204 + i * 1], v[203 + i * -1])) + k.emit(ds_read_b128(v[12 + i * 4:15 + i * 4], v[217], v[0], v[0], 0, 0 + i * 128, 2 + i * 14)) + k.emit(v_mfma_fp4(v[156 + i * 24:159 + i * 24], v[76 + i * 8:79 + i * 8], v[52 + i * -8:55 + i * -8], 3, 0, v[204 + i * 1], v[203 + i * -1])) + k.emit(v_mfma_fp4(v[160 + i * 8:163 + i * 8], v[80:83], v[40 + i * 8:43 + i * 8], 0, 0, v[205], v[202 + i * 1])) + k.emit(ds_read_b128(v[28 + i * 4:31 + i * 4], v[217], v[0], v[0], 0, 64 + i * 128, 2 + i * 14)) + k.emit(v_mfma_fp4(v[164 + i * 8:167 + i * 8], v[80:83], v[44 + i * 8:47 + i * 8], 2, 0, v[205], v[202 + i * 1])) + k.emit(buffer_load_dword(v[210 + i * 1], v[238 + i * 1], s[24:27], 0, 0, 1)) + k.emit(v_mfma_fp4(v[184:187], v[84:87], v[48:51], 1, 0, v[205], v[203])) + k.emit(ds_read_b128(v[20:23], v[217], v[0], v[0], 0, 128, 18)) + k.emit(v_mfma_fp4(v[188:191], v[84:87], v[52:55], 3, 0, v[205], v[203])) + k.emit(s_add_u32(NULL, 0, s[67])) + k.emit(v_mfma_fp4(v[128:131], v[88:91], v[56:59], 0, 3, v[204], v[202])) + k.emit(ds_read_b128(v[36:39], v[217], v[0], v[0], 0, 192, 18)) + k.emit(v_mfma_fp4(v[132:135], v[88:91], v[60:63], 2, 3, v[204], v[202])) + k.emit(buffer_load_dwordx4(v[0:3], v[212], s[12:15], 0, 0, 1, 0, 0, 0, 0, 1)) + k.emit(v_mfma_fp4(v[144:147], v[92:95], v[56:59], 1, 3, v[204], v[202])) + k.emit(ds_read_b32(v[200], v[219], v[0], v[0], 0, 0, 4)) + k.emit(v_mfma_fp4(v[148:151], v[92:95], v[60:63], 3, 3, v[204], v[202])) + k.emit(s_add_u32(NULL, LIT, s[67], 4224)) + k.emit(v_mfma_fp4(v[136:139], v[88:91], v[64:67], 0, 3, v[204], v[203])) + k.emit(ds_read_b32(v[201], v[219], v[0], v[0], 0, 0, 5)) + k.emit(v_mfma_fp4(v[140:143], v[88:91], v[68:71], 2, 3, v[204], v[203])) + k.emit(buffer_load_dwordx4(v[0:3], v[213], s[12:15], 0, 0, 1, 0, 0, 0, 0, 1)) + k.emit(v_mfma_fp4(v[152:155], v[92:95], v[64:67], 1, 3, v[204], v[203])) + k.emit(v_mfma_fp4(v[156:159], v[92:95], v[68:71], 3, 3, v[204], v[203])) + k.emit(s_add_u32(NULL, 0, s[68])) + k.emit(v_mfma_fp4(v[160:163], v[96:99], v[56:59], 0, 3, v[205], v[202])) + k.emit(v_mfma_fp4(v[164:167], v[96:99], v[60:63], 2, 3, v[205], v[202])) + k.emit(buffer_load_dword(v[0], v[218], s[20:23], 0, 0, 1, 0, 0, 0, 0, 1)) + for i in range(2): + k.emit(v_mfma_fp4(v[176 + i * 8:179 + i * 8], v[100:103], v[56 + i * 8:59 + i * 8], 1, 3, v[205], v[202 + i * 1])) + k.emit(v_mfma_fp4(v[180 + i * 8:183 + i * 8], v[100:103], v[60 + i * 8:63 + i * 8], 3, 3, v[205], v[202 + i * 1])) + k.emit(v_mfma_fp4(v[168 + i * 24:171 + i * 24], v[96 + i * 8:99 + i * 8], v[64 + i * -24:67 + i * -24], 0, 3 + i * -3, v[205 + i * 1], v[203 + i * -1])) + k.emit(v_mfma_fp4(v[172 + i * 24:175 + i * 24], v[96 + i * 8:99 + i * 8], v[68 + i * -24:71 + i * -24], 2, 3 + i * -3, v[205 + i * 1], v[203 + i * -1])) + k.emit(s_add_u32(NULL, LIT, s[67], 8448)) + k.emit(v_mfma_fp4(v[208:211], v[108:111], v[40:43], 1, 0, v[206], v[202])) + k.emit(buffer_load_dwordx4(v[0:3], v[214], s[12:15], 0, 0, 1, 0, 0, 0, 0, 1)) + k.emit(v_mfma_fp4(v[212:215], v[108:111], v[44:47], 3, 0, v[206], v[202])) + k.emit(v_mfma_fp4(v[200:203], v[104:107], v[48:51], 0, 0, v[206], v[203])) + k.emit(v_mfma_fp4(v[204:207], v[104:107], v[52:55], 2, 0, v[206], v[203])) + k.emit(s_add_u32(NULL, LIT, s[67], 12672)) + k.emit(v_mfma_fp4(v[216:219], v[108:111], v[48:51], 1, 0, v[206], v[203])) + k.emit(buffer_load_dwordx4(v[0:3], v[215], s[12:15], 0, 0, 1, 0, 0, 0, 0, 1)) + k.emit(v_mfma_fp4(v[220:223], v[108:111], v[52:55], 3, 0, v[206], v[203])) + k.emit(v_mfma_fp4(v[224:227], v[112:115], v[40:43], 0, 0, v[207], v[202])) + k.emit(s_add_u32(s[63], LIT, s[60], 512)) + k.emit(v_mfma_fp4(v[228:231], v[112:115], v[44:47], 2, 0, v[207], v[202])) + k.emit(v_mfma_fp4(v[240:243], v[116:119], v[40:43], 1, 0, v[207], v[202])) + k.emit(s_cmp_lt_u32(s[63], s[61])) + for i in range(2): + k.emit(v_mfma_fp4(v[244 + i * -8:247 + i * -8], v[116 + i * -4:119 + i * -4], v[44 + i * 8:47 + i * 8], 3 + i * -1, 0, v[207], v[202 + i * 1])) + k.emit(v_mfma_fp4(v[232 + i * 16:235 + i * 16], v[112 + i * 4:115 + i * 4], v[48:51], 0 + i * 1, 0, v[207], v[203])) + k.emit(s_cselect_b32(s[70 + i * 2], s[70 + i * 2], 0)) + for i in range(2): + k.emit(v_mfma_fp4(v[252 + i * -48:255 + i * -48], v[116 + i * 4:119 + i * 4], v[52 + i * 16:55 + i * 16], 3 + i * -1, 0 + i * 3, v[207 + i * -1], v[203])) + k.emit(v_mfma_fp4(v[192 + i * 24:195 + i * 24], v[120 + i * 4:123 + i * 4], v[56 + i * 8:59 + i * 8], 0 + i * 1, 3, v[206], v[202 + i * 1])) + k.emit(s_add_u32(s[16 + i * 8], s[16 + i * 8], s[70 + i * 2])) + k.emit(v_mfma_fp4(v[196 + i * 24:199 + i * 24], v[120 + i * 4:123 + i * 4], v[60 + i * 8:63 + i * 8], 2 + i * 1, 3, v[206], v[202 + i * 1])) + k.emit(v_mfma_fp4(v[208 + i * 16:211 + i * 16], v[124 + i * 4:127 + i * 4], v[56:59], 1 + i * -1, 3, v[206 + i * 1], v[202])) + k.emit(s_addc_u32(s[17 + i * 8], 0, s[17 + i * 8])) + k.emit(v_mfma_fp4(v[212 + i * 16:215 + i * 16], v[124 + i * 4:127 + i * 4], v[60:63], 3 + i * -1, 3, v[206 + i * 1], v[202])) + k.emit(v_mfma_fp4(v[200 + i * 40:203 + i * 40], v[120 + i * 12:123 + i * 12], v[64 + i * -8:67 + i * -8], 0 + i * 1, 3, v[206 + i * 1], v[203 + i * -1])) + k.emit(s_sub_u32(s[18 + i * 8], s[18 + i * 8], s[70 + i * 2])) + k.emit(v_mfma_fp4(v[244:247], v[132:135], v[60:63], 3, 3, v[207], v[202])) + k.emit(v_mfma_fp4(v[232:235], v[128:131], v[64:67], 0, 3, v[207], v[203])) + k.emit(s_addk_i32(s[60], 256)) + k.emit(v_mfma_fp4(v[236:239], v[128:131], v[68:71], 2, 3, v[207], v[203])) + k.emit(v_mfma_fp4(v[248:251], v[132:135], v[64:67], 1, 3, v[207], v[203])) + k.emit(s_cmp_lt_i32(s[60], s[61])) + k.emit(v_mfma_fp4(v[252:255], v[132:135], v[68:71], 3, 3, v[207], v[203])) + k.emit(s_cbranch_scc0(1926), target='L0_37CC') + k.emit(s_waitcnt(127)) + k.emit(v_mfma_fp4(v[0:3], v[136:139], v[8:11], 0, 0, v[208], v[200])) + k.emit(s_barrier()) + for i in range(2): + k.emit(v_mfma_fp4(v[4 + i * 32:7 + i * 32], v[136 + i * 8:139 + i * 8], v[12:15], 2, 0, v[208 + i * 1], v[200])) + k.emit(buffer_load_dwordx4(v[72 + i * 8:75 + i * 8], v[220 + i * 2], s[16:19], 0, 0, 1)) + for j6 in range(2): + k.emit(v_mfma_fp4(v[16 + j6 * -8 + i * 32:19 + j6 * -8 + i * 32], v[140 + j6 * -4 + i * 8:143 + j6 * -4 + i * 8], v[8 + j6 * 8:11 + j6 * 8], 1 + j6 * -1, 0, v[208 + i * 1], v[200 + j6 * 1])) + k.emit(ds_read_b128(v[40 + j6 * 16 + i * 8:43 + j6 * 16 + i * 8], v[217], v[0], v[0], 0, 0 + j6 * 64 + i * 128, 33 + i * 16)) + k.emit(v_mfma_fp4(v[20 + j6 * -8 + i * 32:23 + j6 * -8 + i * 32], v[140 + j6 * -4 + i * 8:143 + j6 * -4 + i * 8], v[12 + j6 * 8:15 + j6 * 8], 3 + j6 * -1, 0, v[208 + i * 1], v[200 + j6 * 1])) + k.emit(buffer_load_dwordx4(v[76 + i * 8:79 + i * 8], v[221 + i * 2], s[16:19], 0, 0, 1)) + k.emit(v_mfma_fp4(v[24 + i * 32:27 + i * 32], v[140 + i * 8:143 + i * 8], v[16:19], 1, 0, v[208 + i * 1], v[201])) + k.emit(ds_read_b128(v[44 + i * 8:47 + i * 8], v[217], v[0], v[0], 0, 0 + i * 128, 35 + i * 16)) + k.emit(v_mfma_fp4(v[28 + i * 32:31 + i * 32], v[140 + i * 8:143 + i * 8], v[20:23], 3, 0, v[208 + i * 1], v[201])) + k.emit(v_mfma_fp4(v[32 + i * -32:35 + i * -32], v[144 + i * 8:147 + i * 8], v[8 + i * 16:11 + i * 16], 0, 0 + i * 3, v[209 + i * -1], v[200])) + k.emit(ds_read_b128(v[60 + i * 8:63 + i * 8], v[217], v[0], v[0], 0, 64 + i * 128, 35 + i * 16)) + k.emit(v_mfma_fp4(v[4:7], v[152:155], v[28:31], 2, 3, v[208], v[200])) + k.emit(buffer_load_dwordx4(v[88:91], v[224], s[16:19], 0, 0, 1)) + for i in range(2): + k.emit(v_mfma_fp4(v[16 + i * -8:19 + i * -8], v[156 + i * -4:159 + i * -4], v[24 + i * 8:27 + i * 8], 1 + i * -1, 3, v[208], v[200 + i * 1])) + k.emit(ds_read_b32(v[202 + i * 1], v[219], v[0], v[0], 0, 0, 6 + i * 1)) + k.emit(v_mfma_fp4(v[20 + i * -8:23 + i * -8], v[156 + i * -4:159 + i * -4], v[28 + i * 8:31 + i * 8], 3 + i * -1, 3, v[208], v[200 + i * 1])) + for i in range(2): + k.emit(buffer_load_dwordx4(v[92 + i * 4:95 + i * 4], v[225 + i * 1], s[16:19], 0, 0, 1)) + k.emit(v_mfma_fp4(v[24 + i * 24:27 + i * 24], v[156 + i * 8:159 + i * 8], v[32 + i * -8:35 + i * -8], 1, 3, v[208 + i * 1], v[201 + i * -1])) + k.emit(v_mfma_fp4(v[28 + i * 24:31 + i * 24], v[156 + i * 8:159 + i * 8], v[36 + i * -8:39 + i * -8], 3, 3, v[208 + i * 1], v[201 + i * -1])) + k.emit(v_mfma_fp4(v[32 + i * 8:35 + i * 8], v[160:163], v[24 + i * 8:27 + i * 8], 0, 3, v[209], v[200 + i * 1])) + k.emit(v_mfma_fp4(v[36 + i * 8:39 + i * 8], v[160:163], v[28 + i * 8:31 + i * 8], 2, 3, v[209], v[200 + i * 1])) + k.emit(buffer_load_dwordx4(v[100:103], v[227], s[16:19], 0, 0, 1)) + k.emit(v_mfma_fp4(v[56:59], v[164:167], v[32:35], 1, 3, v[209], v[201])) + k.emit(v_mfma_fp4(v[60:63], v[164:167], v[36:39], 3, 3, v[209], v[201])) + k.emit(s_waitcnt(3965)) + k.emit(v_mfma_fp4(v[64:67], v[168:171], v[8:11], 0, 0, v[210], v[200])) + k.emit(s_add_u32(s[62], LIT, s[60], 512)) + k.emit(v_mfma_fp4(v[68:71], v[168:171], v[12:15], 2, 0, v[210], v[200])) + k.emit(buffer_load_dword(v[204], v[236], s[24:27], 0, 0, 1)) + k.emit(v_mfma_fp4(v[80:83], v[172:175], v[8:11], 1, 0, v[210], v[200])) + k.emit(s_cmp_lt_u32(s[62], s[61])) + k.emit(v_mfma_fp4(v[84:87], v[172:175], v[12:15], 3, 0, v[210], v[200])) + k.emit(v_mfma_fp4(v[72:75], v[168:171], v[16:19], 0, 0, v[210], v[201])) + k.emit(s_cselect_b32(s[69], s[69], 0)) + k.emit(v_mfma_fp4(v[76:79], v[168:171], v[20:23], 2, 0, v[210], v[201])) + k.emit(buffer_load_dword(v[205], v[237], s[24:27], 0, 0, 1)) + k.emit(v_mfma_fp4(v[88:91], v[172:175], v[16:19], 1, 0, v[210], v[201])) + k.emit(s_cselect_b32(s[71], s[71], 0)) + k.emit(v_mfma_fp4(v[92:95], v[172:175], v[20:23], 3, 0, v[210], v[201])) + k.emit(v_mfma_fp4(v[96:99], v[176:179], v[8:11], 0, 0, v[211], v[200])) + k.emit(s_add_u32(s[12], s[12], s[69])) + k.emit(v_mfma_fp4(v[100:103], v[176:179], v[12:15], 2, 0, v[211], v[200])) + k.emit(buffer_load_dwordx4(v[104:107], v[228], s[16:19], 0, 0, 1)) + k.emit(v_mfma_fp4(v[112:115], v[180:183], v[8:11], 1, 0, v[211], v[200])) + k.emit(s_addc_u32(s[13], 0, s[13])) + k.emit(v_mfma_fp4(v[116:119], v[180:183], v[12:15], 3, 0, v[211], v[200])) + k.emit(v_mfma_fp4(v[104:107], v[176:179], v[16:19], 0, 0, v[211], v[201])) + k.emit(s_sub_u32(s[14], s[14], s[69])) + k.emit(v_mfma_fp4(v[108:111], v[176:179], v[20:23], 2, 0, v[211], v[201])) + k.emit(buffer_load_dwordx4(v[108:111], v[229], s[16:19], 0, 0, 1)) + k.emit(v_mfma_fp4(v[120:123], v[180:183], v[16:19], 1, 0, v[211], v[201])) + k.emit(s_add_u32(s[20], s[20], s[71])) + k.emit(v_mfma_fp4(v[124:127], v[180:183], v[20:23], 3, 0, v[211], v[201])) + k.emit(v_mfma_fp4(v[64:67], v[184:187], v[24:27], 0, 3, v[210], v[200])) + k.emit(s_addc_u32(s[21], 0, s[21])) + k.emit(v_mfma_fp4(v[68:71], v[184:187], v[28:31], 2, 3, v[210], v[200])) + k.emit(buffer_load_dwordx4(v[112:115], v[230], s[16:19], 0, 0, 1)) + k.emit(v_mfma_fp4(v[80:83], v[188:191], v[24:27], 1, 3, v[210], v[200])) + k.emit(s_sub_u32(s[22], s[22], s[71])) + k.emit(v_mfma_fp4(v[84:87], v[188:191], v[28:31], 3, 3, v[210], v[200])) + k.emit(v_mfma_fp4(v[72:75], v[184:187], v[32:35], 0, 3, v[210], v[201])) + for i in range(2): + k.emit(v_mfma_fp4(v[76 + i * 24:79 + i * 24], v[184 + i * 8:187 + i * 8], v[36 + i * -8:39 + i * -8], 2, 3, v[210 + i * 1], v[201 + i * -1])) + k.emit(buffer_load_dwordx4(v[116 + i * 4:119 + i * 4], v[231 + i * 1], s[16:19], 0, 0, 1)) + k.emit(v_mfma_fp4(v[88 + i * 24:91 + i * 24], v[188 + i * 8:191 + i * 8], v[32 + i * -8:35 + i * -8], 1, 3, v[210 + i * 1], v[201 + i * -1])) + k.emit(v_mfma_fp4(v[92 + i * 24:95 + i * 24], v[188 + i * 8:191 + i * 8], v[36 + i * -8:39 + i * -8], 3, 3, v[210 + i * 1], v[201 + i * -1])) + k.emit(v_mfma_fp4(v[96 + i * 8:99 + i * 8], v[192:195], v[24 + i * 8:27 + i * 8], 0, 3, v[211], v[200 + i * 1])) + k.emit(v_mfma_fp4(v[108:111], v[192:195], v[36:39], 2, 3, v[211], v[201])) + k.emit(buffer_load_dwordx4(v[124:127], v[233], s[16:19], 0, 0, 1)) + k.emit(v_mfma_fp4(v[120:123], v[196:199], v[32:35], 1, 3, v[211], v[201])) + k.emit(v_mfma_fp4(v[124:127], v[196:199], v[36:39], 3, 3, v[211], v[201])) + k.emit(s_waitcnt(16498)) + k.emit(v_mfma_fp4(v[128:131], v[136:139], v[40:43], 0, 0, v[208], v[202])) + k.emit(s_barrier()) + k.emit(v_mfma_fp4(v[132:135], v[136:139], v[44:47], 2, 0, v[208], v[202])) + k.emit(buffer_load_dwordx4(v[128:131], v[234], s[16:19], 0, 0, 1)) + k.emit(v_mfma_fp4(v[144:147], v[140:143], v[40:43], 1, 0, v[208], v[202])) + k.emit(ds_read_b128(v[8:11], v[216])) + k.emit(v_mfma_fp4(v[148:151], v[140:143], v[44:47], 3, 0, v[208], v[202])) + k.emit(v_mfma_fp4(v[136:139], v[136:139], v[48:51], 0, 0, v[208], v[203])) + k.emit(ds_read_b128(v[24:27], v[216], v[0], v[0], 0, 64)) + k.emit(v_mfma_fp4(v[140:143], v[136:139], v[52:55], 2, 0, v[208], v[203])) + k.emit(buffer_load_dwordx4(v[132:135], v[235], s[16:19], 0, 0, 1)) + for i in range(2): + k.emit(v_mfma_fp4(v[152 + i * 24:155 + i * 24], v[140 + i * 8:143 + i * 8], v[48 + i * -8:51 + i * -8], 1, 0, v[208 + i * 1], v[203 + i * -1])) + k.emit(ds_read_b128(v[12 + i * 4:15 + i * 4], v[216], v[0], v[0], 0, 0 + i * 128, 2 + i * 14)) + k.emit(v_mfma_fp4(v[156 + i * 24:159 + i * 24], v[140 + i * 8:143 + i * 8], v[52 + i * -8:55 + i * -8], 3, 0, v[208 + i * 1], v[203 + i * -1])) + k.emit(v_mfma_fp4(v[160 + i * 8:163 + i * 8], v[144:147], v[40 + i * 8:43 + i * 8], 0, 0, v[209], v[202 + i * 1])) + k.emit(ds_read_b128(v[28 + i * 4:31 + i * 4], v[216], v[0], v[0], 0, 64 + i * 128, 2 + i * 14)) + k.emit(v_mfma_fp4(v[164 + i * 8:167 + i * 8], v[144:147], v[44 + i * 8:47 + i * 8], 2, 0, v[209], v[202 + i * 1])) + k.emit(buffer_load_dword(v[206 + i * 1], v[238 + i * 1], s[24:27], 0, 0, 1)) + k.emit(v_mfma_fp4(v[184:187], v[148:151], v[48:51], 1, 0, v[209], v[203])) + k.emit(ds_read_b128(v[20:23], v[216], v[0], v[0], 0, 128, 18)) + k.emit(v_mfma_fp4(v[188:191], v[148:151], v[52:55], 3, 0, v[209], v[203])) + k.emit(s_add_u32(NULL, LIT, s[67], 16896)) + k.emit(v_mfma_fp4(v[128:131], v[152:155], v[56:59], 0, 3, v[208], v[202])) + k.emit(ds_read_b128(v[36:39], v[216], v[0], v[0], 0, 192, 18)) + k.emit(v_mfma_fp4(v[132:135], v[152:155], v[60:63], 2, 3, v[208], v[202])) + k.emit(buffer_load_dwordx4(v[0:3], v[212], s[12:15], 0, 0, 1, 0, 0, 0, 0, 1)) + k.emit(v_mfma_fp4(v[144:147], v[156:159], v[56:59], 1, 3, v[208], v[202])) + k.emit(ds_read_b32(v[200], v[219])) + k.emit(v_mfma_fp4(v[148:151], v[156:159], v[60:63], 3, 3, v[208], v[202])) + k.emit(s_add_u32(NULL, LIT, s[67], 21120)) + k.emit(v_mfma_fp4(v[136:139], v[152:155], v[64:67], 0, 3, v[208], v[203])) + k.emit(ds_read_b32(v[201], v[219], v[0], v[0], 0, 0, 1)) + k.emit(v_mfma_fp4(v[140:143], v[152:155], v[68:71], 2, 3, v[208], v[203])) + k.emit(buffer_load_dwordx4(v[0:3], v[213], s[12:15], 0, 0, 1, 0, 0, 0, 0, 1)) + k.emit(v_mfma_fp4(v[152:155], v[156:159], v[64:67], 1, 3, v[208], v[203])) + k.emit(v_mfma_fp4(v[156:159], v[156:159], v[68:71], 3, 3, v[208], v[203])) + k.emit(s_add_u32(NULL, LIT, s[68], 1024)) + k.emit(v_mfma_fp4(v[160:163], v[160:163], v[56:59], 0, 3, v[209], v[202])) + k.emit(v_mfma_fp4(v[164:167], v[160:163], v[60:63], 2, 3, v[209], v[202])) + k.emit(buffer_load_dword(v[0], v[218], s[20:23], 0, 0, 1, 0, 0, 0, 0, 1)) + for i in range(2): + k.emit(v_mfma_fp4(v[176 + i * 8:179 + i * 8], v[164:167], v[56 + i * 8:59 + i * 8], 1, 3, v[209], v[202 + i * 1])) + k.emit(v_mfma_fp4(v[180 + i * 8:183 + i * 8], v[164:167], v[60 + i * 8:63 + i * 8], 3, 3, v[209], v[202 + i * 1])) + k.emit(v_mfma_fp4(v[168 + i * 24:171 + i * 24], v[160 + i * 8:163 + i * 8], v[64 + i * -24:67 + i * -24], 0, 3 + i * -3, v[209 + i * 1], v[203 + i * -1])) + k.emit(v_mfma_fp4(v[172 + i * 24:175 + i * 24], v[160 + i * 8:163 + i * 8], v[68 + i * -24:71 + i * -24], 2, 3 + i * -3, v[209 + i * 1], v[203 + i * -1])) + k.emit(s_add_u32(NULL, LIT, s[67], 25344)) + k.emit(v_mfma_fp4(v[208:211], v[172:175], v[40:43], 1, 0, v[210], v[202])) + k.emit(buffer_load_dwordx4(v[0:3], v[214], s[12:15], 0, 0, 1, 0, 0, 0, 0, 1)) + k.emit(v_mfma_fp4(v[212:215], v[172:175], v[44:47], 3, 0, v[210], v[202])) + k.emit(v_mfma_fp4(v[200:203], v[168:171], v[48:51], 0, 0, v[210], v[203])) + k.emit(v_mfma_fp4(v[204:207], v[168:171], v[52:55], 2, 0, v[210], v[203])) + k.emit(s_add_u32(NULL, LIT, s[67], 29568)) + k.emit(v_mfma_fp4(v[216:219], v[172:175], v[48:51], 1, 0, v[210], v[203])) + k.emit(buffer_load_dwordx4(v[0:3], v[215], s[12:15], 0, 0, 1, 0, 0, 0, 0, 1)) + k.emit(v_mfma_fp4(v[220:223], v[172:175], v[52:55], 3, 0, v[210], v[203])) + k.emit(v_mfma_fp4(v[224:227], v[176:179], v[40:43], 0, 0, v[211], v[202])) + k.emit(s_add_u32(s[63], LIT, s[60], 512)) + k.emit(v_mfma_fp4(v[228:231], v[176:179], v[44:47], 2, 0, v[211], v[202])) + k.emit(v_mfma_fp4(v[240:243], v[180:183], v[40:43], 1, 0, v[211], v[202])) + k.emit(s_cmp_lt_u32(s[63], s[61])) + for i in range(2): + k.emit(v_mfma_fp4(v[244 + i * -8:247 + i * -8], v[180 + i * -4:183 + i * -4], v[44 + i * 8:47 + i * 8], 3 + i * -1, 0, v[211], v[202 + i * 1])) + k.emit(v_mfma_fp4(v[232 + i * 16:235 + i * 16], v[176 + i * 4:179 + i * 4], v[48:51], 0 + i * 1, 0, v[211], v[203])) + k.emit(s_cselect_b32(s[70 + i * 2], s[70 + i * 2], 0)) + for i in range(2): + k.emit(v_mfma_fp4(v[252 + i * -48:255 + i * -48], v[180 + i * 4:183 + i * 4], v[52 + i * 16:55 + i * 16], 3 + i * -1, 0 + i * 3, v[211 + i * -1], v[203])) + k.emit(v_mfma_fp4(v[192 + i * 24:195 + i * 24], v[184 + i * 4:187 + i * 4], v[56 + i * 8:59 + i * 8], 0 + i * 1, 3, v[210], v[202 + i * 1])) + k.emit(s_add_u32(s[16 + i * 8], s[16 + i * 8], s[70 + i * 2])) + k.emit(v_mfma_fp4(v[196 + i * 24:199 + i * 24], v[184 + i * 4:187 + i * 4], v[60 + i * 8:63 + i * 8], 2 + i * 1, 3, v[210], v[202 + i * 1])) + k.emit(v_mfma_fp4(v[208 + i * 16:211 + i * 16], v[188 + i * 4:191 + i * 4], v[56:59], 1 + i * -1, 3, v[210 + i * 1], v[202])) + k.emit(s_addc_u32(s[17 + i * 8], 0, s[17 + i * 8])) + k.emit(v_mfma_fp4(v[212 + i * 16:215 + i * 16], v[188 + i * 4:191 + i * 4], v[60:63], 3 + i * -1, 3, v[210 + i * 1], v[202])) + k.emit(v_mfma_fp4(v[200 + i * 40:203 + i * 40], v[184 + i * 12:187 + i * 12], v[64 + i * -8:67 + i * -8], 0 + i * 1, 3, v[210 + i * 1], v[203 + i * -1])) + k.emit(s_sub_u32(s[18 + i * 8], s[18 + i * 8], s[70 + i * 2])) + k.emit(v_mfma_fp4(v[244:247], v[196:199], v[60:63], 3, 3, v[211], v[202])) + k.emit(v_mfma_fp4(v[232:235], v[192:195], v[64:67], 0, 3, v[211], v[203])) + k.emit(s_addk_i32(s[60], 256)) + k.emit(v_mfma_fp4(v[236:239], v[192:195], v[68:71], 2, 3, v[211], v[203])) + k.emit(v_mfma_fp4(v[248:251], v[196:199], v[64:67], 1, 3, v[211], v[203])) + k.emit(s_cmp_lt_i32(s[60], s[61])) + k.emit(v_mfma_fp4(v[252:255], v[196:199], v[68:71], 3, 3, v[211], v[203])) + k.emit(s_cbranch_scc0(1284), target='L0_37CC') + k.emit(s_branch(64253), target='L0_0FB4') + k.label('L0_23C0') + k.emit(s_waitcnt(127)) + k.emit(v_mfma_fp4(v[0:3], v[72:75], v[8:11], 0, 0, v[204], v[200])) + k.emit(s_barrier()) + for i in range(2): + k.emit(v_mfma_fp4(v[4 + i * 32:7 + i * 32], v[72 + i * 8:75 + i * 8], v[12:15], 2, 0, v[204 + i * 1], v[200])) + k.emit(ds_read_b128(v[40 + i * 8:43 + i * 8], v[216], v[0], v[0], 0, 0 + i * 128, 33 + i * 16)) + k.emit(v_mfma_fp4(v[16 + i * 32:19 + i * 32], v[76 + i * 8:79 + i * 8], v[8:11], 1, 0, v[204 + i * 1], v[200])) + k.emit(buffer_load_dwordx4(v[136 + i * 8:139 + i * 8], v[220 + i * 2], s[16:19], 0, 0, 1)) + for j7 in range(2): + k.emit(v_mfma_fp4(v[20 + j7 * -8 + i * 32:23 + j7 * -8 + i * 32], v[76 + j7 * -4 + i * 8:79 + j7 * -4 + i * 8], v[12 + j7 * 8:15 + j7 * 8], 3 + j7 * -1, 0, v[204 + i * 1], v[200 + j7 * 1])) + k.emit(ds_read_b128(v[56 + j7 * -12 + i * 8:59 + j7 * -12 + i * 8], v[216], v[0], v[0], 0, 64 + j7 * -64 + i * 128, 33 + j7 * 2 + i * 16)) + k.emit(v_mfma_fp4(v[8 + j7 * 16 + i * 32:11 + j7 * 16 + i * 32], v[72 + j7 * 4 + i * 8:75 + j7 * 4 + i * 8], v[16:19], 0 + j7 * 1, 0, v[204 + i * 1], v[201])) + k.emit(buffer_load_dwordx4(v[140 + i * 8:143 + i * 8], v[221 + i * 2], s[16:19], 0, 0, 1)) + k.emit(v_mfma_fp4(v[28 + i * 32:31 + i * 32], v[76 + i * 8:79 + i * 8], v[20:23], 3, 0, v[204 + i * 1], v[201])) + k.emit(ds_read_b128(v[60 + i * 8:63 + i * 8], v[216], v[0], v[0], 0, 64 + i * 128, 35 + i * 16)) + k.emit(v_mfma_fp4(v[32 + i * -32:35 + i * -32], v[80 + i * 8:83 + i * 8], v[8 + i * 16:11 + i * 16], 0, 0 + i * 3, v[205 + i * -1], v[200])) + k.emit(v_mfma_fp4(v[4:7], v[88:91], v[28:31], 2, 3, v[204], v[200])) + k.emit(ds_read_b32(v[202], v[219], v[0], v[0], 0, 0, 2)) + k.emit(v_mfma_fp4(v[16:19], v[92:95], v[24:27], 1, 3, v[204], v[200])) + k.emit(buffer_load_dwordx4(v[152:155], v[224], s[16:19], 0, 0, 1)) + k.emit(v_mfma_fp4(v[20:23], v[92:95], v[28:31], 3, 3, v[204], v[200])) + k.emit(ds_read_b32(v[203], v[219], v[0], v[0], 0, 0, 3)) + k.emit(v_mfma_fp4(v[8:11], v[88:91], v[32:35], 0, 3, v[204], v[201])) + for i in range(2): + k.emit(v_mfma_fp4(v[12 + i * 24:15 + i * 24], v[88 + i * 8:91 + i * 8], v[36 + i * -8:39 + i * -8], 2, 3, v[204 + i * 1], v[201 + i * -1])) + k.emit(v_mfma_fp4(v[24 + i * 24:27 + i * 24], v[92 + i * 8:95 + i * 8], v[32 + i * -8:35 + i * -8], 1, 3, v[204 + i * 1], v[201 + i * -1])) + k.emit(buffer_load_dwordx4(v[156 + i * 4:159 + i * 4], v[225 + i * 1], s[16:19], 0, 0, 1)) + k.emit(v_mfma_fp4(v[28 + i * 24:31 + i * 24], v[92 + i * 8:95 + i * 8], v[36 + i * -8:39 + i * -8], 3, 3, v[204 + i * 1], v[201 + i * -1])) + k.emit(v_mfma_fp4(v[32 + i * 8:35 + i * 8], v[96:99], v[24 + i * 8:27 + i * 8], 0, 3, v[205], v[200 + i * 1])) + k.emit(v_mfma_fp4(v[44:47], v[96:99], v[36:39], 2, 3, v[205], v[201])) + k.emit(v_mfma_fp4(v[56:59], v[100:103], v[32:35], 1, 3, v[205], v[201])) + k.emit(buffer_load_dwordx4(v[164:167], v[227], s[16:19], 0, 0, 1)) + k.emit(v_mfma_fp4(v[60:63], v[100:103], v[36:39], 3, 3, v[205], v[201])) + k.emit(s_waitcnt(3965)) + k.emit(v_mfma_fp4(v[64:67], v[104:107], v[8:11], 0, 0, v[206], v[200])) + k.emit(s_add_u32(s[62], LIT, s[60], 512)) + k.emit(v_mfma_fp4(v[68:71], v[104:107], v[12:15], 2, 0, v[206], v[200])) + k.emit(v_mfma_fp4(v[80:83], v[108:111], v[8:11], 1, 0, v[206], v[200])) + k.emit(s_cmp_lt_u32(s[62], s[61])) + k.emit(v_mfma_fp4(v[84:87], v[108:111], v[12:15], 3, 0, v[206], v[200])) + k.emit(buffer_load_dword(v[208], v[236], s[24:27], 0, 0, 1)) + for i in range(2): + k.emit(v_mfma_fp4(v[72 + i * 16:75 + i * 16], v[104 + i * 4:107 + i * 4], v[16:19], 0 + i * 1, 0, v[206], v[201])) + k.emit(s_cselect_b32(s[69 + i * 2], s[69 + i * 2], 0)) + k.emit(v_mfma_fp4(v[76 + i * 16:79 + i * 16], v[104 + i * 4:107 + i * 4], v[20:23], 2 + i * 1, 0, v[206], v[201])) + k.emit(buffer_load_dword(v[209], v[237], s[24:27], 0, 0, 1)) + k.emit(v_mfma_fp4(v[96:99], v[112:115], v[8:11], 0, 0, v[207], v[200])) + k.emit(s_add_u32(s[12], s[12], s[69])) + k.emit(v_mfma_fp4(v[100:103], v[112:115], v[12:15], 2, 0, v[207], v[200])) + k.emit(v_mfma_fp4(v[112:115], v[116:119], v[8:11], 1, 0, v[207], v[200])) + k.emit(s_addc_u32(s[13], 0, s[13])) + k.emit(v_mfma_fp4(v[116:119], v[116:119], v[12:15], 3, 0, v[207], v[200])) + k.emit(buffer_load_dwordx4(v[168:171], v[228], s[16:19], 0, 0, 1)) + k.emit(v_mfma_fp4(v[104:107], v[112:115], v[16:19], 0, 0, v[207], v[201])) + k.emit(s_sub_u32(s[14], s[14], s[69])) + k.emit(v_mfma_fp4(v[108:111], v[112:115], v[20:23], 2, 0, v[207], v[201])) + k.emit(v_mfma_fp4(v[120:123], v[116:119], v[16:19], 1, 0, v[207], v[201])) + k.emit(s_add_u32(s[20], s[20], s[71])) + k.emit(v_mfma_fp4(v[124:127], v[116:119], v[20:23], 3, 0, v[207], v[201])) + k.emit(buffer_load_dwordx4(v[172:175], v[229], s[16:19], 0, 0, 1)) + k.emit(v_mfma_fp4(v[64:67], v[120:123], v[24:27], 0, 3, v[206], v[200])) + k.emit(s_addc_u32(s[21], 0, s[21])) + k.emit(v_mfma_fp4(v[68:71], v[120:123], v[28:31], 2, 3, v[206], v[200])) + k.emit(v_mfma_fp4(v[80:83], v[124:127], v[24:27], 1, 3, v[206], v[200])) + k.emit(s_sub_u32(s[22], s[22], s[71])) + k.emit(v_mfma_fp4(v[84:87], v[124:127], v[28:31], 3, 3, v[206], v[200])) + k.emit(buffer_load_dwordx4(v[176:179], v[230], s[16:19], 0, 0, 1)) + k.emit(v_mfma_fp4(v[72:75], v[120:123], v[32:35], 0, 3, v[206], v[201])) + for i in range(2): + k.emit(v_mfma_fp4(v[76 + i * 24:79 + i * 24], v[120 + i * 8:123 + i * 8], v[36 + i * -8:39 + i * -8], 2, 3, v[206 + i * 1], v[201 + i * -1])) + k.emit(v_mfma_fp4(v[88 + i * 24:91 + i * 24], v[124 + i * 8:127 + i * 8], v[32 + i * -8:35 + i * -8], 1, 3, v[206 + i * 1], v[201 + i * -1])) + k.emit(v_mfma_fp4(v[92 + i * 24:95 + i * 24], v[124 + i * 8:127 + i * 8], v[36 + i * -8:39 + i * -8], 3, 3, v[206 + i * 1], v[201 + i * -1])) + k.emit(buffer_load_dwordx4(v[180 + i * 4:183 + i * 4], v[231 + i * 1], s[16:19], 0, 0, 1)) + k.emit(v_mfma_fp4(v[96 + i * 8:99 + i * 8], v[128:131], v[24 + i * 8:27 + i * 8], 0, 3, v[207], v[200 + i * 1])) + k.emit(v_mfma_fp4(v[108:111], v[128:131], v[36:39], 2, 3, v[207], v[201])) + k.emit(v_mfma_fp4(v[120:123], v[132:135], v[32:35], 1, 3, v[207], v[201])) + k.emit(v_mfma_fp4(v[124:127], v[132:135], v[36:39], 3, 3, v[207], v[201])) + k.emit(buffer_load_dwordx4(v[188:191], v[233], s[16:19], 0, 0, 1)) + k.emit(s_waitcnt(16498)) + k.emit(v_mfma_fp4(v[128:131], v[72:75], v[40:43], 0, 0, v[204], v[202])) + k.emit(s_barrier()) + k.emit(v_mfma_fp4(v[132:135], v[72:75], v[44:47], 2, 0, v[204], v[202])) + k.emit(ds_read_b128(v[8:11], v[217])) + k.emit(v_mfma_fp4(v[144:147], v[76:79], v[40:43], 1, 0, v[204], v[202])) + k.emit(buffer_load_dwordx4(v[192:195], v[234], s[16:19], 0, 0, 1)) + k.emit(v_mfma_fp4(v[148:151], v[76:79], v[44:47], 3, 0, v[204], v[202])) + k.emit(ds_read_b128(v[24:27], v[217], v[0], v[0], 0, 64)) + k.emit(v_mfma_fp4(v[136:139], v[72:75], v[48:51], 0, 0, v[204], v[203])) + k.emit(v_mfma_fp4(v[140:143], v[72:75], v[52:55], 2, 0, v[204], v[203])) + k.emit(ds_read_b128(v[12:15], v[217], v[0], v[0], 0, 0, 2)) + k.emit(v_mfma_fp4(v[152:155], v[76:79], v[48:51], 1, 0, v[204], v[203])) + k.emit(buffer_load_dwordx4(v[196:199], v[235], s[16:19], 0, 0, 1)) + for i in range(2): + k.emit(v_mfma_fp4(v[156 + i * 24:159 + i * 24], v[76 + i * 8:79 + i * 8], v[52 + i * -8:55 + i * -8], 3, 0, v[204 + i * 1], v[203 + i * -1])) + k.emit(ds_read_b128(v[28 + i * 4:31 + i * 4], v[217], v[0], v[0], 0, 64 + i * 128, 2 + i * 14)) + k.emit(v_mfma_fp4(v[160 + i * 8:163 + i * 8], v[80:83], v[40 + i * 8:43 + i * 8], 0, 0, v[205], v[202 + i * 1])) + k.emit(v_mfma_fp4(v[164 + i * 8:167 + i * 8], v[80:83], v[44 + i * 8:47 + i * 8], 2, 0, v[205], v[202 + i * 1])) + k.emit(ds_read_b128(v[16 + i * 4:19 + i * 4], v[217], v[0], v[0], 0, 128, 16 + i * 2)) + k.emit(v_mfma_fp4(v[176 + i * 8:179 + i * 8], v[84:87], v[40 + i * 8:43 + i * 8], 1, 0, v[205], v[202 + i * 1])) + k.emit(buffer_load_dword(v[210 + i * 1], v[238 + i * 1], s[24:27], 0, 0, 1)) + k.emit(v_mfma_fp4(v[188:191], v[84:87], v[52:55], 3, 0, v[205], v[203])) + k.emit(ds_read_b128(v[36:39], v[217], v[0], v[0], 0, 192, 18)) + k.emit(v_mfma_fp4(v[128:131], v[88:91], v[56:59], 0, 3, v[204], v[202])) + k.emit(s_add_u32(NULL, 0, s[67])) + k.emit(v_mfma_fp4(v[132:135], v[88:91], v[60:63], 2, 3, v[204], v[202])) + k.emit(ds_read_b32(v[200], v[219], v[0], v[0], 0, 0, 4)) + k.emit(v_mfma_fp4(v[144:147], v[92:95], v[56:59], 1, 3, v[204], v[202])) + k.emit(buffer_load_dwordx4(v[0:3], v[212], s[12:15], 0, 0, 1, 0, 0, 0, 0, 1)) + k.emit(v_mfma_fp4(v[148:151], v[92:95], v[60:63], 3, 3, v[204], v[202])) + k.emit(ds_read_b32(v[201], v[219], v[0], v[0], 0, 0, 5)) + k.emit(v_mfma_fp4(v[136:139], v[88:91], v[64:67], 0, 3, v[204], v[203])) + k.emit(s_add_u32(NULL, LIT, s[67], 4224)) + k.emit(v_mfma_fp4(v[140:143], v[88:91], v[68:71], 2, 3, v[204], v[203])) + k.emit(v_mfma_fp4(v[152:155], v[92:95], v[64:67], 1, 3, v[204], v[203])) + k.emit(buffer_load_dwordx4(v[0:3], v[213], s[12:15], 0, 0, 1, 0, 0, 0, 0, 1)) + k.emit(v_mfma_fp4(v[156:159], v[92:95], v[68:71], 3, 3, v[204], v[203])) + k.emit(v_mfma_fp4(v[160:163], v[96:99], v[56:59], 0, 3, v[205], v[202])) + k.emit(s_add_u32(NULL, 0, s[68])) + k.emit(v_mfma_fp4(v[164:167], v[96:99], v[60:63], 2, 3, v[205], v[202])) + k.emit(v_mfma_fp4(v[176:179], v[100:103], v[56:59], 1, 3, v[205], v[202])) + k.emit(buffer_load_dword(v[0], v[218], s[20:23], 0, 0, 1, 0, 0, 0, 0, 1)) + for i in range(2): + k.emit(v_mfma_fp4(v[180 + i * 4:183 + i * 4], v[100:103], v[60 + i * 4:63 + i * 4], 3 + i * -2, 3, v[205], v[202 + i * 1])) + k.emit(v_mfma_fp4(v[168 + i * 20:171 + i * 20], v[96 + i * 4:99 + i * 4], v[64 + i * 4:67 + i * 4], 0 + i * 3, 3, v[205], v[203])) + k.emit(v_mfma_fp4(v[172 + i * 20:175 + i * 20], v[96 + i * 8:99 + i * 8], v[68 + i * -28:71 + i * -28], 2 + i * -2, 3 + i * -3, v[205 + i * 1], v[203 + i * -1])) + for i in range(2): + k.emit(s_add_u32(NULL, LIT, s[67], 8448 + i * 4224)) + k.emit(v_mfma_fp4(v[196 + i * 8:199 + i * 8], v[104:107], v[44 + i * 8:47 + i * 8], 2, 0, v[206], v[202 + i * 1])) + k.emit(buffer_load_dwordx4(v[0:3], v[214 + i * 1], s[12:15], 0, 0, 1, 0, 0, 0, 0, 1)) + k.emit(v_mfma_fp4(v[208 + i * 8:211 + i * 8], v[108:111], v[40 + i * 8:43 + i * 8], 1, 0, v[206], v[202 + i * 1])) + k.emit(v_mfma_fp4(v[212 + i * 8:215 + i * 8], v[108:111], v[44 + i * 8:47 + i * 8], 3, 0, v[206], v[202 + i * 1])) + k.emit(v_mfma_fp4(v[200 + i * 24:203 + i * 24], v[104 + i * 8:107 + i * 8], v[48 + i * -8:51 + i * -8], 0, 0, v[206 + i * 1], v[203 + i * -1])) + k.emit(s_add_u32(s[63], LIT, s[60], 512)) + k.emit(v_mfma_fp4(v[228:231], v[112:115], v[44:47], 2, 0, v[207], v[202])) + k.emit(v_mfma_fp4(v[240:243], v[116:119], v[40:43], 1, 0, v[207], v[202])) + k.emit(s_cmp_lt_u32(s[63], s[61])) + for i in range(2): + k.emit(v_mfma_fp4(v[244 + i * -8:247 + i * -8], v[116 + i * -4:119 + i * -4], v[44 + i * 8:47 + i * 8], 3 + i * -1, 0, v[207], v[202 + i * 1])) + k.emit(v_mfma_fp4(v[232 + i * 16:235 + i * 16], v[112 + i * 4:115 + i * 4], v[48:51], 0 + i * 1, 0, v[207], v[203])) + k.emit(s_cselect_b32(s[70 + i * 2], s[70 + i * 2], 0)) + for i in range(2): + k.emit(v_mfma_fp4(v[252 + i * -48:255 + i * -48], v[116 + i * 4:119 + i * 4], v[52 + i * 16:55 + i * 16], 3 + i * -1, 0 + i * 3, v[207 + i * -1], v[203])) + k.emit(v_mfma_fp4(v[192 + i * 24:195 + i * 24], v[120 + i * 4:123 + i * 4], v[56 + i * 8:59 + i * 8], 0 + i * 1, 3, v[206], v[202 + i * 1])) + k.emit(s_add_u32(s[16 + i * 8], s[16 + i * 8], s[70 + i * 2])) + k.emit(v_mfma_fp4(v[196 + i * 24:199 + i * 24], v[120 + i * 4:123 + i * 4], v[60 + i * 8:63 + i * 8], 2 + i * 1, 3, v[206], v[202 + i * 1])) + k.emit(v_mfma_fp4(v[208 + i * 16:211 + i * 16], v[124 + i * 4:127 + i * 4], v[56:59], 1 + i * -1, 3, v[206 + i * 1], v[202])) + k.emit(s_addc_u32(s[17 + i * 8], 0, s[17 + i * 8])) + k.emit(v_mfma_fp4(v[212 + i * 16:215 + i * 16], v[124 + i * 4:127 + i * 4], v[60:63], 3 + i * -1, 3, v[206 + i * 1], v[202])) + k.emit(v_mfma_fp4(v[200 + i * 40:203 + i * 40], v[120 + i * 12:123 + i * 12], v[64 + i * -8:67 + i * -8], 0 + i * 1, 3, v[206 + i * 1], v[203 + i * -1])) + k.emit(s_sub_u32(s[18 + i * 8], s[18 + i * 8], s[70 + i * 2])) + k.emit(v_mfma_fp4(v[244:247], v[132:135], v[60:63], 3, 3, v[207], v[202])) + k.emit(v_mfma_fp4(v[232:235], v[128:131], v[64:67], 0, 3, v[207], v[203])) + k.emit(s_addk_i32(s[60], 256)) + k.emit(v_mfma_fp4(v[236:239], v[128:131], v[68:71], 2, 3, v[207], v[203])) + k.emit(v_mfma_fp4(v[248:251], v[132:135], v[64:67], 1, 3, v[207], v[203])) + k.emit(s_cmp_lt_i32(s[60], s[61])) + k.emit(v_mfma_fp4(v[252:255], v[132:135], v[68:71], 3, 3, v[207], v[203])) + k.emit(s_cbranch_scc0(643), target='L0_37CC') + k.emit(s_waitcnt(127)) + k.emit(v_mfma_fp4(v[0:3], v[136:139], v[8:11], 0, 0, v[208], v[200])) + k.emit(s_barrier()) + for i in range(2): + k.emit(v_mfma_fp4(v[4 + i * 32:7 + i * 32], v[136 + i * 8:139 + i * 8], v[12:15], 2, 0, v[208 + i * 1], v[200])) + k.emit(ds_read_b128(v[40 + i * 8:43 + i * 8], v[217], v[0], v[0], 0, 0 + i * 128, 33 + i * 16)) + k.emit(v_mfma_fp4(v[16 + i * 32:19 + i * 32], v[140 + i * 8:143 + i * 8], v[8:11], 1, 0, v[208 + i * 1], v[200])) + k.emit(buffer_load_dwordx4(v[72 + i * 8:75 + i * 8], v[220 + i * 2], s[16:19], 0, 0, 1)) + for j8 in range(2): + k.emit(v_mfma_fp4(v[20 + j8 * -8 + i * 32:23 + j8 * -8 + i * 32], v[140 + j8 * -4 + i * 8:143 + j8 * -4 + i * 8], v[12 + j8 * 8:15 + j8 * 8], 3 + j8 * -1, 0, v[208 + i * 1], v[200 + j8 * 1])) + k.emit(ds_read_b128(v[56 + j8 * -12 + i * 8:59 + j8 * -12 + i * 8], v[217], v[0], v[0], 0, 64 + j8 * -64 + i * 128, 33 + j8 * 2 + i * 16)) + k.emit(v_mfma_fp4(v[8 + j8 * 16 + i * 32:11 + j8 * 16 + i * 32], v[136 + j8 * 4 + i * 8:139 + j8 * 4 + i * 8], v[16:19], 0 + j8 * 1, 0, v[208 + i * 1], v[201])) + k.emit(buffer_load_dwordx4(v[76 + i * 8:79 + i * 8], v[221 + i * 2], s[16:19], 0, 0, 1)) + k.emit(v_mfma_fp4(v[28 + i * 32:31 + i * 32], v[140 + i * 8:143 + i * 8], v[20:23], 3, 0, v[208 + i * 1], v[201])) + k.emit(ds_read_b128(v[60 + i * 8:63 + i * 8], v[217], v[0], v[0], 0, 64 + i * 128, 35 + i * 16)) + k.emit(v_mfma_fp4(v[32 + i * -32:35 + i * -32], v[144 + i * 8:147 + i * 8], v[8 + i * 16:11 + i * 16], 0, 0 + i * 3, v[209 + i * -1], v[200])) + k.emit(v_mfma_fp4(v[4:7], v[152:155], v[28:31], 2, 3, v[208], v[200])) + k.emit(ds_read_b32(v[202], v[219], v[0], v[0], 0, 0, 6)) + k.emit(v_mfma_fp4(v[16:19], v[156:159], v[24:27], 1, 3, v[208], v[200])) + k.emit(buffer_load_dwordx4(v[88:91], v[224], s[16:19], 0, 0, 1)) + k.emit(v_mfma_fp4(v[20:23], v[156:159], v[28:31], 3, 3, v[208], v[200])) + k.emit(ds_read_b32(v[203], v[219], v[0], v[0], 0, 0, 7)) + k.emit(v_mfma_fp4(v[8:11], v[152:155], v[32:35], 0, 3, v[208], v[201])) + for i in range(2): + k.emit(v_mfma_fp4(v[12 + i * 24:15 + i * 24], v[152 + i * 8:155 + i * 8], v[36 + i * -8:39 + i * -8], 2, 3, v[208 + i * 1], v[201 + i * -1])) + k.emit(v_mfma_fp4(v[24 + i * 24:27 + i * 24], v[156 + i * 8:159 + i * 8], v[32 + i * -8:35 + i * -8], 1, 3, v[208 + i * 1], v[201 + i * -1])) + k.emit(buffer_load_dwordx4(v[92 + i * 4:95 + i * 4], v[225 + i * 1], s[16:19], 0, 0, 1)) + k.emit(v_mfma_fp4(v[28 + i * 24:31 + i * 24], v[156 + i * 8:159 + i * 8], v[36 + i * -8:39 + i * -8], 3, 3, v[208 + i * 1], v[201 + i * -1])) + k.emit(v_mfma_fp4(v[32 + i * 8:35 + i * 8], v[160:163], v[24 + i * 8:27 + i * 8], 0, 3, v[209], v[200 + i * 1])) + k.emit(v_mfma_fp4(v[44:47], v[160:163], v[36:39], 2, 3, v[209], v[201])) + k.emit(v_mfma_fp4(v[56:59], v[164:167], v[32:35], 1, 3, v[209], v[201])) + k.emit(buffer_load_dwordx4(v[100:103], v[227], s[16:19], 0, 0, 1)) + k.emit(v_mfma_fp4(v[60:63], v[164:167], v[36:39], 3, 3, v[209], v[201])) + k.emit(s_waitcnt(3965)) + k.emit(v_mfma_fp4(v[64:67], v[168:171], v[8:11], 0, 0, v[210], v[200])) + k.emit(s_add_u32(s[62], LIT, s[60], 512)) + k.emit(v_mfma_fp4(v[68:71], v[168:171], v[12:15], 2, 0, v[210], v[200])) + k.emit(v_mfma_fp4(v[80:83], v[172:175], v[8:11], 1, 0, v[210], v[200])) + k.emit(s_cmp_lt_u32(s[62], s[61])) + k.emit(v_mfma_fp4(v[84:87], v[172:175], v[12:15], 3, 0, v[210], v[200])) + k.emit(buffer_load_dword(v[204], v[236], s[24:27], 0, 0, 1)) + for i in range(2): + k.emit(v_mfma_fp4(v[72 + i * 16:75 + i * 16], v[168 + i * 4:171 + i * 4], v[16:19], 0 + i * 1, 0, v[210], v[201])) + k.emit(s_cselect_b32(s[69 + i * 2], s[69 + i * 2], 0)) + k.emit(v_mfma_fp4(v[76 + i * 16:79 + i * 16], v[168 + i * 4:171 + i * 4], v[20:23], 2 + i * 1, 0, v[210], v[201])) + k.emit(buffer_load_dword(v[205], v[237], s[24:27], 0, 0, 1)) + k.emit(v_mfma_fp4(v[96:99], v[176:179], v[8:11], 0, 0, v[211], v[200])) + k.emit(s_add_u32(s[12], s[12], s[69])) + k.emit(v_mfma_fp4(v[100:103], v[176:179], v[12:15], 2, 0, v[211], v[200])) + k.emit(v_mfma_fp4(v[112:115], v[180:183], v[8:11], 1, 0, v[211], v[200])) + k.emit(s_addc_u32(s[13], 0, s[13])) + k.emit(v_mfma_fp4(v[116:119], v[180:183], v[12:15], 3, 0, v[211], v[200])) + k.emit(buffer_load_dwordx4(v[104:107], v[228], s[16:19], 0, 0, 1)) + k.emit(v_mfma_fp4(v[104:107], v[176:179], v[16:19], 0, 0, v[211], v[201])) + k.emit(s_sub_u32(s[14], s[14], s[69])) + k.emit(v_mfma_fp4(v[108:111], v[176:179], v[20:23], 2, 0, v[211], v[201])) + k.emit(v_mfma_fp4(v[120:123], v[180:183], v[16:19], 1, 0, v[211], v[201])) + k.emit(s_add_u32(s[20], s[20], s[71])) + k.emit(v_mfma_fp4(v[124:127], v[180:183], v[20:23], 3, 0, v[211], v[201])) + k.emit(buffer_load_dwordx4(v[108:111], v[229], s[16:19], 0, 0, 1)) + k.emit(v_mfma_fp4(v[64:67], v[184:187], v[24:27], 0, 3, v[210], v[200])) + k.emit(s_addc_u32(s[21], 0, s[21])) + k.emit(v_mfma_fp4(v[68:71], v[184:187], v[28:31], 2, 3, v[210], v[200])) + k.emit(v_mfma_fp4(v[80:83], v[188:191], v[24:27], 1, 3, v[210], v[200])) + k.emit(s_sub_u32(s[22], s[22], s[71])) + k.emit(v_mfma_fp4(v[84:87], v[188:191], v[28:31], 3, 3, v[210], v[200])) + k.emit(buffer_load_dwordx4(v[112:115], v[230], s[16:19], 0, 0, 1)) + k.emit(v_mfma_fp4(v[72:75], v[184:187], v[32:35], 0, 3, v[210], v[201])) + for i in range(2): + k.emit(v_mfma_fp4(v[76 + i * 24:79 + i * 24], v[184 + i * 8:187 + i * 8], v[36 + i * -8:39 + i * -8], 2, 3, v[210 + i * 1], v[201 + i * -1])) + k.emit(v_mfma_fp4(v[88 + i * 24:91 + i * 24], v[188 + i * 8:191 + i * 8], v[32 + i * -8:35 + i * -8], 1, 3, v[210 + i * 1], v[201 + i * -1])) + k.emit(v_mfma_fp4(v[92 + i * 24:95 + i * 24], v[188 + i * 8:191 + i * 8], v[36 + i * -8:39 + i * -8], 3, 3, v[210 + i * 1], v[201 + i * -1])) + k.emit(buffer_load_dwordx4(v[116 + i * 4:119 + i * 4], v[231 + i * 1], s[16:19], 0, 0, 1)) + k.emit(v_mfma_fp4(v[96 + i * 8:99 + i * 8], v[192:195], v[24 + i * 8:27 + i * 8], 0, 3, v[211], v[200 + i * 1])) + k.emit(v_mfma_fp4(v[108:111], v[192:195], v[36:39], 2, 3, v[211], v[201])) + k.emit(v_mfma_fp4(v[120:123], v[196:199], v[32:35], 1, 3, v[211], v[201])) + k.emit(v_mfma_fp4(v[124:127], v[196:199], v[36:39], 3, 3, v[211], v[201])) + k.emit(buffer_load_dwordx4(v[124:127], v[233], s[16:19], 0, 0, 1)) + k.emit(s_waitcnt(16498)) + k.emit(v_mfma_fp4(v[128:131], v[136:139], v[40:43], 0, 0, v[208], v[202])) + k.emit(s_barrier()) + k.emit(v_mfma_fp4(v[132:135], v[136:139], v[44:47], 2, 0, v[208], v[202])) + k.emit(ds_read_b128(v[8:11], v[216])) + k.emit(v_mfma_fp4(v[144:147], v[140:143], v[40:43], 1, 0, v[208], v[202])) + k.emit(buffer_load_dwordx4(v[128:131], v[234], s[16:19], 0, 0, 1)) + k.emit(v_mfma_fp4(v[148:151], v[140:143], v[44:47], 3, 0, v[208], v[202])) + k.emit(ds_read_b128(v[24:27], v[216], v[0], v[0], 0, 64)) + k.emit(v_mfma_fp4(v[136:139], v[136:139], v[48:51], 0, 0, v[208], v[203])) + k.emit(v_mfma_fp4(v[140:143], v[136:139], v[52:55], 2, 0, v[208], v[203])) + k.emit(ds_read_b128(v[12:15], v[216], v[0], v[0], 0, 0, 2)) + k.emit(v_mfma_fp4(v[152:155], v[140:143], v[48:51], 1, 0, v[208], v[203])) + k.emit(buffer_load_dwordx4(v[132:135], v[235], s[16:19], 0, 0, 1)) + for i in range(2): + k.emit(v_mfma_fp4(v[156 + i * 24:159 + i * 24], v[140 + i * 8:143 + i * 8], v[52 + i * -8:55 + i * -8], 3, 0, v[208 + i * 1], v[203 + i * -1])) + k.emit(ds_read_b128(v[28 + i * 4:31 + i * 4], v[216], v[0], v[0], 0, 64 + i * 128, 2 + i * 14)) + k.emit(v_mfma_fp4(v[160 + i * 8:163 + i * 8], v[144:147], v[40 + i * 8:43 + i * 8], 0, 0, v[209], v[202 + i * 1])) + k.emit(v_mfma_fp4(v[164 + i * 8:167 + i * 8], v[144:147], v[44 + i * 8:47 + i * 8], 2, 0, v[209], v[202 + i * 1])) + k.emit(ds_read_b128(v[16 + i * 4:19 + i * 4], v[216], v[0], v[0], 0, 128, 16 + i * 2)) + k.emit(v_mfma_fp4(v[176 + i * 8:179 + i * 8], v[148:151], v[40 + i * 8:43 + i * 8], 1, 0, v[209], v[202 + i * 1])) + k.emit(buffer_load_dword(v[206 + i * 1], v[238 + i * 1], s[24:27], 0, 0, 1)) + k.emit(v_mfma_fp4(v[188:191], v[148:151], v[52:55], 3, 0, v[209], v[203])) + k.emit(ds_read_b128(v[36:39], v[216], v[0], v[0], 0, 192, 18)) + k.emit(v_mfma_fp4(v[128:131], v[152:155], v[56:59], 0, 3, v[208], v[202])) + k.emit(s_add_u32(NULL, LIT, s[67], 16896)) + k.emit(v_mfma_fp4(v[132:135], v[152:155], v[60:63], 2, 3, v[208], v[202])) + k.emit(ds_read_b32(v[200], v[219])) + k.emit(v_mfma_fp4(v[144:147], v[156:159], v[56:59], 1, 3, v[208], v[202])) + k.emit(buffer_load_dwordx4(v[0:3], v[212], s[12:15], 0, 0, 1, 0, 0, 0, 0, 1)) + k.emit(v_mfma_fp4(v[148:151], v[156:159], v[60:63], 3, 3, v[208], v[202])) + k.emit(ds_read_b32(v[201], v[219], v[0], v[0], 0, 0, 1)) + k.emit(v_mfma_fp4(v[136:139], v[152:155], v[64:67], 0, 3, v[208], v[203])) + k.emit(s_add_u32(NULL, LIT, s[67], 21120)) + k.emit(v_mfma_fp4(v[140:143], v[152:155], v[68:71], 2, 3, v[208], v[203])) + k.emit(v_mfma_fp4(v[152:155], v[156:159], v[64:67], 1, 3, v[208], v[203])) + k.emit(buffer_load_dwordx4(v[0:3], v[213], s[12:15], 0, 0, 1, 0, 0, 0, 0, 1)) + k.emit(v_mfma_fp4(v[156:159], v[156:159], v[68:71], 3, 3, v[208], v[203])) + k.emit(v_mfma_fp4(v[160:163], v[160:163], v[56:59], 0, 3, v[209], v[202])) + k.emit(s_add_u32(NULL, LIT, s[68], 1024)) + k.emit(v_mfma_fp4(v[164:167], v[160:163], v[60:63], 2, 3, v[209], v[202])) + k.emit(v_mfma_fp4(v[176:179], v[164:167], v[56:59], 1, 3, v[209], v[202])) + k.emit(buffer_load_dword(v[0], v[218], s[20:23], 0, 0, 1, 0, 0, 0, 0, 1)) + for i in range(2): + k.emit(v_mfma_fp4(v[180 + i * 4:183 + i * 4], v[164:167], v[60 + i * 4:63 + i * 4], 3 + i * -2, 3, v[209], v[202 + i * 1])) + k.emit(v_mfma_fp4(v[168 + i * 20:171 + i * 20], v[160 + i * 4:163 + i * 4], v[64 + i * 4:67 + i * 4], 0 + i * 3, 3, v[209], v[203])) + k.emit(v_mfma_fp4(v[172 + i * 20:175 + i * 20], v[160 + i * 8:163 + i * 8], v[68 + i * -28:71 + i * -28], 2 + i * -2, 3 + i * -3, v[209 + i * 1], v[203 + i * -1])) + for i in range(2): + k.emit(s_add_u32(NULL, LIT, s[67], 25344 + i * 4224)) + k.emit(v_mfma_fp4(v[196 + i * 8:199 + i * 8], v[168:171], v[44 + i * 8:47 + i * 8], 2, 0, v[210], v[202 + i * 1])) + k.emit(buffer_load_dwordx4(v[0:3], v[214 + i * 1], s[12:15], 0, 0, 1, 0, 0, 0, 0, 1)) + k.emit(v_mfma_fp4(v[208 + i * 8:211 + i * 8], v[172:175], v[40 + i * 8:43 + i * 8], 1, 0, v[210], v[202 + i * 1])) + k.emit(v_mfma_fp4(v[212 + i * 8:215 + i * 8], v[172:175], v[44 + i * 8:47 + i * 8], 3, 0, v[210], v[202 + i * 1])) + k.emit(v_mfma_fp4(v[200 + i * 24:203 + i * 24], v[168 + i * 8:171 + i * 8], v[48 + i * -8:51 + i * -8], 0, 0, v[210 + i * 1], v[203 + i * -1])) + k.emit(s_add_u32(s[63], LIT, s[60], 512)) + k.emit(v_mfma_fp4(v[228:231], v[176:179], v[44:47], 2, 0, v[211], v[202])) + k.emit(v_mfma_fp4(v[240:243], v[180:183], v[40:43], 1, 0, v[211], v[202])) + k.emit(s_cmp_lt_u32(s[63], s[61])) + for i in range(2): + k.emit(v_mfma_fp4(v[244 + i * -8:247 + i * -8], v[180 + i * -4:183 + i * -4], v[44 + i * 8:47 + i * 8], 3 + i * -1, 0, v[211], v[202 + i * 1])) + k.emit(v_mfma_fp4(v[232 + i * 16:235 + i * 16], v[176 + i * 4:179 + i * 4], v[48:51], 0 + i * 1, 0, v[211], v[203])) + k.emit(s_cselect_b32(s[70 + i * 2], s[70 + i * 2], 0)) + for i in range(2): + k.emit(v_mfma_fp4(v[252 + i * -48:255 + i * -48], v[180 + i * 4:183 + i * 4], v[52 + i * 16:55 + i * 16], 3 + i * -1, 0 + i * 3, v[211 + i * -1], v[203])) + k.emit(v_mfma_fp4(v[192 + i * 24:195 + i * 24], v[184 + i * 4:187 + i * 4], v[56 + i * 8:59 + i * 8], 0 + i * 1, 3, v[210], v[202 + i * 1])) + k.emit(s_add_u32(s[16 + i * 8], s[16 + i * 8], s[70 + i * 2])) + k.emit(v_mfma_fp4(v[196 + i * 24:199 + i * 24], v[184 + i * 4:187 + i * 4], v[60 + i * 8:63 + i * 8], 2 + i * 1, 3, v[210], v[202 + i * 1])) + k.emit(v_mfma_fp4(v[208 + i * 16:211 + i * 16], v[188 + i * 4:191 + i * 4], v[56:59], 1 + i * -1, 3, v[210 + i * 1], v[202])) + k.emit(s_addc_u32(s[17 + i * 8], 0, s[17 + i * 8])) + k.emit(v_mfma_fp4(v[212 + i * 16:215 + i * 16], v[188 + i * 4:191 + i * 4], v[60:63], 3 + i * -1, 3, v[210 + i * 1], v[202])) + k.emit(v_mfma_fp4(v[200 + i * 40:203 + i * 40], v[184 + i * 12:187 + i * 12], v[64 + i * -8:67 + i * -8], 0 + i * 1, 3, v[210 + i * 1], v[203 + i * -1])) + k.emit(s_sub_u32(s[18 + i * 8], s[18 + i * 8], s[70 + i * 2])) + k.emit(v_mfma_fp4(v[244:247], v[196:199], v[60:63], 3, 3, v[211], v[202])) + k.emit(v_mfma_fp4(v[232:235], v[192:195], v[64:67], 0, 3, v[211], v[203])) + k.emit(s_addk_i32(s[60], 256)) + k.emit(v_mfma_fp4(v[236:239], v[192:195], v[68:71], 2, 3, v[211], v[203])) + k.emit(v_mfma_fp4(v[248:251], v[196:199], v[64:67], 1, 3, v[211], v[203])) + k.emit(s_cmp_lt_i32(s[60], s[61])) + k.emit(v_mfma_fp4(v[252:255], v[196:199], v[68:71], 3, 3, v[211], v[203])) + k.emit(s_cbranch_scc0(1), target='L0_37CC') + k.emit(s_branch(64253), target='L0_23C0') + k.label('L0_37CC') + k.emit(s_waitcnt()) + k.emit(s_barrier()) + k.emit(s_cmp_eq_u32(s[65], 0)) + k.emit(s_cbranch_scc1(2149), target='L0_5970') + k.emit(v_lshrrev_b32_e32(v[4], 4)) + k.emit(v_mul_i32_i24_e64(v[4], v[4], 8)) + k.emit(v_and_b32_e64(v[5], v[0], 15)) + k.emit(v_lshlrev_b32_e32(v[5], 8, v[5])) + k.emit(v_add_i32(v[4], v[4], v[5])) + k.emit(s_mul_i32(s[62], s[46], LIT, 16384)) + k.emit(s_add_i32(s[62], s[62], 0)) + k.emit(v_add_i32(v[4], v[4], s[62])) + for i in range(2): + for j9 in range(4): + k.emit(v_accvgpr_read(v[8 + j9 * 1 + i * 4], v[0 + j9 * 1 + i * 16])) + k.emit(v_mul_f32_e32(v[8 + j9 * 1 + i * 4], s[41], v[8 + j9 * 1 + i * 4])) + for i in range(4): + k.emit(v_cvt_pk_bf16_f32(v[16 + i * 1], v[8 + i * 2], v[9 + i * 2])) + k.emit(ds_write_b64(v[0], v[4], v[16:17])) + k.emit(ds_write_b64(v[0], v[4], v[18:19], v[0], 0, 32)) + for i in range(2): + for j10 in range(4): + k.emit(v_accvgpr_read(v[8 + j10 * 1 + i * 4], v[32 + j10 * 1 + i * 16])) + k.emit(v_mul_f32_e32(v[8 + j10 * 1 + i * 4], s[41], v[8 + j10 * 1 + i * 4])) + for i in range(4): + k.emit(v_cvt_pk_bf16_f32(v[16 + i * 1], v[8 + i * 2], v[9 + i * 2])) + k.emit(ds_write_b64(v[0], v[4], v[16:17], v[0], 0, 64)) + k.emit(ds_write_b64(v[0], v[4], v[18:19], v[0], 0, 96)) + for i in range(3): + for j11 in range(2): + k.emit(v_accvgpr_read(v[8], v[4 + j11 * 32 + i * 4])) + k.emit(v_mul_f32_e32(v[8], s[41], v[8])) + k.emit(v_accvgpr_read(v[9], v[5 + j11 * 32 + i * 4])) + k.emit(v_mul_f32_e32(v[9], s[41], v[9])) + k.emit(v_accvgpr_read(v[10], v[6 + j11 * 32 + i * 4])) + k.emit(v_mul_f32_e32(v[10], s[41], v[10])) + k.emit(v_accvgpr_read(v[11], v[7 + j11 * 32 + i * 4])) + k.emit(v_mul_f32_e32(v[11], s[41], v[11])) + k.emit(v_accvgpr_read(v[12], v[20 + j11 * 32 + i * 4])) + k.emit(v_mul_f32_e32(v[12], s[41], v[12])) + k.emit(v_accvgpr_read(v[13], v[21 + j11 * 32 + i * 4])) + k.emit(v_mul_f32_e32(v[13], s[41], v[13])) + k.emit(v_accvgpr_read(v[14], v[22 + j11 * 32 + i * 4])) + k.emit(v_mul_f32_e32(v[14], s[41], v[14])) + k.emit(v_accvgpr_read(v[15], v[23 + j11 * 32 + i * 4])) + k.emit(v_mul_f32_e32(v[15], s[41], v[15])) + k.emit(v_cvt_pk_bf16_f32(v[16], v[8], v[9])) + k.emit(v_cvt_pk_bf16_f32(v[17], v[10], v[11])) + k.emit(v_cvt_pk_bf16_f32(v[18], v[12], v[13])) + k.emit(v_cvt_pk_bf16_f32(v[19], v[14], v[15])) + k.emit(ds_write_b64(v[0], v[4], v[16:17], v[0], 0, 0 + j11 * 64, 16 + i * 16)) + k.emit(ds_write_b64(v[0], v[4], v[18:19], v[0], 0, 32 + j11 * 64, 16 + i * 16)) + for i in range(2): + for j12 in range(2): + k.emit(v_accvgpr_read(v[8 + j12 * 4], v[64 + j12 * 16 + i * 32])) + k.emit(v_mul_f32_e32(v[8 + j12 * 4], s[41], v[8 + j12 * 4])) + k.emit(v_accvgpr_read(v[9 + j12 * 4], v[65 + j12 * 16 + i * 32])) + k.emit(v_mul_f32_e32(v[9 + j12 * 4], s[41], v[9 + j12 * 4])) + k.emit(v_accvgpr_read(v[10 + j12 * 4], v[66 + j12 * 16 + i * 32])) + k.emit(v_mul_f32_e32(v[10 + j12 * 4], s[41], v[10 + j12 * 4])) + k.emit(v_accvgpr_read(v[11 + j12 * 4], v[67 + j12 * 16 + i * 32])) + k.emit(v_mul_f32_e32(v[11 + j12 * 4], s[41], v[11 + j12 * 4])) + for j13 in range(4): + k.emit(v_cvt_pk_bf16_f32(v[16 + j13 * 1], v[8 + j13 * 2], v[9 + j13 * 2])) + k.emit(ds_write_b64(v[0], v[4], v[16:17], v[0], 0, 128 + i * 64)) + k.emit(ds_write_b64(v[0], v[4], v[18:19], v[0], 0, 160 + i * 64)) + for i in range(3): + for j14 in range(2): + k.emit(v_accvgpr_read(v[8], v[68 + j14 * 32 + i * 4])) + k.emit(v_mul_f32_e32(v[8], s[41], v[8])) + k.emit(v_accvgpr_read(v[9], v[69 + j14 * 32 + i * 4])) + k.emit(v_mul_f32_e32(v[9], s[41], v[9])) + k.emit(v_accvgpr_read(v[10], v[70 + j14 * 32 + i * 4])) + k.emit(v_mul_f32_e32(v[10], s[41], v[10])) + k.emit(v_accvgpr_read(v[11], v[71 + j14 * 32 + i * 4])) + k.emit(v_mul_f32_e32(v[11], s[41], v[11])) + k.emit(v_accvgpr_read(v[12], v[84 + j14 * 32 + i * 4])) + k.emit(v_mul_f32_e32(v[12], s[41], v[12])) + k.emit(v_accvgpr_read(v[13], v[85 + j14 * 32 + i * 4])) + k.emit(v_mul_f32_e32(v[13], s[41], v[13])) + k.emit(v_accvgpr_read(v[14], v[86 + j14 * 32 + i * 4])) + k.emit(v_mul_f32_e32(v[14], s[41], v[14])) + k.emit(v_accvgpr_read(v[15], v[87 + j14 * 32 + i * 4])) + k.emit(v_mul_f32_e32(v[15], s[41], v[15])) + k.emit(v_cvt_pk_bf16_f32(v[16], v[8], v[9])) + k.emit(v_cvt_pk_bf16_f32(v[17], v[10], v[11])) + k.emit(v_cvt_pk_bf16_f32(v[18], v[12], v[13])) + k.emit(v_cvt_pk_bf16_f32(v[19], v[14], v[15])) + k.emit(ds_write_b64(v[0], v[4], v[16:17], v[0], 0, 128 + j14 * 64, 16 + i * 16)) + k.emit(ds_write_b64(v[0], v[4], v[18:19], v[0], 0, 160 + j14 * 64, 16 + i * 16)) + k.emit(s_waitcnt(49279)) + k.emit(v_mul_i32_i24_e64(v[4], v[0], 4)) + k.emit(v_add_i32(v[4], v[4], s[62])) + k.emit(s_mul_i32(s[63], s[36], 0)) + k.emit(v_add_u32_e32(v[244], s[63], v[240])) + k.emit(ds_read_b32(v[16], v[4])) + for i in range(15): + k.emit(ds_read_b32(v[17], v[4], v[0], v[0], 0, 0, 1 + i * 4)) + k.emit(ds_read_b32(v[18], v[4], v[0], v[0], 0, 0, 2 + i * 4)) + k.emit(ds_read_b32(v[19], v[4], v[0], v[0], 0, 0, 3 + i * 4)) + for j15 in range(4): + k.emit(s_waitcnt(50047 + j15 * -256)) + k.emit(buffer_atomic_pk_add_bf16(v[16 + j15 * 1], v[244], s[4:7], 0, 0, 1)) + k.emit(v_add_u32_e64(v[244], v[244], s[36])) + k.emit(s_mul_i32(s[63], s[36], 4 + i * 4)) + k.emit(v_add_u32_e32(v[244], s[63], v[240])) + k.emit(ds_read_b32(v[16], v[4], v[0], v[0], 0, 0, 4 + i * 4)) + k.emit(ds_read_b32(v[17], v[4], v[0], v[0], 0, 0, 61)) + k.emit(ds_read_b32(v[18], v[4], v[0], v[0], 0, 0, 62)) + k.emit(ds_read_b32(v[19], v[4], v[0], v[0], 0, 0, 63)) + for i in range(4): + k.emit(s_waitcnt(50047 + i * -256)) + k.emit(buffer_atomic_pk_add_bf16(v[16 + i * 1], v[244], s[4:7], 0, 0, 1)) + k.emit(v_add_u32_e64(v[244], v[244], s[36])) + k.emit(v_lshrrev_b32_e32(v[4], 4)) + k.emit(v_mul_i32_i24_e64(v[4], v[4], 8)) + k.emit(v_and_b32_e64(v[5], v[0], 15)) + k.emit(v_lshlrev_b32_e32(v[5], 8, v[5])) + k.emit(v_add_i32(v[4], v[4], v[5])) + k.emit(s_mul_i32(s[62], s[46], LIT, 16384)) + k.emit(s_add_i32(s[62], s[62], 0)) + k.emit(v_add_i32(v[4], v[4], s[62])) + for i in range(2): + for j16 in range(4): + k.emit(v_accvgpr_read(v[8 + j16 * 1 + i * 4], v[128 + j16 * 1 + i * 16])) + k.emit(v_mul_f32_e32(v[8 + j16 * 1 + i * 4], s[41], v[8 + j16 * 1 + i * 4])) + for i in range(4): + k.emit(v_cvt_pk_bf16_f32(v[16 + i * 1], v[8 + i * 2], v[9 + i * 2])) + k.emit(ds_write_b64(v[0], v[4], v[16:17])) + k.emit(ds_write_b64(v[0], v[4], v[18:19], v[0], 0, 32)) + for i in range(2): + for j17 in range(4): + k.emit(v_accvgpr_read(v[8 + j17 * 1 + i * 4], v[160 + j17 * 1 + i * 16])) + k.emit(v_mul_f32_e32(v[8 + j17 * 1 + i * 4], s[41], v[8 + j17 * 1 + i * 4])) + for i in range(4): + k.emit(v_cvt_pk_bf16_f32(v[16 + i * 1], v[8 + i * 2], v[9 + i * 2])) + k.emit(ds_write_b64(v[0], v[4], v[16:17], v[0], 0, 64)) + k.emit(ds_write_b64(v[0], v[4], v[18:19], v[0], 0, 96)) + for i in range(3): + for j18 in range(2): + k.emit(v_accvgpr_read(v[8], v[132 + j18 * 32 + i * 4])) + k.emit(v_mul_f32_e32(v[8], s[41], v[8])) + k.emit(v_accvgpr_read(v[9], v[133 + j18 * 32 + i * 4])) + k.emit(v_mul_f32_e32(v[9], s[41], v[9])) + k.emit(v_accvgpr_read(v[10], v[134 + j18 * 32 + i * 4])) + k.emit(v_mul_f32_e32(v[10], s[41], v[10])) + k.emit(v_accvgpr_read(v[11], v[135 + j18 * 32 + i * 4])) + k.emit(v_mul_f32_e32(v[11], s[41], v[11])) + k.emit(v_accvgpr_read(v[12], v[148 + j18 * 32 + i * 4])) + k.emit(v_mul_f32_e32(v[12], s[41], v[12])) + k.emit(v_accvgpr_read(v[13], v[149 + j18 * 32 + i * 4])) + k.emit(v_mul_f32_e32(v[13], s[41], v[13])) + k.emit(v_accvgpr_read(v[14], v[150 + j18 * 32 + i * 4])) + k.emit(v_mul_f32_e32(v[14], s[41], v[14])) + k.emit(v_accvgpr_read(v[15], v[151 + j18 * 32 + i * 4])) + k.emit(v_mul_f32_e32(v[15], s[41], v[15])) + k.emit(v_cvt_pk_bf16_f32(v[16], v[8], v[9])) + k.emit(v_cvt_pk_bf16_f32(v[17], v[10], v[11])) + k.emit(v_cvt_pk_bf16_f32(v[18], v[12], v[13])) + k.emit(v_cvt_pk_bf16_f32(v[19], v[14], v[15])) + k.emit(ds_write_b64(v[0], v[4], v[16:17], v[0], 0, 0 + j18 * 64, 16 + i * 16)) + k.emit(ds_write_b64(v[0], v[4], v[18:19], v[0], 0, 32 + j18 * 64, 16 + i * 16)) + for i in range(2): + for j19 in range(2): + k.emit(v_accvgpr_read(v[8 + j19 * 4], v[192 + j19 * 16 + i * 32])) + k.emit(v_mul_f32_e32(v[8 + j19 * 4], s[41], v[8 + j19 * 4])) + k.emit(v_accvgpr_read(v[9 + j19 * 4], v[193 + j19 * 16 + i * 32])) + k.emit(v_mul_f32_e32(v[9 + j19 * 4], s[41], v[9 + j19 * 4])) + k.emit(v_accvgpr_read(v[10 + j19 * 4], v[194 + j19 * 16 + i * 32])) + k.emit(v_mul_f32_e32(v[10 + j19 * 4], s[41], v[10 + j19 * 4])) + k.emit(v_accvgpr_read(v[11 + j19 * 4], v[195 + j19 * 16 + i * 32])) + k.emit(v_mul_f32_e32(v[11 + j19 * 4], s[41], v[11 + j19 * 4])) + for j20 in range(4): + k.emit(v_cvt_pk_bf16_f32(v[16 + j20 * 1], v[8 + j20 * 2], v[9 + j20 * 2])) + k.emit(ds_write_b64(v[0], v[4], v[16:17], v[0], 0, 128 + i * 64)) + k.emit(ds_write_b64(v[0], v[4], v[18:19], v[0], 0, 160 + i * 64)) + for i in range(3): + for j21 in range(2): + k.emit(v_accvgpr_read(v[8], v[196 + j21 * 32 + i * 4])) + k.emit(v_mul_f32_e32(v[8], s[41], v[8])) + k.emit(v_accvgpr_read(v[9], v[197 + j21 * 32 + i * 4])) + k.emit(v_mul_f32_e32(v[9], s[41], v[9])) + k.emit(v_accvgpr_read(v[10], v[198 + j21 * 32 + i * 4])) + k.emit(v_mul_f32_e32(v[10], s[41], v[10])) + k.emit(v_accvgpr_read(v[11], v[199 + j21 * 32 + i * 4])) + k.emit(v_mul_f32_e32(v[11], s[41], v[11])) + k.emit(v_accvgpr_read(v[12], v[212 + j21 * 32 + i * 4])) + k.emit(v_mul_f32_e32(v[12], s[41], v[12])) + k.emit(v_accvgpr_read(v[13], v[213 + j21 * 32 + i * 4])) + k.emit(v_mul_f32_e32(v[13], s[41], v[13])) + k.emit(v_accvgpr_read(v[14], v[214 + j21 * 32 + i * 4])) + k.emit(v_mul_f32_e32(v[14], s[41], v[14])) + k.emit(v_accvgpr_read(v[15], v[215 + j21 * 32 + i * 4])) + k.emit(v_mul_f32_e32(v[15], s[41], v[15])) + k.emit(v_cvt_pk_bf16_f32(v[16], v[8], v[9])) + k.emit(v_cvt_pk_bf16_f32(v[17], v[10], v[11])) + k.emit(v_cvt_pk_bf16_f32(v[18], v[12], v[13])) + k.emit(v_cvt_pk_bf16_f32(v[19], v[14], v[15])) + k.emit(ds_write_b64(v[0], v[4], v[16:17], v[0], 0, 128 + j21 * 64, 16 + i * 16)) + k.emit(ds_write_b64(v[0], v[4], v[18:19], v[0], 0, 160 + j21 * 64, 16 + i * 16)) + k.emit(s_waitcnt(49279)) + k.emit(v_mul_i32_i24_e64(v[4], v[0], 4)) + k.emit(v_add_i32(v[4], v[4], s[62])) + k.emit(s_mul_i32(s[63], s[36], 0)) + k.emit(v_add_u32_e32(v[244], s[63], v[242])) + k.emit(ds_read_b32(v[16], v[4])) + for i in range(15): + k.emit(ds_read_b32(v[17], v[4], v[0], v[0], 0, 0, 1 + i * 4)) + k.emit(ds_read_b32(v[18], v[4], v[0], v[0], 0, 0, 2 + i * 4)) + k.emit(ds_read_b32(v[19], v[4], v[0], v[0], 0, 0, 3 + i * 4)) + for j22 in range(4): + k.emit(s_waitcnt(50047 + j22 * -256)) + k.emit(buffer_atomic_pk_add_bf16(v[16 + j22 * 1], v[244], s[4:7], 0, 0, 1)) + k.emit(v_add_u32_e64(v[244], v[244], s[36])) + k.emit(s_mul_i32(s[63], s[36], 4 + i * 4)) + k.emit(v_add_u32_e32(v[244], s[63], v[242])) + k.emit(ds_read_b32(v[16], v[4], v[0], v[0], 0, 0, 4 + i * 4)) + k.emit(ds_read_b32(v[17], v[4], v[0], v[0], 0, 0, 61)) + k.emit(ds_read_b32(v[18], v[4], v[0], v[0], 0, 0, 62)) + k.emit(ds_read_b32(v[19], v[4], v[0], v[0], 0, 0, 63)) + for i in range(4): + k.emit(s_waitcnt(50047 + i * -256)) + k.emit(buffer_atomic_pk_add_bf16(v[16 + i * 1], v[244], s[4:7], 0, 0, 1)) + k.emit(v_add_u32_e64(v[244], v[244], s[36])) + k.emit(s_branch(1344), target='L0_6E70') + k.label('L0_5970') + k.emit(s_mul_i32(s[62], s[36], 0)) + k.emit(v_add_u32_e32(v[244], s[62], v[240])) + for i in range(2): + for j23 in range(2): + k.emit(v_accvgpr_read(v[8 + j23 * 4], v[0 + j23 * 16 + i * 32])) + k.emit(v_mul_f32_e32(v[8 + j23 * 4], s[41], v[8 + j23 * 4])) + k.emit(v_accvgpr_read(v[9 + j23 * 4], v[1 + j23 * 16 + i * 32])) + k.emit(v_mul_f32_e32(v[9 + j23 * 4], s[41], v[9 + j23 * 4])) + k.emit(v_accvgpr_read(v[10 + j23 * 4], v[2 + j23 * 16 + i * 32])) + k.emit(v_mul_f32_e32(v[10 + j23 * 4], s[41], v[10 + j23 * 4])) + k.emit(v_accvgpr_read(v[11 + j23 * 4], v[3 + j23 * 16 + i * 32])) + k.emit(v_mul_f32_e32(v[11 + j23 * 4], s[41], v[11 + j23 * 4])) + for j24 in range(4): + k.emit(v_cvt_pk_bf16_f32(v[16 + j24 * 1], v[8 + j24 * 2], v[9 + j24 * 2])) + k.emit(s_nop(1)) + k.emit(v_permlane16_swap_b32_e32(v[16], v[18])) + k.emit(s_nop(1)) + k.emit(v_permlane16_swap_b32_e32(v[17], v[19])) + k.emit(s_nop(1)) + k.emit(buffer_store_dwordx4(v[16:19], v[244], s[4:7], 0, 0, 1)) + k.emit(v_add_i32(v[244], v[244], 64)) + k.emit(s_mul_i32(s[62], s[36], 16)) + k.emit(v_add_u32_e32(v[244], s[62], v[240])) + for i in range(2): + for j25 in range(2): + k.emit(v_accvgpr_read(v[8 + j25 * 4], v[4 + j25 * 16 + i * 32])) + k.emit(v_mul_f32_e32(v[8 + j25 * 4], s[41], v[8 + j25 * 4])) + k.emit(v_accvgpr_read(v[9 + j25 * 4], v[5 + j25 * 16 + i * 32])) + k.emit(v_mul_f32_e32(v[9 + j25 * 4], s[41], v[9 + j25 * 4])) + k.emit(v_accvgpr_read(v[10 + j25 * 4], v[6 + j25 * 16 + i * 32])) + k.emit(v_mul_f32_e32(v[10 + j25 * 4], s[41], v[10 + j25 * 4])) + k.emit(v_accvgpr_read(v[11 + j25 * 4], v[7 + j25 * 16 + i * 32])) + k.emit(v_mul_f32_e32(v[11 + j25 * 4], s[41], v[11 + j25 * 4])) + for j26 in range(4): + k.emit(v_cvt_pk_bf16_f32(v[16 + j26 * 1], v[8 + j26 * 2], v[9 + j26 * 2])) + k.emit(s_nop(1)) + k.emit(v_permlane16_swap_b32_e32(v[16], v[18])) + k.emit(s_nop(1)) + k.emit(v_permlane16_swap_b32_e32(v[17], v[19])) + k.emit(s_nop(1)) + k.emit(buffer_store_dwordx4(v[16:19], v[244], s[4:7], 0, 0, 1)) + k.emit(v_add_i32(v[244], v[244], 64)) + k.emit(s_mul_i32(s[62], s[36], 32)) + k.emit(v_add_u32_e32(v[244], s[62], v[240])) + for i in range(2): + for j27 in range(2): + k.emit(v_accvgpr_read(v[8 + j27 * 4], v[8 + j27 * 16 + i * 32])) + k.emit(v_mul_f32_e32(v[8 + j27 * 4], s[41], v[8 + j27 * 4])) + k.emit(v_accvgpr_read(v[9 + j27 * 4], v[9 + j27 * 16 + i * 32])) + k.emit(v_mul_f32_e32(v[9 + j27 * 4], s[41], v[9 + j27 * 4])) + k.emit(v_accvgpr_read(v[10 + j27 * 4], v[10 + j27 * 16 + i * 32])) + k.emit(v_mul_f32_e32(v[10 + j27 * 4], s[41], v[10 + j27 * 4])) + k.emit(v_accvgpr_read(v[11 + j27 * 4], v[11 + j27 * 16 + i * 32])) + k.emit(v_mul_f32_e32(v[11 + j27 * 4], s[41], v[11 + j27 * 4])) + for j28 in range(4): + k.emit(v_cvt_pk_bf16_f32(v[16 + j28 * 1], v[8 + j28 * 2], v[9 + j28 * 2])) + k.emit(s_nop(1)) + k.emit(v_permlane16_swap_b32_e32(v[16], v[18])) + k.emit(s_nop(1)) + k.emit(v_permlane16_swap_b32_e32(v[17], v[19])) + k.emit(s_nop(1)) + k.emit(buffer_store_dwordx4(v[16:19], v[244], s[4:7], 0, 0, 1)) + k.emit(v_add_i32(v[244], v[244], 64)) + k.emit(s_mul_i32(s[62], s[36], 48)) + k.emit(v_add_u32_e32(v[244], s[62], v[240])) + for i in range(2): + for j29 in range(2): + k.emit(v_accvgpr_read(v[8 + j29 * 4], v[12 + j29 * 16 + i * 32])) + k.emit(v_mul_f32_e32(v[8 + j29 * 4], s[41], v[8 + j29 * 4])) + k.emit(v_accvgpr_read(v[9 + j29 * 4], v[13 + j29 * 16 + i * 32])) + k.emit(v_mul_f32_e32(v[9 + j29 * 4], s[41], v[9 + j29 * 4])) + k.emit(v_accvgpr_read(v[10 + j29 * 4], v[14 + j29 * 16 + i * 32])) + k.emit(v_mul_f32_e32(v[10 + j29 * 4], s[41], v[10 + j29 * 4])) + k.emit(v_accvgpr_read(v[11 + j29 * 4], v[15 + j29 * 16 + i * 32])) + k.emit(v_mul_f32_e32(v[11 + j29 * 4], s[41], v[11 + j29 * 4])) + for j30 in range(4): + k.emit(v_cvt_pk_bf16_f32(v[16 + j30 * 1], v[8 + j30 * 2], v[9 + j30 * 2])) + k.emit(s_nop(1)) + k.emit(v_permlane16_swap_b32_e32(v[16], v[18])) + k.emit(s_nop(1)) + k.emit(v_permlane16_swap_b32_e32(v[17], v[19])) + k.emit(s_nop(1)) + k.emit(buffer_store_dwordx4(v[16:19], v[244], s[4:7], 0, 0, 1)) + k.emit(v_add_i32(v[244], v[244], 64)) + k.emit(s_mul_i32(s[62], s[36], 0)) + k.emit(v_add_u32_e32(v[244], s[62], v[241])) + for i in range(2): + for j31 in range(2): + k.emit(v_accvgpr_read(v[8 + j31 * 4], v[64 + j31 * 16 + i * 32])) + k.emit(v_mul_f32_e32(v[8 + j31 * 4], s[41], v[8 + j31 * 4])) + k.emit(v_accvgpr_read(v[9 + j31 * 4], v[65 + j31 * 16 + i * 32])) + k.emit(v_mul_f32_e32(v[9 + j31 * 4], s[41], v[9 + j31 * 4])) + k.emit(v_accvgpr_read(v[10 + j31 * 4], v[66 + j31 * 16 + i * 32])) + k.emit(v_mul_f32_e32(v[10 + j31 * 4], s[41], v[10 + j31 * 4])) + k.emit(v_accvgpr_read(v[11 + j31 * 4], v[67 + j31 * 16 + i * 32])) + k.emit(v_mul_f32_e32(v[11 + j31 * 4], s[41], v[11 + j31 * 4])) + for j32 in range(4): + k.emit(v_cvt_pk_bf16_f32(v[16 + j32 * 1], v[8 + j32 * 2], v[9 + j32 * 2])) + k.emit(s_nop(1)) + k.emit(v_permlane16_swap_b32_e32(v[16], v[18])) + k.emit(s_nop(1)) + k.emit(v_permlane16_swap_b32_e32(v[17], v[19])) + k.emit(s_nop(1)) + k.emit(buffer_store_dwordx4(v[16:19], v[244], s[4:7], 0, 0, 1)) + k.emit(v_add_i32(v[244], v[244], 64)) + k.emit(s_mul_i32(s[62], s[36], 16)) + k.emit(v_add_u32_e32(v[244], s[62], v[241])) + for i in range(2): + for j33 in range(2): + k.emit(v_accvgpr_read(v[8 + j33 * 4], v[68 + j33 * 16 + i * 32])) + k.emit(v_mul_f32_e32(v[8 + j33 * 4], s[41], v[8 + j33 * 4])) + k.emit(v_accvgpr_read(v[9 + j33 * 4], v[69 + j33 * 16 + i * 32])) + k.emit(v_mul_f32_e32(v[9 + j33 * 4], s[41], v[9 + j33 * 4])) + k.emit(v_accvgpr_read(v[10 + j33 * 4], v[70 + j33 * 16 + i * 32])) + k.emit(v_mul_f32_e32(v[10 + j33 * 4], s[41], v[10 + j33 * 4])) + k.emit(v_accvgpr_read(v[11 + j33 * 4], v[71 + j33 * 16 + i * 32])) + k.emit(v_mul_f32_e32(v[11 + j33 * 4], s[41], v[11 + j33 * 4])) + for j34 in range(4): + k.emit(v_cvt_pk_bf16_f32(v[16 + j34 * 1], v[8 + j34 * 2], v[9 + j34 * 2])) + k.emit(s_nop(1)) + k.emit(v_permlane16_swap_b32_e32(v[16], v[18])) + k.emit(s_nop(1)) + k.emit(v_permlane16_swap_b32_e32(v[17], v[19])) + k.emit(s_nop(1)) + k.emit(buffer_store_dwordx4(v[16:19], v[244], s[4:7], 0, 0, 1)) + k.emit(v_add_i32(v[244], v[244], 64)) + k.emit(s_mul_i32(s[62], s[36], 32)) + k.emit(v_add_u32_e32(v[244], s[62], v[241])) + for i in range(2): + for j35 in range(2): + k.emit(v_accvgpr_read(v[8 + j35 * 4], v[72 + j35 * 16 + i * 32])) + k.emit(v_mul_f32_e32(v[8 + j35 * 4], s[41], v[8 + j35 * 4])) + k.emit(v_accvgpr_read(v[9 + j35 * 4], v[73 + j35 * 16 + i * 32])) + k.emit(v_mul_f32_e32(v[9 + j35 * 4], s[41], v[9 + j35 * 4])) + k.emit(v_accvgpr_read(v[10 + j35 * 4], v[74 + j35 * 16 + i * 32])) + k.emit(v_mul_f32_e32(v[10 + j35 * 4], s[41], v[10 + j35 * 4])) + k.emit(v_accvgpr_read(v[11 + j35 * 4], v[75 + j35 * 16 + i * 32])) + k.emit(v_mul_f32_e32(v[11 + j35 * 4], s[41], v[11 + j35 * 4])) + for j36 in range(4): + k.emit(v_cvt_pk_bf16_f32(v[16 + j36 * 1], v[8 + j36 * 2], v[9 + j36 * 2])) + k.emit(s_nop(1)) + k.emit(v_permlane16_swap_b32_e32(v[16], v[18])) + k.emit(s_nop(1)) + k.emit(v_permlane16_swap_b32_e32(v[17], v[19])) + k.emit(s_nop(1)) + k.emit(buffer_store_dwordx4(v[16:19], v[244], s[4:7], 0, 0, 1)) + k.emit(v_add_i32(v[244], v[244], 64)) + k.emit(s_mul_i32(s[62], s[36], 48)) + k.emit(v_add_u32_e32(v[244], s[62], v[241])) + for i in range(2): + for j37 in range(2): + k.emit(v_accvgpr_read(v[8 + j37 * 4], v[76 + j37 * 16 + i * 32])) + k.emit(v_mul_f32_e32(v[8 + j37 * 4], s[41], v[8 + j37 * 4])) + k.emit(v_accvgpr_read(v[9 + j37 * 4], v[77 + j37 * 16 + i * 32])) + k.emit(v_mul_f32_e32(v[9 + j37 * 4], s[41], v[9 + j37 * 4])) + k.emit(v_accvgpr_read(v[10 + j37 * 4], v[78 + j37 * 16 + i * 32])) + k.emit(v_mul_f32_e32(v[10 + j37 * 4], s[41], v[10 + j37 * 4])) + k.emit(v_accvgpr_read(v[11 + j37 * 4], v[79 + j37 * 16 + i * 32])) + k.emit(v_mul_f32_e32(v[11 + j37 * 4], s[41], v[11 + j37 * 4])) + for j38 in range(4): + k.emit(v_cvt_pk_bf16_f32(v[16 + j38 * 1], v[8 + j38 * 2], v[9 + j38 * 2])) + k.emit(s_nop(1)) + k.emit(v_permlane16_swap_b32_e32(v[16], v[18])) + k.emit(s_nop(1)) + k.emit(v_permlane16_swap_b32_e32(v[17], v[19])) + k.emit(s_nop(1)) + k.emit(buffer_store_dwordx4(v[16:19], v[244], s[4:7], 0, 0, 1)) + k.emit(v_add_i32(v[244], v[244], 64)) + k.emit(s_mul_i32(s[62], s[36], 0)) + k.emit(v_add_u32_e32(v[244], s[62], v[242])) + for i in range(2): + for j39 in range(2): + k.emit(v_accvgpr_read(v[8 + j39 * 4], v[128 + j39 * 16 + i * 32])) + k.emit(v_mul_f32_e32(v[8 + j39 * 4], s[41], v[8 + j39 * 4])) + k.emit(v_accvgpr_read(v[9 + j39 * 4], v[129 + j39 * 16 + i * 32])) + k.emit(v_mul_f32_e32(v[9 + j39 * 4], s[41], v[9 + j39 * 4])) + k.emit(v_accvgpr_read(v[10 + j39 * 4], v[130 + j39 * 16 + i * 32])) + k.emit(v_mul_f32_e32(v[10 + j39 * 4], s[41], v[10 + j39 * 4])) + k.emit(v_accvgpr_read(v[11 + j39 * 4], v[131 + j39 * 16 + i * 32])) + k.emit(v_mul_f32_e32(v[11 + j39 * 4], s[41], v[11 + j39 * 4])) + for j40 in range(4): + k.emit(v_cvt_pk_bf16_f32(v[16 + j40 * 1], v[8 + j40 * 2], v[9 + j40 * 2])) + k.emit(s_nop(1)) + k.emit(v_permlane16_swap_b32_e32(v[16], v[18])) + k.emit(s_nop(1)) + k.emit(v_permlane16_swap_b32_e32(v[17], v[19])) + k.emit(s_nop(1)) + k.emit(buffer_store_dwordx4(v[16:19], v[244], s[4:7], 0, 0, 1)) + k.emit(v_add_i32(v[244], v[244], 64)) + k.emit(s_mul_i32(s[62], s[36], 16)) + k.emit(v_add_u32_e32(v[244], s[62], v[242])) + for i in range(2): + for j41 in range(2): + k.emit(v_accvgpr_read(v[8 + j41 * 4], v[132 + j41 * 16 + i * 32])) + k.emit(v_mul_f32_e32(v[8 + j41 * 4], s[41], v[8 + j41 * 4])) + k.emit(v_accvgpr_read(v[9 + j41 * 4], v[133 + j41 * 16 + i * 32])) + k.emit(v_mul_f32_e32(v[9 + j41 * 4], s[41], v[9 + j41 * 4])) + k.emit(v_accvgpr_read(v[10 + j41 * 4], v[134 + j41 * 16 + i * 32])) + k.emit(v_mul_f32_e32(v[10 + j41 * 4], s[41], v[10 + j41 * 4])) + k.emit(v_accvgpr_read(v[11 + j41 * 4], v[135 + j41 * 16 + i * 32])) + k.emit(v_mul_f32_e32(v[11 + j41 * 4], s[41], v[11 + j41 * 4])) + for j42 in range(4): + k.emit(v_cvt_pk_bf16_f32(v[16 + j42 * 1], v[8 + j42 * 2], v[9 + j42 * 2])) + k.emit(s_nop(1)) + k.emit(v_permlane16_swap_b32_e32(v[16], v[18])) + k.emit(s_nop(1)) + k.emit(v_permlane16_swap_b32_e32(v[17], v[19])) + k.emit(s_nop(1)) + k.emit(buffer_store_dwordx4(v[16:19], v[244], s[4:7], 0, 0, 1)) + k.emit(v_add_i32(v[244], v[244], 64)) + k.emit(s_mul_i32(s[62], s[36], 32)) + for i in range(2): + k.emit(v_add_u32_e32(v[244], s[62], v[242 + i * 1])) + for j43 in range(2): + k.emit(v_accvgpr_read(v[8], v[136 + j43 * 32 + i * 56])) + k.emit(v_mul_f32_e32(v[8], s[41], v[8])) + k.emit(v_accvgpr_read(v[9], v[137 + j43 * 32 + i * 56])) + k.emit(v_mul_f32_e32(v[9], s[41], v[9])) + k.emit(v_accvgpr_read(v[10], v[138 + j43 * 32 + i * 56])) + k.emit(v_mul_f32_e32(v[10], s[41], v[10])) + k.emit(v_accvgpr_read(v[11], v[139 + j43 * 32 + i * 56])) + k.emit(v_mul_f32_e32(v[11], s[41], v[11])) + k.emit(v_accvgpr_read(v[12], v[152 + j43 * 32 + i * 56])) + k.emit(v_mul_f32_e32(v[12], s[41], v[12])) + k.emit(v_accvgpr_read(v[13], v[153 + j43 * 32 + i * 56])) + k.emit(v_mul_f32_e32(v[13], s[41], v[13])) + k.emit(v_accvgpr_read(v[14], v[154 + j43 * 32 + i * 56])) + k.emit(v_mul_f32_e32(v[14], s[41], v[14])) + k.emit(v_accvgpr_read(v[15], v[155 + j43 * 32 + i * 56])) + k.emit(v_mul_f32_e32(v[15], s[41], v[15])) + k.emit(v_cvt_pk_bf16_f32(v[16], v[8], v[9])) + k.emit(v_cvt_pk_bf16_f32(v[17], v[10], v[11])) + k.emit(v_cvt_pk_bf16_f32(v[18], v[12], v[13])) + k.emit(v_cvt_pk_bf16_f32(v[19], v[14], v[15])) + k.emit(s_nop(1)) + k.emit(v_permlane16_swap_b32_e32(v[16], v[18])) + k.emit(s_nop(1)) + k.emit(v_permlane16_swap_b32_e32(v[17], v[19])) + k.emit(s_nop(1)) + k.emit(buffer_store_dwordx4(v[16:19], v[244], s[4:7], 0, 0, 1)) + k.emit(v_add_i32(v[244], v[244], 64)) + k.emit(s_mul_i32(s[62], s[36], 48 + i * -32)) + k.emit(v_add_u32_e32(v[244], s[62], v[242 + i * 1])) + for j44 in range(2): + k.emit(v_accvgpr_read(v[8], v[140 + j44 * 32 + i * 56])) + k.emit(v_mul_f32_e32(v[8], s[41], v[8])) + k.emit(v_accvgpr_read(v[9], v[141 + j44 * 32 + i * 56])) + k.emit(v_mul_f32_e32(v[9], s[41], v[9])) + k.emit(v_accvgpr_read(v[10], v[142 + j44 * 32 + i * 56])) + k.emit(v_mul_f32_e32(v[10], s[41], v[10])) + k.emit(v_accvgpr_read(v[11], v[143 + j44 * 32 + i * 56])) + k.emit(v_mul_f32_e32(v[11], s[41], v[11])) + k.emit(v_accvgpr_read(v[12], v[156 + j44 * 32 + i * 56])) + k.emit(v_mul_f32_e32(v[12], s[41], v[12])) + k.emit(v_accvgpr_read(v[13], v[157 + j44 * 32 + i * 56])) + k.emit(v_mul_f32_e32(v[13], s[41], v[13])) + k.emit(v_accvgpr_read(v[14], v[158 + j44 * 32 + i * 56])) + k.emit(v_mul_f32_e32(v[14], s[41], v[14])) + k.emit(v_accvgpr_read(v[15], v[159 + j44 * 32 + i * 56])) + k.emit(v_mul_f32_e32(v[15], s[41], v[15])) + k.emit(v_cvt_pk_bf16_f32(v[16], v[8], v[9])) + k.emit(v_cvt_pk_bf16_f32(v[17], v[10], v[11])) + k.emit(v_cvt_pk_bf16_f32(v[18], v[12], v[13])) + k.emit(v_cvt_pk_bf16_f32(v[19], v[14], v[15])) + k.emit(s_nop(1)) + k.emit(v_permlane16_swap_b32_e32(v[16], v[18])) + k.emit(s_nop(1)) + k.emit(v_permlane16_swap_b32_e32(v[17], v[19])) + k.emit(s_nop(1)) + k.emit(buffer_store_dwordx4(v[16:19], v[244], s[4:7], 0, 0, 1)) + k.emit(v_add_i32(v[244], v[244], 64)) + k.emit(s_mul_i32(s[62], s[36], 0 + i * 32)) + k.emit(v_add_u32_e32(v[244], s[62], v[243])) + for i in range(2): + for j45 in range(2): + k.emit(v_accvgpr_read(v[8 + j45 * 4], v[200 + j45 * 16 + i * 32])) + k.emit(v_mul_f32_e32(v[8 + j45 * 4], s[41], v[8 + j45 * 4])) + k.emit(v_accvgpr_read(v[9 + j45 * 4], v[201 + j45 * 16 + i * 32])) + k.emit(v_mul_f32_e32(v[9 + j45 * 4], s[41], v[9 + j45 * 4])) + k.emit(v_accvgpr_read(v[10 + j45 * 4], v[202 + j45 * 16 + i * 32])) + k.emit(v_mul_f32_e32(v[10 + j45 * 4], s[41], v[10 + j45 * 4])) + k.emit(v_accvgpr_read(v[11 + j45 * 4], v[203 + j45 * 16 + i * 32])) + k.emit(v_mul_f32_e32(v[11 + j45 * 4], s[41], v[11 + j45 * 4])) + for j46 in range(4): + k.emit(v_cvt_pk_bf16_f32(v[16 + j46 * 1], v[8 + j46 * 2], v[9 + j46 * 2])) + k.emit(s_nop(1)) + k.emit(v_permlane16_swap_b32_e32(v[16], v[18])) + k.emit(s_nop(1)) + k.emit(v_permlane16_swap_b32_e32(v[17], v[19])) + k.emit(s_nop(1)) + k.emit(buffer_store_dwordx4(v[16:19], v[244], s[4:7], 0, 0, 1)) + k.emit(v_add_i32(v[244], v[244], 64)) + k.emit(s_mul_i32(s[62], s[36], 48)) + k.emit(v_add_u32_e32(v[244], s[62], v[243])) + for i in range(2): + for j47 in range(2): + k.emit(v_accvgpr_read(v[8 + j47 * 4], v[204 + j47 * 16 + i * 32])) + k.emit(v_mul_f32_e32(v[8 + j47 * 4], s[41], v[8 + j47 * 4])) + k.emit(v_accvgpr_read(v[9 + j47 * 4], v[205 + j47 * 16 + i * 32])) + k.emit(v_mul_f32_e32(v[9 + j47 * 4], s[41], v[9 + j47 * 4])) + k.emit(v_accvgpr_read(v[10 + j47 * 4], v[206 + j47 * 16 + i * 32])) + k.emit(v_mul_f32_e32(v[10 + j47 * 4], s[41], v[10 + j47 * 4])) + k.emit(v_accvgpr_read(v[11 + j47 * 4], v[207 + j47 * 16 + i * 32])) + k.emit(v_mul_f32_e32(v[11 + j47 * 4], s[41], v[11 + j47 * 4])) + for j48 in range(4): + k.emit(v_cvt_pk_bf16_f32(v[16 + j48 * 1], v[8 + j48 * 2], v[9 + j48 * 2])) + k.emit(s_nop(1)) + k.emit(v_permlane16_swap_b32_e32(v[16], v[18])) + k.emit(s_nop(1)) + k.emit(v_permlane16_swap_b32_e32(v[17], v[19])) + k.emit(s_nop(1)) + k.emit(buffer_store_dwordx4(v[16:19], v[244], s[4:7], 0, 0, 1)) + k.emit(v_add_i32(v[244], v[244], 64)) + k.label('L0_6E70') + k.emit(s_waitcnt()) + k.emit(s_endpgm()) + elif (tile_m, tile_n) == (192, 256): + k.emit(s_and_b32(s[1], s[1], LIT, 65535)) + k.emit(s_load_dwordx2(s[4:5], s[0:1], s[0], 0, 0, 0, 0, 1)) + k.emit(s_mov_b32(s[8], 0)) + k.emit(s_mov_b32(s[9], 0)) + k.emit(s_load_dwordx2(s[12:13], s[0:1], s[0], 8, 0, 0, 0, 1)) + k.emit(s_load_dwordx2(s[16:17], s[0:1], s[0], 16, 0, 0, 0, 1)) + k.emit(s_mov_b32(s[41], 1.0)) + k.emit(s_mov_b32(s[42], 0)) + k.emit(s_mov_b32(s[36], N)) + k.emit(s_mov_b32(s[37], K)) + k.emit(s_mov_b32(s[38], K)) + k.emit(s_mov_b32(s[43], M)) + k.emit(s_mov_b32(s[44], N)) + k.emit(s_mov_b32(s[45], K)) + k.emit(s_load_dwordx2(s[20:21], s[0:1], s[0], 24, 0, 0, 0, 1)) + k.emit(s_load_dwordx2(s[24:25], s[0:1], s[0], 32, 0, 0, 0, 1)) + k.emit(s_mov_b32(s[39], scale_k)) + k.emit(s_mov_b32(s[40], scale_k)) + k.emit(v_lshrrev_b32_e32(v[1], 10)) + k.emit(v_lshrrev_b32_e32(v[2], 10, v[1])) + k.emit(v_and_b32_e32(v[2], LIT, v[2], 1023)) + k.emit(v_and_b32_e32(v[1], LIT, v[1], 1023)) + k.emit(v_and_b32_e32(v[0], LIT, v[0], 1023)) + k.emit(v_lshrrev_b32_e32(v[3], 6)) + k.emit(v_and_b32_e32(v[0], 63)) + k.emit(v_readfirstlane_b32_e32(v[46], v[3])) + k.emit(s_waitcnt(49279)) + k.emit(s_mul_i32(s[63], LIT, 8, 192)) + k.emit(v_cvt_f32_u32_e32(v[4], s[63])) + k.emit(s_sub_i32(s[62], 0, s[63])) + k.emit(v_rcp_iflag_f32_e32(v[4], v[4])) + k.emit(s_nop()) + k.emit(v_mul_f32_e32(v[4], LIT, v[4], 1333788670)) + k.emit(v_cvt_u32_f32_e32(v[4], v[4])) + k.emit(v_mul_lo_u32(v[5], s[62], v[4])) + k.emit(v_mul_hi_u32(v[5], v[4], v[5])) + k.emit(v_add_u32_e32(v[4], v[4], v[5])) + k.emit(v_mul_hi_u32(v[4], s[43], v[4])) + k.emit(v_mul_lo_u32(v[5], v[4], s[63])) + k.emit(v_sub_u32_e32(v[7], s[43], v[5])) + k.emit(v_add_u32_e32(v[6], 1, v[4])) + k.emit(v_cmp_le_u32_e32(s[63], v[7])) + k.emit(v_subrev_u32_e32(v[5], s[63], v[7])) + k.emit(s_nop()) + k.emit(v_cndmask_b32_e32(v[4], v[4], v[6])) + k.emit(v_cndmask_b32_e32(v[7], v[7], v[5])) + k.emit(v_add_u32_e32(v[5], 1, v[4])) + k.emit(v_cmp_le_u32_e32(s[63], v[7])) + k.emit(s_nop(1)) + k.emit(v_cndmask_b32_e32(v[7], v[4], v[5])) + k.emit(s_nop(3)) + k.emit(v_readfirstlane_b32_e32(v[62], v[7])) + k.emit(s_nop(3)) + k.emit(s_lshl_b32(s[62], s[62], 3)) + k.emit(s_cmp_lt_i32(s[3], s[62])) + k.emit(s_cbranch_scc0(33), target='L1_01C0') + k.emit(s_add_u32(s[49], s[44], LIT, 255)) + k.emit(s_lshr_b32(s[48], s[49], 8)) + k.emit(s_mul_i32(s[49], s[48], s[3])) + k.emit(s_add_i32(s[49], s[49], s[2])) + k.emit(s_lshr_b32(s[63], s[44], 13)) + k.emit(s_lshl_b32(s[47], s[63], 5)) + k.emit(s_mul_i32(s[62], s[62], s[47])) + k.emit(s_cmp_lt_i32(s[49], s[62])) + k.emit(s_cbranch_scc0(13), target='L1_0198') + k.emit(s_and_b32(s[62], s[49], LIT, 255)) + k.emit(s_and_b32(s[47], s[62], 31)) + k.emit(s_lshr_b32(s[48], s[62], 5)) + k.emit(s_lshr_b32(s[49], s[49], 8)) + k.label('L1_0178') + k.emit(s_cmp_lt_i32(s[49], s[63])) + k.emit(s_cbranch_scc1(3), target='L1_018C') + k.emit(s_sub_i32(s[49], s[49], s[63])) + k.emit(s_add_i32(s[48], s[48], 8)) + k.emit(s_branch(65531), target='L1_0178') + k.label('L1_018C') + k.emit(s_mul_i32(s[49], s[49], 32)) + k.emit(s_add_i32(s[47], s[47], s[49])) + k.emit(s_branch(12), target='L1_01C8') + k.label('L1_0198') + k.emit(s_sub_i32(s[49], s[49], s[62])) + k.emit(s_sub_i32(s[63], s[48], s[47])) + k.emit(s_mov_b32(s[48], 0)) + k.label('L1_01A4') + k.emit(s_cmp_lt_i32(s[49], s[63])) + k.emit(s_cbranch_scc1(3), target='L1_01B8') + k.emit(s_sub_i32(s[49], s[49], s[63])) + k.emit(s_add_i32(s[48], s[48], 1)) + k.emit(s_branch(65531), target='L1_01A4') + k.label('L1_01B8') + k.emit(s_add_i32(s[47], s[47], s[49])) + k.emit(s_branch(2), target='L1_01C8') + k.label('L1_01C0') + k.emit(s_mov_b32(s[47], s[2])) + k.emit(s_mov_b32(s[48], s[3])) + k.label('L1_01C8') + k.emit(s_lshr_b32(s[37], s[37], 1)) + k.emit(s_mul_i32(s[62], s[48], LIT, 192)) + k.emit(s_mul_hi_u32(s[63], s[37], s[62])) + k.emit(s_add_u32(s[13], s[13], s[63])) + k.emit(s_mul_i32(s[63], s[37], s[62])) + k.emit(s_add_u32(s[12], s[12], s[63])) + k.emit(s_addc_u32(s[13], s[13], 0)) + k.emit(s_sub_i32(s[63], s[43], s[62])) + k.emit(s_cmp_lt_u32(s[63], LIT, 192)) + k.emit(s_cselect_b32(s[62], s[63], LIT, 192)) + k.emit(s_mul_i32(s[14], s[37], s[62])) + k.emit(s_mov_b32(s[15], LIT, 131072)) + k.emit(v_lshrrev_b32_e32(v[4], 3)) + k.emit(v_lshrrev_b32_e32(v[5], 2, v[4])) + k.emit(v_lshlrev_b32_e32(v[5], 4, v[5])) + k.emit(v_and_b32_e32(v[4], 3, v[4])) + k.emit(v_lshrrev_b32_e32(v[6], 1, v[4])) + k.emit(v_lshlrev_b32_e32(v[6], 2, v[6])) + k.emit(v_add_u32_e32(v[5], v[5], v[6])) + k.emit(v_and_b32_e32(v[4], 1, v[4])) + k.emit(v_add_u32_e32(v[5], v[5], v[4])) + k.emit(v_mul_lo_u32(v[178], s[37], v[5])) + k.emit(v_and_b32_e32(v[4], 7)) + k.emit(v_lshlrev_b32_e32(v[4], 4, v[4])) + k.emit(v_add_u32_e32(v[178], v[4], v[178])) + k.emit(s_lshr_b32(s[62], s[46], 1)) + k.emit(s_mul_i32(s[62], s[62], 8)) + k.emit(s_and_b32(s[63], s[46], 1)) + k.emit(s_mul_i32(s[63], s[63], 2)) + k.emit(s_add_u32(s[62], s[62], s[63])) + k.emit(s_mul_i32(s[62], s[37], s[62])) + k.emit(v_add_u32_e32(v[178], s[62], v[178])) + k.emit(s_mul_i32(s[62], s[37], 32)) + for i in range(5): + k.emit(v_add_u32_e32(v[179 + i * 1], s[62], v[178 + i * 1])) + k.emit(s_mul_i32(s[64], LIT, s[46], 1056)) + k.emit(s_add_u32(s[64], LIT, s[64], 4096)) + k.emit(v_and_b32_e32(v[4], 15)) + k.emit(v_lshrrev_b32_e32(v[5], 3, v[4])) + k.emit(v_mul_i32_i24_e32(v[5], 2, v[5])) + k.emit(v_and_b32_e32(v[4], 3)) + k.emit(v_lshrrev_b32_e32(v[6], 1, v[4])) + k.emit(v_add_u32_e32(v[4], v[5], v[6])) + k.emit(v_mul_i32_i24_e32(v[184], LIT, v[4], 1056)) + k.emit(v_and_b32_e32(v[4], 7)) + k.emit(v_lshrrev_b32_e32(v[5], 2, v[4])) + k.emit(v_mul_i32_i24_e32(v[5], LIT, v[5], 256)) + k.emit(v_add_u32_e32(v[184], v[5], v[184])) + k.emit(v_and_b32_e32(v[4], 1)) + k.emit(v_mul_i32_i24_e32(v[6], LIT, v[4], 128)) + k.emit(v_add_u32_e32(v[184], v[6], v[184])) + k.emit(v_lshrrev_b32_e32(v[4], 4)) + k.emit(v_mul_i32_i24_e32(v[4], 16, v[4])) + k.emit(v_add_u32_e32(v[184], v[4], v[184])) + k.emit(v_add_u32_e32(v[184], LIT, v[184], 4096)) + k.emit(v_add_u32_e32(v[185], LIT, v[184], 25344)) + k.emit(s_mul_i32(s[62], s[48], LIT, 192)) + k.emit(s_mul_hi_u32(s[63], s[39], s[62])) + k.emit(s_add_u32(s[21], s[21], s[63])) + k.emit(s_mul_i32(s[63], s[39], s[62])) + k.emit(s_add_u32(s[20], s[20], s[63])) + k.emit(s_addc_u32(s[21], s[21], 0)) + k.emit(s_add_u32(s[63], s[43], 31)) + k.emit(s_lshr_b32(s[63], s[63], 5)) + k.emit(s_lshl_b32(s[63], s[63], 5)) + k.emit(s_sub_i32(s[63], s[63], s[62])) + k.emit(s_cmp_lt_u32(s[63], LIT, 192)) + k.emit(s_cselect_b32(s[62], s[63], LIT, 192)) + k.emit(s_mul_i32(s[22], s[39], s[62])) + k.emit(s_mov_b32(s[23], LIT, 131072)) + k.emit(v_lshlrev_b32_e32(v[186], 2)) + k.emit(s_mul_i32(s[63], s[46], 32)) + k.emit(s_mul_i32(s[63], s[63], s[39])) + k.emit(v_add_u32_e32(v[186], s[63], v[186])) + k.emit(s_mul_i32(s[63], LIT, s[39], 128)) + k.emit(v_add_u32_e32(v[187], s[63], v[186])) + k.emit(s_mul_i32(s[65], s[46], LIT, 256)) + k.emit(s_add_i32(s[65], s[65], 0)) + k.emit(v_lshlrev_b32_e32(v[188], 2)) + k.emit(v_add_u32_e32(v[188], 0, v[188])) + k.emit(s_lshr_b32(s[38], s[38], 1)) + k.emit(s_mul_i32(s[62], s[47], LIT, 256)) + k.emit(s_mul_hi_u32(s[63], s[38], s[62])) + k.emit(s_add_u32(s[17], s[17], s[63])) + k.emit(s_mul_i32(s[63], s[38], s[62])) + k.emit(s_add_u32(s[16], s[16], s[63])) + k.emit(s_addc_u32(s[17], s[17], 0)) + k.emit(s_sub_i32(s[63], s[44], s[62])) + k.emit(s_cmp_lt_u32(s[63], LIT, 256)) + k.emit(s_cselect_b32(s[62], s[63], LIT, 256)) + k.emit(s_mul_i32(s[18], s[38], s[62])) + k.emit(s_mov_b32(s[19], LIT, 131072)) + k.emit(v_lshlrev_b32_e32(v[189], 4)) + k.emit(s_mul_i32(s[63], s[46], 64)) + k.emit(s_mul_i32(s[62], s[63], s[38])) + k.emit(v_add_u32_e32(v[189], s[62], v[189])) + k.emit(s_mul_i32(s[62], 16, s[38])) + k.emit(v_add_u32_e32(v[190], s[62], v[189])) + k.emit(v_add_u32_e32(v[191], s[62], v[190])) + k.emit(v_add_u32_e32(v[192], s[62], v[191])) + k.emit(s_mul_i32(s[62], s[47], LIT, 256)) + k.emit(s_mul_hi_u32(s[63], s[40], s[62])) + k.emit(s_add_u32(s[25], s[25], s[63])) + k.emit(s_mul_i32(s[63], s[40], s[62])) + k.emit(s_add_u32(s[24], s[24], s[63])) + k.emit(s_addc_u32(s[25], s[25], 0)) + k.emit(s_sub_i32(s[63], s[44], s[62])) + k.emit(s_cmp_lt_u32(s[63], LIT, 256)) + k.emit(s_cselect_b32(s[62], s[63], LIT, 256)) + k.emit(s_mul_i32(s[26], s[40], s[62])) + k.emit(s_mov_b32(s[27], LIT, 131072)) + k.emit(v_lshlrev_b32_e32(v[193], 2)) + k.emit(s_mul_i32(s[63], s[46], 64)) + k.emit(s_mul_i32(s[63], s[63], s[40])) + k.emit(v_add_u32_e32(v[193], s[63], v[193])) + k.emit(s_mul_i32(s[62], 32, s[40])) + k.emit(v_add_u32_e32(v[194], s[62], v[193])) + k.emit(s_mov_b32(s[66], LIT, 128)) + k.emit(s_mov_b32(s[67], LIT, 2048)) + k.emit(s_mov_b32(s[68], LIT, 256)) + k.emit(s_mov_b32(s[69], LIT, 256)) + k.emit(s_mov_b32(s[60], 0)) + k.emit(s_mov_b32(s[61], s[45])) + k.emit(s_add_u32(NULL, 0, s[65])) + k.emit(buffer_load_dword(v[0], v[186], s[20:23], 0, 0, 1, 0, 0, 0, 0, 1)) + for i in range(8): + k.emit(v_accvgpr_write(v[0 + i * 1], 0)) + k.emit(s_add_u32(NULL, LIT, s[65], 1024)) + k.emit(buffer_load_dword(v[0], v[187], s[20:23], 0, 0, 1, 0, 0, 0, 0, 1)) + for i in range(8): + k.emit(v_accvgpr_write(v[8 + i * 1], 0)) + k.emit(s_add_u32(NULL, 0, s[64])) + for i in range(5): + k.emit(buffer_load_dwordx4(v[0:3], v[178 + i * 1], s[12:15], 0, 0, 1, 0, 0, 0, 0, 1)) + for j49 in range(8): + k.emit(v_accvgpr_write(v[16 + j49 * 1 + i * 8], 0)) + k.emit(s_add_u32(NULL, LIT, s[64], 4224 + i * 4224)) + k.emit(buffer_load_dwordx4(v[0:3], v[183], s[12:15], 0, 0, 1, 0, 0, 0, 0, 1)) + for i in range(8): + k.emit(v_accvgpr_write(v[56 + i * 1], 0)) + k.emit(s_add_u32(s[62], LIT, s[60], 256)) + k.emit(s_cmp_lt_u32(s[62], s[61])) + k.emit(s_cselect_b32(s[66], s[66], 0)) + k.emit(s_cselect_b32(s[68], s[68], 0)) + for i in range(2): + k.emit(s_add_u32(s[12 + i * 8], s[12 + i * 8], s[66 + i * 2])) + k.emit(s_addc_u32(s[13 + i * 8], 0, s[13 + i * 8])) + k.emit(s_sub_u32(s[14 + i * 8], s[14 + i * 8], s[66 + i * 2])) + for i in range(2): + for j50 in range(2): + k.emit(buffer_load_dwordx4(v[104 + j50 * 8 + i * 16:107 + j50 * 8 + i * 16], v[189 + i * 2], s[16:19], 0, 0 + j50 * 1024, 1)) + k.emit(v_accvgpr_write(v[64 + j50 * 16 + i * 40], 0)) + k.emit(v_accvgpr_write(v[65 + j50 * 16 + i * 40], 0)) + k.emit(v_accvgpr_write(v[66 + j50 * 16 + i * 40], 0)) + k.emit(v_accvgpr_write(v[67 + j50 * 16 + i * 40], 0)) + k.emit(v_accvgpr_write(v[68 + j50 * 16 + i * 40], 0)) + k.emit(v_accvgpr_write(v[69 + j50 * 16 + i * 40], 0)) + k.emit(v_accvgpr_write(v[70 + j50 * 16 + i * 40], 0)) + k.emit(v_accvgpr_write(v[71 + j50 * 16 + i * 40], 0)) + k.emit(buffer_load_dwordx4(v[108 + j50 * 8 + i * 16:111 + j50 * 8 + i * 16], v[190 + i * 2], s[16:19], 0, 0 + j50 * 1024, 1)) + k.emit(v_accvgpr_write(v[72 + j50 * 16 + i * 40], 0)) + k.emit(v_accvgpr_write(v[73 + j50 * 16 + i * 40], 0)) + k.emit(v_accvgpr_write(v[74 + j50 * 16 + i * 40], 0)) + k.emit(v_accvgpr_write(v[75 + j50 * 16 + i * 40], 0)) + k.emit(v_accvgpr_write(v[76 + j50 * 16 + i * 40], 0)) + k.emit(v_accvgpr_write(v[77 + j50 * 16 + i * 40], 0)) + k.emit(v_accvgpr_write(v[78 + j50 * 16 + i * 40], 0)) + k.emit(v_accvgpr_write(v[79 + j50 * 16 + i * 40], 0)) + k.emit(buffer_load_dword(v[174 + i * 1], v[193 + i * 1], s[24:27], 0, 0, 1)) + for j51 in range(8): + k.emit(v_accvgpr_write(v[96 + j51 * 1 + i * 40], 0)) + k.emit(s_add_u32(s[63], LIT, s[60], 256)) + k.emit(s_cmp_lt_u32(s[63], s[61])) + k.emit(s_cselect_b32(s[67], s[67], 0)) + k.emit(s_cselect_b32(s[69], s[69], 0)) + for i in range(2): + k.emit(s_add_u32(s[16 + i * 8], s[16 + i * 8], s[67 + i * 2])) + k.emit(s_addc_u32(s[17 + i * 8], 0, s[17 + i * 8])) + k.emit(s_sub_u32(s[18 + i * 8], s[18 + i * 8], s[67 + i * 2])) + for i in range(2): + k.emit(s_add_u32(NULL, LIT, s[65], 2048 + i * 1024)) + k.emit(buffer_load_dword(v[0], v[186 + i * 1], s[20:23], 0, 0, 1, 0, 0, 0, 0, 1)) + for j52 in range(8): + k.emit(v_accvgpr_write(v[144 + j52 * 1 + i * 8], 0)) + for i in range(4): + k.emit(s_add_u32(NULL, LIT, s[64], 25344 + i * 4224)) + k.emit(buffer_load_dwordx4(v[0:3], v[178 + i * 1], s[12:15], 0, 0, 1, 0, 0, 0, 0, 1)) + for j53 in range(8): + k.emit(v_accvgpr_write(v[160 + j53 * 1 + i * 8], 0)) + k.emit(s_add_u32(NULL, LIT, s[64], 42240)) + k.emit(buffer_load_dwordx4(v[0:3], v[182], s[12:15], 0, 0, 1, 0, 0, 0, 0, 1)) + k.emit(s_add_u32(NULL, LIT, s[64], 46464)) + k.emit(buffer_load_dwordx4(v[0:3], v[183], s[12:15], 0, 0, 1, 0, 0, 0, 0, 1)) + k.emit(s_add_u32(s[62], LIT, s[60], 512)) + k.emit(s_cmp_lt_u32(s[62], s[61])) + k.emit(s_cselect_b32(s[66], s[66], 0)) + k.emit(s_cselect_b32(s[68], s[68], 0)) + for i in range(2): + k.emit(s_add_u32(s[12 + i * 8], s[12 + i * 8], s[66 + i * 2])) + k.emit(s_addc_u32(s[13 + i * 8], 0, s[13 + i * 8])) + k.emit(s_sub_u32(s[14 + i * 8], s[14 + i * 8], s[66 + i * 2])) + k.emit(buffer_load_dwordx4(v[136:139], v[189], s[16:19], 0, 0, 1)) + k.emit(buffer_load_dwordx4(v[140:143], v[190], s[16:19], 0, 0, 1)) + k.emit(buffer_load_dwordx4(v[144:147], v[189], s[16:19], 0, 1024, 1)) + k.emit(buffer_load_dwordx4(v[148:151], v[190], s[16:19], 0, 1024, 1)) + k.emit(buffer_load_dword(v[176], v[193], s[24:27], 0, 0, 1)) + k.emit(s_waitcnt(20347)) + k.emit(s_barrier()) + k.emit(ds_read_b128(v[8:11], v[184])) + k.emit(ds_read_b128(v[16:19], v[184], v[0], v[0], 0, 64)) + k.emit(ds_read_b128(v[12:15], v[184], v[0], v[0], 0, 0, 2)) + k.emit(ds_read_b128(v[20:23], v[184], v[0], v[0], 0, 64, 2)) + k.emit(ds_read_b32(v[168], v[188])) + k.emit(ds_read_b128(v[24:27], v[184], v[0], v[0], 0, 128, 16)) + k.emit(ds_read_b128(v[32:35], v[184], v[0], v[0], 0, 192, 16)) + k.emit(ds_read_b128(v[28:31], v[184], v[0], v[0], 0, 128, 18)) + k.emit(ds_read_b128(v[36:39], v[184], v[0], v[0], 0, 192, 18)) + k.emit(ds_read_b32(v[169], v[188], v[0], v[0], 0, 0, 1)) + for i in range(5): + k.emit(s_nop()) + k.emit(s_lshl_b32(s[36], s[36], 1)) + k.emit(s_and_b32(s[5], s[5], LIT, 65535)) + k.emit(s_or_b32(s[5], s[5], LIT, 262144)) + k.emit(s_mul_i32(s[62], s[48], LIT, 192)) + k.emit(s_mul_hi_u32(s[63], s[36], s[62])) + k.emit(s_add_u32(s[5], s[5], s[63])) + k.emit(s_mul_i32(s[63], s[36], s[62])) + k.emit(s_add_u32(s[4], s[4], s[63])) + k.emit(s_addc_u32(s[5], s[5], 0)) + k.emit(s_mul_i32(s[63], s[47], LIT, 256)) + k.emit(s_lshl_b32(s[63], s[63], 1)) + k.emit(s_add_u32(s[4], s[4], s[63])) + k.emit(s_addc_u32(s[5], s[5], 0)) + k.emit(s_sub_i32(s[62], s[43], s[62])) + k.emit(s_cmp_lt_u32(s[62], LIT, 192)) + k.emit(s_cselect_b32(s[62], s[62], LIT, 192)) + k.emit(s_mul_i32(s[62], s[36], s[62])) + k.emit(s_sub_i32(s[62], s[62], s[63])) + k.emit(s_mov_b32(s[6], s[62])) + k.emit(s_mov_b32(s[7], LIT, 131072)) + k.emit(v_lshrrev_b32_e32(v[4], 3)) + k.emit(v_mul_lo_u32(v[195], v[4], s[36])) + k.emit(s_mul_i32(s[62], s[46], 64)) + k.emit(s_lshl_b32(s[62], s[62], 1)) + k.emit(v_and_b32_e32(v[4], 7)) + k.emit(v_mul_i32_i24_e32(v[4], 16, v[4])) + k.emit(v_add_u32_e32(v[4], s[62], v[4])) + k.emit(v_add_u32_e32(v[195], v[195], v[4])) + k.emit(s_mul_i32(s[62], s[36], 8)) + for i in range(23): + k.emit(v_add_u32_e32(v[196 + i * 1], s[62], v[195 + i * 1])) + k.emit(s_cmp_lt_i32(s[46], 2)) + k.emit(s_cbranch_scc0(1063), target='L1_1E2C') + k.label('L1_0D90') + k.emit(s_waitcnt(20338)) + k.emit(s_barrier()) + for i in range(2): + k.emit(s_waitcnt(50559)) + for j54 in range(2): + k.emit(v_mfma_fp4(v[0 + i * 32:3 + i * 32], v[104 + j54 * 8:107 + j54 * 8], v[8 + j54 * 8 + i * 16:11 + j54 * 8 + i * 16], 0, 0 + j54 * 3, v[174], v[168 + i * 1])) + k.emit(ds_read_b128(v[40 + j54 * 4 + i * 16:43 + j54 * 4 + i * 16], v[184], v[0], v[0], 0, 0 + i * 128, 33 + j54 * 2 + i * 16)) + k.emit(v_mfma_fp4(v[4 + i * 32:7 + i * 32], v[104 + j54 * 8:107 + j54 * 8], v[12 + j54 * 8 + i * 16:15 + j54 * 8 + i * 16], 2, 0 + j54 * 3, v[174], v[168 + i * 1])) + k.emit(buffer_load_dwordx4(v[152 + j54 * 4 + i * 8:155 + j54 * 4 + i * 8], v[191 + j54 * 1], s[16:19], 0, 0 + i * 1024, 1)) + k.emit(v_mfma_fp4(v[8 + i * 32:11 + i * 32], v[108 + j54 * 8:111 + j54 * 8], v[8 + j54 * 8 + i * 16:11 + j54 * 8 + i * 16], 1, 0 + j54 * 3, v[174], v[168 + i * 1])) + k.emit(ds_read_b128(v[48 + j54 * 4 + i * 16:51 + j54 * 4 + i * 16], v[184], v[0], v[0], 0, 64 + i * 128, 33 + j54 * 2 + i * 16)) + k.emit(v_mfma_fp4(v[12 + i * 32:15 + i * 32], v[108 + j54 * 8:111 + j54 * 8], v[12 + j54 * 8 + i * 16:15 + j54 * 8 + i * 16], 3, 0 + j54 * 3, v[174], v[168 + i * 1])) + k.emit(ds_read_b32(v[170 + i * 1], v[188], v[0], v[0], 0, 0, 2 + i * 1)) + k.emit(s_waitcnt(50559)) + k.emit(v_mfma_fp4(v[64:67], v[104:107], v[40:43], 0, 0, v[174], v[170])) + k.emit(ds_read_b128(v[72:75], v[184], v[0], v[0], 0, 0, 66)) + k.emit(v_mfma_fp4(v[68:71], v[104:107], v[44:47], 2, 0, v[174], v[170])) + k.emit(buffer_load_dword(v[177], v[194], s[24:27], 0, 0, 1)) + k.emit(v_mfma_fp4(v[72:75], v[108:111], v[40:43], 1, 0, v[174], v[170])) + k.emit(s_add_u32(s[63], LIT, s[60], 512)) + k.emit(ds_read_b128(v[80:83], v[184], v[0], v[0], 0, 64, 66)) + k.emit(v_mfma_fp4(v[76:79], v[108:111], v[44:47], 3, 0, v[174], v[170])) + k.emit(s_cmp_lt_u32(s[63], s[61])) + k.emit(v_mfma_fp4(v[64:67], v[112:115], v[48:51], 0, 3, v[174], v[170])) + k.emit(s_cselect_b32(s[67], s[67], 0)) + k.emit(ds_read_b128(v[76:79], v[184], v[0], v[0], 0, 0, 68)) + k.emit(v_mfma_fp4(v[68:71], v[112:115], v[52:55], 2, 3, v[174], v[170])) + k.emit(s_cselect_b32(s[69], s[69], 0)) + k.emit(v_mfma_fp4(v[72:75], v[116:119], v[48:51], 1, 3, v[174], v[170])) + k.emit(s_add_u32(s[16], s[16], s[67])) + k.emit(ds_read_b128(v[84:87], v[184], v[0], v[0], 0, 64, 68)) + k.emit(v_mfma_fp4(v[76:79], v[116:119], v[52:55], 3, 3, v[174], v[170])) + k.emit(s_addc_u32(s[17], 0, s[17])) + k.emit(s_sub_u32(s[18], s[18], s[67])) + k.emit(ds_read_b32(v[172], v[188], v[0], v[0], 0, 0, 4)) + k.emit(s_waitcnt(50559)) + k.emit(v_mfma_fp4(v[96:99], v[104:107], v[56:59], 0, 0, v[174], v[171])) + k.emit(s_add_u32(s[24], s[24], s[69])) + k.emit(ds_read_b128(v[88:91], v[184], v[0], v[0], 0, 128, 82)) + k.emit(v_mfma_fp4(v[100:103], v[104:107], v[60:63], 2, 0, v[174], v[171])) + k.emit(s_addc_u32(s[25], 0, s[25])) + k.emit(v_mfma_fp4(v[104:107], v[108:111], v[56:59], 1, 0, v[174], v[171])) + k.emit(s_sub_u32(s[26], s[26], s[69])) + for i in range(2): + k.emit(ds_read_b128(v[96 + i * -4:99 + i * -4], v[184], v[0], v[0], 0, 192 + i * -64, 82 + i * 2)) + k.emit(v_mfma_fp4(v[108 + i * -8:111 + i * -8], v[108 + i * 4:111 + i * 4], v[60 + i * 8:63 + i * 8], 3 + i * -1, 0 + i * 3, v[174], v[171])) + k.emit(v_mfma_fp4(v[96 + i * 8:99 + i * 8], v[112 + i * 4:115 + i * 4], v[64:67], 0 + i * 1, 3, v[174], v[171])) + k.emit(ds_read_b128(v[100:103], v[184], v[0], v[0], 0, 192, 84)) + k.emit(v_mfma_fp4(v[108:111], v[116:119], v[68:71], 3, 3, v[174], v[171])) + k.emit(ds_read_b32(v[173], v[188], v[0], v[0], 0, 0, 5)) + k.emit(s_barrier()) + k.emit(s_waitcnt(50559)) + k.emit(v_mfma_fp4(v[128:131], v[104:107], v[72:75], 0, 0, v[174], v[172])) + k.emit(v_mfma_fp4(v[132:135], v[104:107], v[76:79], 2, 0, v[174], v[172])) + k.emit(s_add_u32(NULL, 0, s[65])) + k.emit(buffer_load_dword(v[0], v[186], s[20:23], 0, 0, 1, 0, 0, 0, 0, 1)) + k.emit(v_mfma_fp4(v[136:139], v[108:111], v[72:75], 1, 0, v[174], v[172])) + k.emit(v_mfma_fp4(v[140:143], v[108:111], v[76:79], 3, 0, v[174], v[172])) + k.emit(v_mfma_fp4(v[128:131], v[112:115], v[80:83], 0, 3, v[174], v[172])) + k.emit(v_mfma_fp4(v[132:135], v[112:115], v[84:87], 2, 3, v[174], v[172])) + k.emit(s_add_u32(NULL, LIT, s[65], 1024)) + k.emit(buffer_load_dword(v[0], v[187], s[20:23], 0, 0, 1, 0, 0, 0, 0, 1)) + k.emit(v_mfma_fp4(v[136:139], v[116:119], v[80:83], 1, 3, v[174], v[172])) + k.emit(v_mfma_fp4(v[140:143], v[116:119], v[84:87], 3, 3, v[174], v[172])) + k.emit(s_waitcnt(49279)) + k.emit(v_mfma_fp4(v[160:163], v[104:107], v[88:91], 0, 0, v[174], v[173])) + k.emit(v_mfma_fp4(v[164:167], v[104:107], v[92:95], 2, 0, v[174], v[173])) + k.emit(s_add_u32(NULL, 0, s[64])) + k.emit(buffer_load_dwordx4(v[0:3], v[178], s[12:15], 0, 0, 1, 0, 0, 0, 0, 1)) + k.emit(v_mfma_fp4(v[168:171], v[108:111], v[88:91], 1, 0, v[174], v[173])) + k.emit(v_mfma_fp4(v[172:175], v[108:111], v[92:95], 3, 0, v[174], v[173])) + k.emit(v_mfma_fp4(v[160:163], v[112:115], v[96:99], 0, 3, v[174], v[173])) + k.emit(v_mfma_fp4(v[164:167], v[112:115], v[100:103], 2, 3, v[174], v[173])) + k.emit(s_add_u32(NULL, LIT, s[64], 4224)) + k.emit(buffer_load_dwordx4(v[0:3], v[179], s[12:15], 0, 0, 1, 0, 0, 0, 0, 1)) + k.emit(v_mfma_fp4(v[168:171], v[116:119], v[96:99], 1, 3, v[174], v[173])) + k.emit(v_mfma_fp4(v[172:175], v[116:119], v[100:103], 3, 3, v[174], v[173])) + k.emit(s_waitcnt(20338)) + k.emit(s_barrier()) + k.emit(v_mfma_fp4(v[16:19], v[120:123], v[8:11], 0, 0, v[175], v[168])) + k.emit(v_mfma_fp4(v[20:23], v[120:123], v[12:15], 2, 0, v[175], v[168])) + k.emit(s_add_u32(NULL, LIT, s[64], 8448)) + k.emit(buffer_load_dwordx4(v[0:3], v[180], s[12:15], 0, 0, 1, 0, 0, 0, 0, 1)) + k.emit(v_mfma_fp4(v[24:27], v[124:127], v[8:11], 1, 0, v[175], v[168])) + k.emit(v_mfma_fp4(v[28:31], v[124:127], v[12:15], 3, 0, v[175], v[168])) + for i in range(2): + k.emit(v_mfma_fp4(v[16 + i * 32:19 + i * 32], v[128 + i * -8:131 + i * -8], v[16 + i * 8:19 + i * 8], 0, 3 + i * -3, v[175], v[168 + i * 1])) + k.emit(v_mfma_fp4(v[20 + i * 32:23 + i * 32], v[128 + i * -8:131 + i * -8], v[20 + i * 8:23 + i * 8], 2, 3 + i * -3, v[175], v[168 + i * 1])) + k.emit(s_add_u32(NULL, LIT, s[64], 12672 + i * 4224)) + k.emit(buffer_load_dwordx4(v[0:3], v[181 + i * 1], s[12:15], 0, 0, 1, 0, 0, 0, 0, 1)) + k.emit(v_mfma_fp4(v[24 + i * 32:27 + i * 32], v[132 + i * -8:135 + i * -8], v[16 + i * 8:19 + i * 8], 1, 3 + i * -3, v[175], v[168 + i * 1])) + k.emit(v_mfma_fp4(v[28 + i * 32:31 + i * 32], v[132 + i * -8:135 + i * -8], v[20 + i * 8:23 + i * 8], 3, 3 + i * -3, v[175], v[168 + i * 1])) + k.emit(v_mfma_fp4(v[48:51], v[128:131], v[32:35], 0, 3, v[175], v[169])) + k.emit(v_mfma_fp4(v[52:55], v[128:131], v[36:39], 2, 3, v[175], v[169])) + k.emit(s_add_u32(NULL, LIT, s[64], 21120)) + k.emit(buffer_load_dwordx4(v[0:3], v[183], s[12:15], 0, 0, 1, 0, 0, 0, 0, 1)) + k.emit(v_mfma_fp4(v[56:59], v[132:135], v[32:35], 1, 3, v[175], v[169])) + k.emit(s_add_u32(s[62], LIT, s[60], 768)) + k.emit(v_mfma_fp4(v[60:63], v[132:135], v[36:39], 3, 3, v[175], v[169])) + k.emit(s_cmp_lt_u32(s[62], s[61])) + k.emit(v_mfma_fp4(v[80:83], v[120:123], v[40:43], 0, 0, v[175], v[170])) + k.emit(s_cselect_b32(s[66], s[66], 0)) + k.emit(v_mfma_fp4(v[84:87], v[120:123], v[44:47], 2, 0, v[175], v[170])) + k.emit(s_cselect_b32(s[68], s[68], 0)) + k.emit(buffer_load_dwordx4(v[104:107], v[189], s[16:19], 0, 0, 1)) + k.emit(v_mfma_fp4(v[88:91], v[124:127], v[40:43], 1, 0, v[175], v[170])) + k.emit(s_add_u32(s[12], s[12], s[66])) + k.emit(v_mfma_fp4(v[92:95], v[124:127], v[44:47], 3, 0, v[175], v[170])) + k.emit(s_addc_u32(s[13], 0, s[13])) + k.emit(v_mfma_fp4(v[80:83], v[128:131], v[48:51], 0, 3, v[175], v[170])) + k.emit(s_sub_u32(s[14], s[14], s[66])) + k.emit(v_mfma_fp4(v[84:87], v[128:131], v[52:55], 2, 3, v[175], v[170])) + k.emit(s_add_u32(s[20], s[20], s[68])) + k.emit(buffer_load_dwordx4(v[108:111], v[190], s[16:19], 0, 0, 1)) + k.emit(v_mfma_fp4(v[88:91], v[132:135], v[48:51], 1, 3, v[175], v[170])) + k.emit(s_addc_u32(s[21], 0, s[21])) + k.emit(v_mfma_fp4(v[92:95], v[132:135], v[52:55], 3, 3, v[175], v[170])) + k.emit(s_sub_u32(s[22], s[22], s[68])) + k.emit(v_mfma_fp4(v[112:115], v[120:123], v[56:59], 0, 0, v[175], v[171])) + k.emit(s_addk_i32(s[60], 256)) + k.emit(v_mfma_fp4(v[116:119], v[120:123], v[60:63], 2, 0, v[175], v[171])) + k.emit(s_cmp_lt_i32(s[60], s[61])) + k.emit(buffer_load_dwordx4(v[112:115], v[189], s[16:19], 0, 1024, 1)) + k.emit(v_mfma_fp4(v[120:123], v[124:127], v[56:59], 1, 0, v[175], v[171])) + k.emit(v_mfma_fp4(v[124:127], v[124:127], v[60:63], 3, 0, v[175], v[171])) + k.emit(v_mfma_fp4(v[112:115], v[128:131], v[64:67], 0, 3, v[175], v[171])) + k.emit(v_mfma_fp4(v[116:119], v[128:131], v[68:71], 2, 3, v[175], v[171])) + k.emit(buffer_load_dwordx4(v[116:119], v[190], s[16:19], 0, 1024, 1)) + k.emit(v_mfma_fp4(v[120:123], v[132:135], v[64:67], 1, 3, v[175], v[171])) + k.emit(v_mfma_fp4(v[124:127], v[132:135], v[68:71], 3, 3, v[175], v[171])) + k.emit(v_mfma_fp4(v[144:147], v[120:123], v[72:75], 0, 0, v[175], v[172])) + k.emit(ds_read_b128(v[8:11], v[185])) + k.emit(v_mfma_fp4(v[148:151], v[120:123], v[76:79], 2, 0, v[175], v[172])) + k.emit(buffer_load_dword(v[174], v[193], s[24:27], 0, 0, 1)) + k.emit(v_mfma_fp4(v[152:155], v[124:127], v[72:75], 1, 0, v[175], v[172])) + k.emit(ds_read_b128(v[16:19], v[185], v[0], v[0], 0, 64)) + for i in range(2): + k.emit(v_mfma_fp4(v[156 + i * -8:159 + i * -8], v[124 + i * 4:127 + i * 4], v[76 + i * 8:79 + i * 8], 3 + i * -1, 0 + i * 3, v[175], v[172])) + k.emit(v_mfma_fp4(v[144 + i * 8:147 + i * 8], v[128 + i * 4:131 + i * 4], v[80:83], 0 + i * 1, 3, v[175], v[172])) + k.emit(ds_read_b128(v[12 + i * 8:15 + i * 8], v[185], v[0], v[0], 0, 0 + i * 64, 2)) + k.emit(v_mfma_fp4(v[156:159], v[132:135], v[84:87], 3, 3, v[175], v[172])) + k.emit(ds_read_b32(v[168], v[188], v[0], v[0], 0, 0, 8)) + for i in range(2): + for j55 in range(2): + k.emit(v_mfma_fp4(v[176 + j55 * 8:179 + j55 * 8], v[120 + j55 * 4 + i * 8:123 + j55 * 4 + i * 8], v[88 + i * 8:91 + i * 8], 0 + j55 * 1, 0 + i * 3, v[175], v[173])) + k.emit(ds_read_b128(v[24 + j55 * 8 + i * 4:27 + j55 * 8 + i * 4], v[185], v[0], v[0], 0, 128 + j55 * 64, 16 + i * 2)) + k.emit(v_mfma_fp4(v[180 + j55 * 8:183 + j55 * 8], v[120 + j55 * 4 + i * 8:123 + j55 * 4 + i * 8], v[92 + i * 8:95 + i * 8], 2 + j55 * 1, 0 + i * 3, v[175], v[173])) + k.emit(ds_read_b32(v[169], v[188], v[0], v[0], 0, 0, 9)) + k.emit(s_cbranch_scc0(1596), target='L1_2EC8') + k.emit(s_waitcnt(20338)) + k.emit(s_barrier()) + for i in range(2): + k.emit(s_waitcnt(50559)) + for j56 in range(2): + k.emit(v_mfma_fp4(v[0 + i * 32:3 + i * 32], v[136 + j56 * 8:139 + j56 * 8], v[8 + j56 * 8 + i * 16:11 + j56 * 8 + i * 16], 0, 0 + j56 * 3, v[176], v[168 + i * 1])) + k.emit(ds_read_b128(v[40 + j56 * 4 + i * 16:43 + j56 * 4 + i * 16], v[185], v[0], v[0], 0, 0 + i * 128, 33 + j56 * 2 + i * 16)) + k.emit(v_mfma_fp4(v[4 + i * 32:7 + i * 32], v[136 + j56 * 8:139 + j56 * 8], v[12 + j56 * 8 + i * 16:15 + j56 * 8 + i * 16], 2, 0 + j56 * 3, v[176], v[168 + i * 1])) + k.emit(buffer_load_dwordx4(v[120 + j56 * 4 + i * 8:123 + j56 * 4 + i * 8], v[191 + j56 * 1], s[16:19], 0, 0 + i * 1024, 1)) + k.emit(v_mfma_fp4(v[8 + i * 32:11 + i * 32], v[140 + j56 * 8:143 + j56 * 8], v[8 + j56 * 8 + i * 16:11 + j56 * 8 + i * 16], 1, 0 + j56 * 3, v[176], v[168 + i * 1])) + k.emit(ds_read_b128(v[48 + j56 * 4 + i * 16:51 + j56 * 4 + i * 16], v[185], v[0], v[0], 0, 64 + i * 128, 33 + j56 * 2 + i * 16)) + k.emit(v_mfma_fp4(v[12 + i * 32:15 + i * 32], v[140 + j56 * 8:143 + j56 * 8], v[12 + j56 * 8 + i * 16:15 + j56 * 8 + i * 16], 3, 0 + j56 * 3, v[176], v[168 + i * 1])) + k.emit(ds_read_b32(v[170 + i * 1], v[188], v[0], v[0], 0, 0, 10 + i * 1)) + k.emit(s_waitcnt(50559)) + k.emit(v_mfma_fp4(v[64:67], v[136:139], v[40:43], 0, 0, v[176], v[170])) + k.emit(ds_read_b128(v[72:75], v[185], v[0], v[0], 0, 0, 66)) + k.emit(v_mfma_fp4(v[68:71], v[136:139], v[44:47], 2, 0, v[176], v[170])) + k.emit(buffer_load_dword(v[175], v[194], s[24:27], 0, 0, 1)) + k.emit(v_mfma_fp4(v[72:75], v[140:143], v[40:43], 1, 0, v[176], v[170])) + k.emit(s_add_u32(s[63], LIT, s[60], 512)) + k.emit(ds_read_b128(v[80:83], v[185], v[0], v[0], 0, 64, 66)) + k.emit(v_mfma_fp4(v[76:79], v[140:143], v[44:47], 3, 0, v[176], v[170])) + k.emit(s_cmp_lt_u32(s[63], s[61])) + k.emit(v_mfma_fp4(v[64:67], v[144:147], v[48:51], 0, 3, v[176], v[170])) + k.emit(s_cselect_b32(s[67], s[67], 0)) + k.emit(ds_read_b128(v[76:79], v[185], v[0], v[0], 0, 0, 68)) + k.emit(v_mfma_fp4(v[68:71], v[144:147], v[52:55], 2, 3, v[176], v[170])) + k.emit(s_cselect_b32(s[69], s[69], 0)) + k.emit(v_mfma_fp4(v[72:75], v[148:151], v[48:51], 1, 3, v[176], v[170])) + k.emit(s_add_u32(s[16], s[16], s[67])) + k.emit(ds_read_b128(v[84:87], v[185], v[0], v[0], 0, 64, 68)) + k.emit(v_mfma_fp4(v[76:79], v[148:151], v[52:55], 3, 3, v[176], v[170])) + k.emit(s_addc_u32(s[17], 0, s[17])) + k.emit(s_sub_u32(s[18], s[18], s[67])) + k.emit(ds_read_b32(v[172], v[188], v[0], v[0], 0, 0, 12)) + k.emit(s_waitcnt(50559)) + k.emit(v_mfma_fp4(v[96:99], v[136:139], v[56:59], 0, 0, v[176], v[171])) + k.emit(s_add_u32(s[24], s[24], s[69])) + k.emit(ds_read_b128(v[88:91], v[185], v[0], v[0], 0, 128, 82)) + k.emit(v_mfma_fp4(v[100:103], v[136:139], v[60:63], 2, 0, v[176], v[171])) + k.emit(s_addc_u32(s[25], 0, s[25])) + k.emit(v_mfma_fp4(v[104:107], v[140:143], v[56:59], 1, 0, v[176], v[171])) + k.emit(s_sub_u32(s[26], s[26], s[69])) + for i in range(2): + k.emit(ds_read_b128(v[96 + i * -4:99 + i * -4], v[185], v[0], v[0], 0, 192 + i * -64, 82 + i * 2)) + k.emit(v_mfma_fp4(v[108 + i * -8:111 + i * -8], v[140 + i * 4:143 + i * 4], v[60 + i * 8:63 + i * 8], 3 + i * -1, 0 + i * 3, v[176], v[171])) + k.emit(v_mfma_fp4(v[96 + i * 8:99 + i * 8], v[144 + i * 4:147 + i * 4], v[64:67], 0 + i * 1, 3, v[176], v[171])) + k.emit(ds_read_b128(v[100:103], v[185], v[0], v[0], 0, 192, 84)) + k.emit(v_mfma_fp4(v[108:111], v[148:151], v[68:71], 3, 3, v[176], v[171])) + k.emit(ds_read_b32(v[173], v[188], v[0], v[0], 0, 0, 13)) + k.emit(s_barrier()) + k.emit(s_waitcnt(50559)) + for i in range(2): + k.emit(v_mfma_fp4(v[128:131], v[136 + i * 8:139 + i * 8], v[72 + i * 8:75 + i * 8], 0, 0 + i * 3, v[176], v[172])) + k.emit(v_mfma_fp4(v[132:135], v[136 + i * 8:139 + i * 8], v[76 + i * 8:79 + i * 8], 2, 0 + i * 3, v[176], v[172])) + k.emit(s_add_u32(NULL, LIT, s[65], 2048 + i * 1024)) + k.emit(buffer_load_dword(v[0], v[186 + i * 1], s[20:23], 0, 0, 1, 0, 0, 0, 0, 1)) + k.emit(v_mfma_fp4(v[136:139], v[140 + i * 8:143 + i * 8], v[72 + i * 8:75 + i * 8], 1, 0 + i * 3, v[176], v[172])) + k.emit(v_mfma_fp4(v[140:143], v[140 + i * 8:143 + i * 8], v[76 + i * 8:79 + i * 8], 3, 0 + i * 3, v[176], v[172])) + k.emit(s_waitcnt(49279)) + for i in range(2): + k.emit(v_mfma_fp4(v[160:163], v[136 + i * 8:139 + i * 8], v[88 + i * 8:91 + i * 8], 0, 0 + i * 3, v[176], v[173])) + k.emit(v_mfma_fp4(v[164:167], v[136 + i * 8:139 + i * 8], v[92 + i * 8:95 + i * 8], 2, 0 + i * 3, v[176], v[173])) + k.emit(s_add_u32(NULL, LIT, s[64], 25344 + i * 4224)) + k.emit(buffer_load_dwordx4(v[0:3], v[178 + i * 1], s[12:15], 0, 0, 1, 0, 0, 0, 0, 1)) + k.emit(v_mfma_fp4(v[168:171], v[140 + i * 8:143 + i * 8], v[88 + i * 8:91 + i * 8], 1, 0 + i * 3, v[176], v[173])) + k.emit(v_mfma_fp4(v[172:175], v[140 + i * 8:143 + i * 8], v[92 + i * 8:95 + i * 8], 3, 0 + i * 3, v[176], v[173])) + k.emit(s_waitcnt(20338)) + k.emit(s_barrier()) + k.emit(v_mfma_fp4(v[16:19], v[152:155], v[8:11], 0, 0, v[177], v[168])) + k.emit(v_mfma_fp4(v[20:23], v[152:155], v[12:15], 2, 0, v[177], v[168])) + k.emit(s_add_u32(NULL, LIT, s[64], 33792)) + k.emit(buffer_load_dwordx4(v[0:3], v[180], s[12:15], 0, 0, 1, 0, 0, 0, 0, 1)) + k.emit(v_mfma_fp4(v[24:27], v[156:159], v[8:11], 1, 0, v[177], v[168])) + k.emit(v_mfma_fp4(v[28:31], v[156:159], v[12:15], 3, 0, v[177], v[168])) + for i in range(2): + k.emit(v_mfma_fp4(v[16 + i * 32:19 + i * 32], v[160 + i * -8:163 + i * -8], v[16 + i * 8:19 + i * 8], 0, 3 + i * -3, v[177], v[168 + i * 1])) + k.emit(v_mfma_fp4(v[20 + i * 32:23 + i * 32], v[160 + i * -8:163 + i * -8], v[20 + i * 8:23 + i * 8], 2, 3 + i * -3, v[177], v[168 + i * 1])) + k.emit(s_add_u32(NULL, LIT, s[64], 38016 + i * 4224)) + k.emit(buffer_load_dwordx4(v[0:3], v[181 + i * 1], s[12:15], 0, 0, 1, 0, 0, 0, 0, 1)) + k.emit(v_mfma_fp4(v[24 + i * 32:27 + i * 32], v[164 + i * -8:167 + i * -8], v[16 + i * 8:19 + i * 8], 1, 3 + i * -3, v[177], v[168 + i * 1])) + k.emit(v_mfma_fp4(v[28 + i * 32:31 + i * 32], v[164 + i * -8:167 + i * -8], v[20 + i * 8:23 + i * 8], 3, 3 + i * -3, v[177], v[168 + i * 1])) + k.emit(v_mfma_fp4(v[48:51], v[160:163], v[32:35], 0, 3, v[177], v[169])) + k.emit(v_mfma_fp4(v[52:55], v[160:163], v[36:39], 2, 3, v[177], v[169])) + k.emit(s_add_u32(NULL, LIT, s[64], 46464)) + k.emit(buffer_load_dwordx4(v[0:3], v[183], s[12:15], 0, 0, 1, 0, 0, 0, 0, 1)) + k.emit(v_mfma_fp4(v[56:59], v[164:167], v[32:35], 1, 3, v[177], v[169])) + k.emit(s_add_u32(s[62], LIT, s[60], 768)) + k.emit(v_mfma_fp4(v[60:63], v[164:167], v[36:39], 3, 3, v[177], v[169])) + k.emit(s_cmp_lt_u32(s[62], s[61])) + k.emit(v_mfma_fp4(v[80:83], v[152:155], v[40:43], 0, 0, v[177], v[170])) + k.emit(s_cselect_b32(s[66], s[66], 0)) + k.emit(v_mfma_fp4(v[84:87], v[152:155], v[44:47], 2, 0, v[177], v[170])) + k.emit(s_cselect_b32(s[68], s[68], 0)) + k.emit(buffer_load_dwordx4(v[136:139], v[189], s[16:19], 0, 0, 1)) + k.emit(v_mfma_fp4(v[88:91], v[156:159], v[40:43], 1, 0, v[177], v[170])) + k.emit(s_add_u32(s[12], s[12], s[66])) + k.emit(v_mfma_fp4(v[92:95], v[156:159], v[44:47], 3, 0, v[177], v[170])) + k.emit(s_addc_u32(s[13], 0, s[13])) + k.emit(v_mfma_fp4(v[80:83], v[160:163], v[48:51], 0, 3, v[177], v[170])) + k.emit(s_sub_u32(s[14], s[14], s[66])) + k.emit(v_mfma_fp4(v[84:87], v[160:163], v[52:55], 2, 3, v[177], v[170])) + k.emit(s_add_u32(s[20], s[20], s[68])) + k.emit(buffer_load_dwordx4(v[140:143], v[190], s[16:19], 0, 0, 1)) + k.emit(v_mfma_fp4(v[88:91], v[164:167], v[48:51], 1, 3, v[177], v[170])) + k.emit(s_addc_u32(s[21], 0, s[21])) + k.emit(v_mfma_fp4(v[92:95], v[164:167], v[52:55], 3, 3, v[177], v[170])) + k.emit(s_sub_u32(s[22], s[22], s[68])) + k.emit(v_mfma_fp4(v[112:115], v[152:155], v[56:59], 0, 0, v[177], v[171])) + k.emit(s_addk_i32(s[60], 256)) + k.emit(v_mfma_fp4(v[116:119], v[152:155], v[60:63], 2, 0, v[177], v[171])) + k.emit(s_cmp_lt_i32(s[60], s[61])) + k.emit(buffer_load_dwordx4(v[144:147], v[189], s[16:19], 0, 1024, 1)) + k.emit(v_mfma_fp4(v[120:123], v[156:159], v[56:59], 1, 0, v[177], v[171])) + k.emit(v_mfma_fp4(v[124:127], v[156:159], v[60:63], 3, 0, v[177], v[171])) + k.emit(v_mfma_fp4(v[112:115], v[160:163], v[64:67], 0, 3, v[177], v[171])) + k.emit(v_mfma_fp4(v[116:119], v[160:163], v[68:71], 2, 3, v[177], v[171])) + k.emit(buffer_load_dwordx4(v[148:151], v[190], s[16:19], 0, 1024, 1)) + k.emit(v_mfma_fp4(v[120:123], v[164:167], v[64:67], 1, 3, v[177], v[171])) + k.emit(v_mfma_fp4(v[124:127], v[164:167], v[68:71], 3, 3, v[177], v[171])) + k.emit(v_mfma_fp4(v[144:147], v[152:155], v[72:75], 0, 0, v[177], v[172])) + k.emit(ds_read_b128(v[8:11], v[184])) + k.emit(v_mfma_fp4(v[148:151], v[152:155], v[76:79], 2, 0, v[177], v[172])) + k.emit(buffer_load_dword(v[176], v[193], s[24:27], 0, 0, 1)) + k.emit(v_mfma_fp4(v[152:155], v[156:159], v[72:75], 1, 0, v[177], v[172])) + k.emit(ds_read_b128(v[16:19], v[184], v[0], v[0], 0, 64)) + for i in range(2): + k.emit(v_mfma_fp4(v[156 + i * -8:159 + i * -8], v[156 + i * 4:159 + i * 4], v[76 + i * 8:79 + i * 8], 3 + i * -1, 0 + i * 3, v[177], v[172])) + k.emit(v_mfma_fp4(v[144 + i * 8:147 + i * 8], v[160 + i * 4:163 + i * 4], v[80:83], 0 + i * 1, 3, v[177], v[172])) + k.emit(ds_read_b128(v[12 + i * 8:15 + i * 8], v[184], v[0], v[0], 0, 0 + i * 64, 2)) + k.emit(v_mfma_fp4(v[156:159], v[164:167], v[84:87], 3, 3, v[177], v[172])) + k.emit(ds_read_b32(v[168], v[188])) + for i in range(2): + for j57 in range(2): + k.emit(v_mfma_fp4(v[176 + j57 * 8:179 + j57 * 8], v[152 + j57 * 4 + i * 8:155 + j57 * 4 + i * 8], v[88 + i * 8:91 + i * 8], 0 + j57 * 1, 0 + i * 3, v[177], v[173])) + k.emit(ds_read_b128(v[24 + j57 * 8 + i * 4:27 + j57 * 8 + i * 4], v[184], v[0], v[0], 0, 128 + j57 * 64, 16 + i * 2)) + k.emit(v_mfma_fp4(v[180 + j57 * 8:183 + j57 * 8], v[152 + j57 * 4 + i * 8:155 + j57 * 4 + i * 8], v[92 + i * 8:95 + i * 8], 2 + j57 * 1, 0 + i * 3, v[177], v[173])) + k.emit(ds_read_b32(v[169], v[188], v[0], v[0], 0, 0, 1)) + k.emit(s_cbranch_scc0(1064), target='L1_2EC8') + k.emit(s_branch(64473), target='L1_0D90') + k.label('L1_1E2C') + k.emit(s_waitcnt(20338)) + k.emit(s_barrier()) + for i in range(2): + k.emit(s_waitcnt(50559)) + for j58 in range(2): + k.emit(v_mfma_fp4(v[0 + i * 32:3 + i * 32], v[104 + j58 * 8:107 + j58 * 8], v[8 + j58 * 8 + i * 16:11 + j58 * 8 + i * 16], 0, 0 + j58 * 3, v[174], v[168 + i * 1])) + k.emit(buffer_load_dwordx4(v[152 + j58 * 4 + i * 8:155 + j58 * 4 + i * 8], v[191 + j58 * 1], s[16:19], 0, 0 + i * 1024, 1)) + k.emit(v_mfma_fp4(v[4 + i * 32:7 + i * 32], v[104 + j58 * 8:107 + j58 * 8], v[12 + j58 * 8 + i * 16:15 + j58 * 8 + i * 16], 2, 0 + j58 * 3, v[174], v[168 + i * 1])) + k.emit(ds_read_b128(v[40 + j58 * 4 + i * 16:43 + j58 * 4 + i * 16], v[184], v[0], v[0], 0, 0 + i * 128, 33 + j58 * 2 + i * 16)) + k.emit(v_mfma_fp4(v[8 + i * 32:11 + i * 32], v[108 + j58 * 8:111 + j58 * 8], v[8 + j58 * 8 + i * 16:11 + j58 * 8 + i * 16], 1, 0 + j58 * 3, v[174], v[168 + i * 1])) + k.emit(v_mfma_fp4(v[12 + i * 32:15 + i * 32], v[108 + j58 * 8:111 + j58 * 8], v[12 + j58 * 8 + i * 16:15 + j58 * 8 + i * 16], 3, 0 + j58 * 3, v[174], v[168 + i * 1])) + k.emit(ds_read_b128(v[48 + j58 * 4 + i * 16:51 + j58 * 4 + i * 16], v[184], v[0], v[0], 0, 64 + i * 128, 33 + j58 * 2 + i * 16)) + k.emit(ds_read_b32(v[170 + i * 1], v[188], v[0], v[0], 0, 0, 2 + i * 1)) + k.emit(s_waitcnt(50559)) + k.emit(v_mfma_fp4(v[64:67], v[104:107], v[40:43], 0, 0, v[174], v[170])) + k.emit(buffer_load_dword(v[177], v[194], s[24:27], 0, 0, 1)) + k.emit(v_mfma_fp4(v[68:71], v[104:107], v[44:47], 2, 0, v[174], v[170])) + k.emit(s_add_u32(s[63], LIT, s[60], 512)) + k.emit(ds_read_b128(v[72:75], v[184], v[0], v[0], 0, 0, 66)) + k.emit(v_mfma_fp4(v[72:75], v[108:111], v[40:43], 1, 0, v[174], v[170])) + k.emit(s_cmp_lt_u32(s[63], s[61])) + k.emit(v_mfma_fp4(v[76:79], v[108:111], v[44:47], 3, 0, v[174], v[170])) + k.emit(s_cselect_b32(s[67], s[67], 0)) + k.emit(ds_read_b128(v[80:83], v[184], v[0], v[0], 0, 64, 66)) + k.emit(v_mfma_fp4(v[64:67], v[112:115], v[48:51], 0, 3, v[174], v[170])) + k.emit(s_cselect_b32(s[69], s[69], 0)) + k.emit(v_mfma_fp4(v[68:71], v[112:115], v[52:55], 2, 3, v[174], v[170])) + k.emit(s_add_u32(s[16], s[16], s[67])) + k.emit(ds_read_b128(v[76:79], v[184], v[0], v[0], 0, 0, 68)) + k.emit(v_mfma_fp4(v[72:75], v[116:119], v[48:51], 1, 3, v[174], v[170])) + k.emit(s_addc_u32(s[17], 0, s[17])) + k.emit(v_mfma_fp4(v[76:79], v[116:119], v[52:55], 3, 3, v[174], v[170])) + k.emit(s_sub_u32(s[18], s[18], s[67])) + k.emit(ds_read_b128(v[84:87], v[184], v[0], v[0], 0, 64, 68)) + k.emit(ds_read_b32(v[172], v[188], v[0], v[0], 0, 0, 4)) + k.emit(s_waitcnt(50559)) + k.emit(v_mfma_fp4(v[96:99], v[104:107], v[56:59], 0, 0, v[174], v[171])) + k.emit(s_add_u32(s[24], s[24], s[69])) + k.emit(v_mfma_fp4(v[100:103], v[104:107], v[60:63], 2, 0, v[174], v[171])) + k.emit(s_addc_u32(s[25], 0, s[25])) + k.emit(ds_read_b128(v[88:91], v[184], v[0], v[0], 0, 128, 82)) + k.emit(v_mfma_fp4(v[104:107], v[108:111], v[56:59], 1, 0, v[174], v[171])) + k.emit(s_sub_u32(s[26], s[26], s[69])) + for i in range(2): + k.emit(v_mfma_fp4(v[108 + i * -8:111 + i * -8], v[108 + i * 4:111 + i * 4], v[60 + i * 8:63 + i * 8], 3 + i * -1, 0 + i * 3, v[174], v[171])) + k.emit(ds_read_b128(v[96 + i * -4:99 + i * -4], v[184], v[0], v[0], 0, 192 + i * -64, 82 + i * 2)) + k.emit(v_mfma_fp4(v[96 + i * 8:99 + i * 8], v[112 + i * 4:115 + i * 4], v[64:67], 0 + i * 1, 3, v[174], v[171])) + k.emit(v_mfma_fp4(v[108:111], v[116:119], v[68:71], 3, 3, v[174], v[171])) + k.emit(ds_read_b128(v[100:103], v[184], v[0], v[0], 0, 192, 84)) + k.emit(ds_read_b32(v[173], v[188], v[0], v[0], 0, 0, 5)) + k.emit(s_barrier()) + k.emit(s_waitcnt(50559)) + k.emit(v_mfma_fp4(v[128:131], v[104:107], v[72:75], 0, 0, v[174], v[172])) + k.emit(s_add_u32(NULL, 0, s[65])) + k.emit(buffer_load_dword(v[0], v[186], s[20:23], 0, 0, 1, 0, 0, 0, 0, 1)) + k.emit(v_mfma_fp4(v[132:135], v[104:107], v[76:79], 2, 0, v[174], v[172])) + k.emit(v_mfma_fp4(v[136:139], v[108:111], v[72:75], 1, 0, v[174], v[172])) + k.emit(v_mfma_fp4(v[140:143], v[108:111], v[76:79], 3, 0, v[174], v[172])) + k.emit(v_mfma_fp4(v[128:131], v[112:115], v[80:83], 0, 3, v[174], v[172])) + k.emit(s_add_u32(NULL, LIT, s[65], 1024)) + k.emit(buffer_load_dword(v[0], v[187], s[20:23], 0, 0, 1, 0, 0, 0, 0, 1)) + k.emit(v_mfma_fp4(v[132:135], v[112:115], v[84:87], 2, 3, v[174], v[172])) + k.emit(v_mfma_fp4(v[136:139], v[116:119], v[80:83], 1, 3, v[174], v[172])) + k.emit(v_mfma_fp4(v[140:143], v[116:119], v[84:87], 3, 3, v[174], v[172])) + k.emit(s_waitcnt(49279)) + k.emit(v_mfma_fp4(v[160:163], v[104:107], v[88:91], 0, 0, v[174], v[173])) + k.emit(s_add_u32(NULL, 0, s[64])) + k.emit(buffer_load_dwordx4(v[0:3], v[178], s[12:15], 0, 0, 1, 0, 0, 0, 0, 1)) + k.emit(v_mfma_fp4(v[164:167], v[104:107], v[92:95], 2, 0, v[174], v[173])) + k.emit(v_mfma_fp4(v[168:171], v[108:111], v[88:91], 1, 0, v[174], v[173])) + k.emit(v_mfma_fp4(v[172:175], v[108:111], v[92:95], 3, 0, v[174], v[173])) + k.emit(v_mfma_fp4(v[160:163], v[112:115], v[96:99], 0, 3, v[174], v[173])) + k.emit(s_add_u32(NULL, LIT, s[64], 4224)) + k.emit(buffer_load_dwordx4(v[0:3], v[179], s[12:15], 0, 0, 1, 0, 0, 0, 0, 1)) + k.emit(v_mfma_fp4(v[164:167], v[112:115], v[100:103], 2, 3, v[174], v[173])) + k.emit(v_mfma_fp4(v[168:171], v[116:119], v[96:99], 1, 3, v[174], v[173])) + k.emit(v_mfma_fp4(v[172:175], v[116:119], v[100:103], 3, 3, v[174], v[173])) + k.emit(s_waitcnt(20338)) + k.emit(s_barrier()) + k.emit(v_mfma_fp4(v[16:19], v[120:123], v[8:11], 0, 0, v[175], v[168])) + k.emit(s_add_u32(NULL, LIT, s[64], 8448)) + k.emit(buffer_load_dwordx4(v[0:3], v[180], s[12:15], 0, 0, 1, 0, 0, 0, 0, 1)) + k.emit(v_mfma_fp4(v[20:23], v[120:123], v[12:15], 2, 0, v[175], v[168])) + k.emit(v_mfma_fp4(v[24:27], v[124:127], v[8:11], 1, 0, v[175], v[168])) + for i in range(2): + k.emit(v_mfma_fp4(v[28:31], v[124 + i * 8:127 + i * 8], v[12 + i * 8:15 + i * 8], 3, 0 + i * 3, v[175], v[168])) + k.emit(v_mfma_fp4(v[16 + i * 32:19 + i * 32], v[128 + i * -8:131 + i * -8], v[16 + i * 8:19 + i * 8], 0, 3 + i * -3, v[175], v[168 + i * 1])) + k.emit(s_add_u32(NULL, LIT, s[64], 12672 + i * 4224)) + k.emit(buffer_load_dwordx4(v[0:3], v[181 + i * 1], s[12:15], 0, 0, 1, 0, 0, 0, 0, 1)) + k.emit(v_mfma_fp4(v[20 + i * 32:23 + i * 32], v[128 + i * -8:131 + i * -8], v[20 + i * 8:23 + i * 8], 2, 3 + i * -3, v[175], v[168 + i * 1])) + k.emit(v_mfma_fp4(v[24 + i * 32:27 + i * 32], v[132 + i * -8:135 + i * -8], v[16 + i * 8:19 + i * 8], 1, 3 + i * -3, v[175], v[168 + i * 1])) + k.emit(v_mfma_fp4(v[60:63], v[124:127], v[28:31], 3, 0, v[175], v[169])) + k.emit(v_mfma_fp4(v[48:51], v[128:131], v[32:35], 0, 3, v[175], v[169])) + k.emit(s_add_u32(NULL, LIT, s[64], 21120)) + k.emit(buffer_load_dwordx4(v[0:3], v[183], s[12:15], 0, 0, 1, 0, 0, 0, 0, 1)) + k.emit(v_mfma_fp4(v[52:55], v[128:131], v[36:39], 2, 3, v[175], v[169])) + k.emit(s_add_u32(s[62], LIT, s[60], 768)) + k.emit(v_mfma_fp4(v[56:59], v[132:135], v[32:35], 1, 3, v[175], v[169])) + k.emit(s_cmp_lt_u32(s[62], s[61])) + k.emit(v_mfma_fp4(v[60:63], v[132:135], v[36:39], 3, 3, v[175], v[169])) + k.emit(s_cselect_b32(s[66], s[66], 0)) + k.emit(v_mfma_fp4(v[80:83], v[120:123], v[40:43], 0, 0, v[175], v[170])) + k.emit(s_cselect_b32(s[68], s[68], 0)) + k.emit(buffer_load_dwordx4(v[104:107], v[189], s[16:19], 0, 0, 1)) + k.emit(v_mfma_fp4(v[84:87], v[120:123], v[44:47], 2, 0, v[175], v[170])) + k.emit(s_add_u32(s[12], s[12], s[66])) + k.emit(v_mfma_fp4(v[88:91], v[124:127], v[40:43], 1, 0, v[175], v[170])) + k.emit(s_addc_u32(s[13], 0, s[13])) + k.emit(v_mfma_fp4(v[92:95], v[124:127], v[44:47], 3, 0, v[175], v[170])) + k.emit(s_sub_u32(s[14], s[14], s[66])) + k.emit(v_mfma_fp4(v[80:83], v[128:131], v[48:51], 0, 3, v[175], v[170])) + k.emit(s_add_u32(s[20], s[20], s[68])) + k.emit(buffer_load_dwordx4(v[108:111], v[190], s[16:19], 0, 0, 1)) + k.emit(v_mfma_fp4(v[84:87], v[128:131], v[52:55], 2, 3, v[175], v[170])) + k.emit(s_addc_u32(s[21], 0, s[21])) + k.emit(v_mfma_fp4(v[88:91], v[132:135], v[48:51], 1, 3, v[175], v[170])) + k.emit(s_sub_u32(s[22], s[22], s[68])) + k.emit(v_mfma_fp4(v[92:95], v[132:135], v[52:55], 3, 3, v[175], v[170])) + k.emit(s_addk_i32(s[60], 256)) + k.emit(v_mfma_fp4(v[112:115], v[120:123], v[56:59], 0, 0, v[175], v[171])) + k.emit(s_cmp_lt_i32(s[60], s[61])) + for i in range(2): + k.emit(buffer_load_dwordx4(v[112 + i * 4:115 + i * 4], v[189 + i * 1], s[16:19], 0, 1024, 1)) + k.emit(v_mfma_fp4(v[116:119], v[120 + i * 8:123 + i * 8], v[60 + i * 8:63 + i * 8], 2, 0 + i * 3, v[175], v[171])) + k.emit(v_mfma_fp4(v[120:123], v[124 + i * 8:127 + i * 8], v[56 + i * 8:59 + i * 8], 1, 0 + i * 3, v[175], v[171])) + k.emit(v_mfma_fp4(v[124:127], v[124 + i * 8:127 + i * 8], v[60 + i * 8:63 + i * 8], 3, 0 + i * 3, v[175], v[171])) + k.emit(v_mfma_fp4(v[112 + i * 32:115 + i * 32], v[128 + i * -8:131 + i * -8], v[64 + i * 8:67 + i * 8], 0, 3 + i * -3, v[175], v[171 + i * 1])) + k.emit(buffer_load_dword(v[174], v[193], s[24:27], 0, 0, 1)) + k.emit(v_mfma_fp4(v[148:151], v[120:123], v[76:79], 2, 0, v[175], v[172])) + k.emit(ds_read_b128(v[8:11], v[185])) + k.emit(v_mfma_fp4(v[152:155], v[124:127], v[72:75], 1, 0, v[175], v[172])) + k.emit(v_mfma_fp4(v[156:159], v[124:127], v[76:79], 3, 0, v[175], v[172])) + k.emit(ds_read_b128(v[16:19], v[185], v[0], v[0], 0, 64)) + for i in range(2): + k.emit(v_mfma_fp4(v[144 + i * 8:147 + i * 8], v[128 + i * 4:131 + i * 4], v[80:83], 0 + i * 1, 3, v[175], v[172])) + k.emit(v_mfma_fp4(v[148 + i * 8:151 + i * 8], v[128 + i * 4:131 + i * 4], v[84:87], 2 + i * 1, 3, v[175], v[172])) + k.emit(ds_read_b128(v[12 + i * 8:15 + i * 8], v[185], v[0], v[0], 0, 0 + i * 64, 2)) + k.emit(ds_read_b32(v[168], v[188], v[0], v[0], 0, 0, 8)) + for i in range(2): + for j59 in range(2): + k.emit(v_mfma_fp4(v[176 + j59 * 8:179 + j59 * 8], v[120 + j59 * 4 + i * 8:123 + j59 * 4 + i * 8], v[88 + i * 8:91 + i * 8], 0 + j59 * 1, 0 + i * 3, v[175], v[173])) + k.emit(v_mfma_fp4(v[180 + j59 * 8:183 + j59 * 8], v[120 + j59 * 4 + i * 8:123 + j59 * 4 + i * 8], v[92 + i * 8:95 + i * 8], 2 + j59 * 1, 0 + i * 3, v[175], v[173])) + k.emit(ds_read_b128(v[24 + j59 * 8 + i * 4:27 + j59 * 8 + i * 4], v[185], v[0], v[0], 0, 128 + j59 * 64, 16 + i * 2)) + k.emit(ds_read_b32(v[169], v[188], v[0], v[0], 0, 0, 9)) + k.emit(s_cbranch_scc0(533), target='L1_2EC8') + k.emit(s_waitcnt(20338)) + k.emit(s_barrier()) + for i in range(2): + k.emit(s_waitcnt(50559)) + for j60 in range(2): + k.emit(v_mfma_fp4(v[0 + i * 32:3 + i * 32], v[136 + j60 * 8:139 + j60 * 8], v[8 + j60 * 8 + i * 16:11 + j60 * 8 + i * 16], 0, 0 + j60 * 3, v[176], v[168 + i * 1])) + k.emit(buffer_load_dwordx4(v[120 + j60 * 4 + i * 8:123 + j60 * 4 + i * 8], v[191 + j60 * 1], s[16:19], 0, 0 + i * 1024, 1)) + k.emit(v_mfma_fp4(v[4 + i * 32:7 + i * 32], v[136 + j60 * 8:139 + j60 * 8], v[12 + j60 * 8 + i * 16:15 + j60 * 8 + i * 16], 2, 0 + j60 * 3, v[176], v[168 + i * 1])) + k.emit(ds_read_b128(v[40 + j60 * 4 + i * 16:43 + j60 * 4 + i * 16], v[185], v[0], v[0], 0, 0 + i * 128, 33 + j60 * 2 + i * 16)) + k.emit(v_mfma_fp4(v[8 + i * 32:11 + i * 32], v[140 + j60 * 8:143 + j60 * 8], v[8 + j60 * 8 + i * 16:11 + j60 * 8 + i * 16], 1, 0 + j60 * 3, v[176], v[168 + i * 1])) + k.emit(v_mfma_fp4(v[12 + i * 32:15 + i * 32], v[140 + j60 * 8:143 + j60 * 8], v[12 + j60 * 8 + i * 16:15 + j60 * 8 + i * 16], 3, 0 + j60 * 3, v[176], v[168 + i * 1])) + k.emit(ds_read_b128(v[48 + j60 * 4 + i * 16:51 + j60 * 4 + i * 16], v[185], v[0], v[0], 0, 64 + i * 128, 33 + j60 * 2 + i * 16)) + k.emit(ds_read_b32(v[170 + i * 1], v[188], v[0], v[0], 0, 0, 10 + i * 1)) + k.emit(s_waitcnt(50559)) + k.emit(v_mfma_fp4(v[64:67], v[136:139], v[40:43], 0, 0, v[176], v[170])) + k.emit(buffer_load_dword(v[175], v[194], s[24:27], 0, 0, 1)) + k.emit(v_mfma_fp4(v[68:71], v[136:139], v[44:47], 2, 0, v[176], v[170])) + k.emit(s_add_u32(s[63], LIT, s[60], 512)) + k.emit(ds_read_b128(v[72:75], v[185], v[0], v[0], 0, 0, 66)) + k.emit(v_mfma_fp4(v[72:75], v[140:143], v[40:43], 1, 0, v[176], v[170])) + k.emit(s_cmp_lt_u32(s[63], s[61])) + k.emit(v_mfma_fp4(v[76:79], v[140:143], v[44:47], 3, 0, v[176], v[170])) + k.emit(s_cselect_b32(s[67], s[67], 0)) + k.emit(ds_read_b128(v[80:83], v[185], v[0], v[0], 0, 64, 66)) + k.emit(v_mfma_fp4(v[64:67], v[144:147], v[48:51], 0, 3, v[176], v[170])) + k.emit(s_cselect_b32(s[69], s[69], 0)) + k.emit(v_mfma_fp4(v[68:71], v[144:147], v[52:55], 2, 3, v[176], v[170])) + k.emit(s_add_u32(s[16], s[16], s[67])) + k.emit(ds_read_b128(v[76:79], v[185], v[0], v[0], 0, 0, 68)) + k.emit(v_mfma_fp4(v[72:75], v[148:151], v[48:51], 1, 3, v[176], v[170])) + k.emit(s_addc_u32(s[17], 0, s[17])) + k.emit(v_mfma_fp4(v[76:79], v[148:151], v[52:55], 3, 3, v[176], v[170])) + k.emit(s_sub_u32(s[18], s[18], s[67])) + k.emit(ds_read_b128(v[84:87], v[185], v[0], v[0], 0, 64, 68)) + k.emit(ds_read_b32(v[172], v[188], v[0], v[0], 0, 0, 12)) + k.emit(s_waitcnt(50559)) + k.emit(v_mfma_fp4(v[96:99], v[136:139], v[56:59], 0, 0, v[176], v[171])) + k.emit(s_add_u32(s[24], s[24], s[69])) + k.emit(v_mfma_fp4(v[100:103], v[136:139], v[60:63], 2, 0, v[176], v[171])) + k.emit(s_addc_u32(s[25], 0, s[25])) + k.emit(ds_read_b128(v[88:91], v[185], v[0], v[0], 0, 128, 82)) + k.emit(v_mfma_fp4(v[104:107], v[140:143], v[56:59], 1, 0, v[176], v[171])) + k.emit(s_sub_u32(s[26], s[26], s[69])) + for i in range(2): + k.emit(v_mfma_fp4(v[108 + i * -8:111 + i * -8], v[140 + i * 4:143 + i * 4], v[60 + i * 8:63 + i * 8], 3 + i * -1, 0 + i * 3, v[176], v[171])) + k.emit(ds_read_b128(v[96 + i * -4:99 + i * -4], v[185], v[0], v[0], 0, 192 + i * -64, 82 + i * 2)) + k.emit(v_mfma_fp4(v[96 + i * 8:99 + i * 8], v[144 + i * 4:147 + i * 4], v[64:67], 0 + i * 1, 3, v[176], v[171])) + k.emit(v_mfma_fp4(v[108:111], v[148:151], v[68:71], 3, 3, v[176], v[171])) + k.emit(ds_read_b128(v[100:103], v[185], v[0], v[0], 0, 192, 84)) + k.emit(ds_read_b32(v[173], v[188], v[0], v[0], 0, 0, 13)) + k.emit(s_barrier()) + k.emit(s_waitcnt(50559)) + for i in range(2): + k.emit(v_mfma_fp4(v[128:131], v[136 + i * 8:139 + i * 8], v[72 + i * 8:75 + i * 8], 0, 0 + i * 3, v[176], v[172])) + k.emit(s_add_u32(NULL, LIT, s[65], 2048 + i * 1024)) + k.emit(buffer_load_dword(v[0], v[186 + i * 1], s[20:23], 0, 0, 1, 0, 0, 0, 0, 1)) + k.emit(v_mfma_fp4(v[132:135], v[136 + i * 8:139 + i * 8], v[76 + i * 8:79 + i * 8], 2, 0 + i * 3, v[176], v[172])) + k.emit(v_mfma_fp4(v[136:139], v[140 + i * 8:143 + i * 8], v[72 + i * 8:75 + i * 8], 1, 0 + i * 3, v[176], v[172])) + k.emit(v_mfma_fp4(v[140:143], v[140 + i * 8:143 + i * 8], v[76 + i * 8:79 + i * 8], 3, 0 + i * 3, v[176], v[172])) + k.emit(s_waitcnt(49279)) + for i in range(2): + k.emit(v_mfma_fp4(v[160:163], v[136 + i * 8:139 + i * 8], v[88 + i * 8:91 + i * 8], 0, 0 + i * 3, v[176], v[173])) + k.emit(s_add_u32(NULL, LIT, s[64], 25344 + i * 4224)) + k.emit(buffer_load_dwordx4(v[0:3], v[178 + i * 1], s[12:15], 0, 0, 1, 0, 0, 0, 0, 1)) + k.emit(v_mfma_fp4(v[164:167], v[136 + i * 8:139 + i * 8], v[92 + i * 8:95 + i * 8], 2, 0 + i * 3, v[176], v[173])) + k.emit(v_mfma_fp4(v[168:171], v[140 + i * 8:143 + i * 8], v[88 + i * 8:91 + i * 8], 1, 0 + i * 3, v[176], v[173])) + k.emit(v_mfma_fp4(v[172:175], v[140 + i * 8:143 + i * 8], v[92 + i * 8:95 + i * 8], 3, 0 + i * 3, v[176], v[173])) + k.emit(s_waitcnt(20338)) + k.emit(s_barrier()) + k.emit(v_mfma_fp4(v[16:19], v[152:155], v[8:11], 0, 0, v[177], v[168])) + k.emit(s_add_u32(NULL, LIT, s[64], 33792)) + k.emit(buffer_load_dwordx4(v[0:3], v[180], s[12:15], 0, 0, 1, 0, 0, 0, 0, 1)) + k.emit(v_mfma_fp4(v[20:23], v[152:155], v[12:15], 2, 0, v[177], v[168])) + k.emit(v_mfma_fp4(v[24:27], v[156:159], v[8:11], 1, 0, v[177], v[168])) + for i in range(2): + k.emit(v_mfma_fp4(v[28:31], v[156 + i * 8:159 + i * 8], v[12 + i * 8:15 + i * 8], 3, 0 + i * 3, v[177], v[168])) + k.emit(v_mfma_fp4(v[16 + i * 32:19 + i * 32], v[160 + i * -8:163 + i * -8], v[16 + i * 8:19 + i * 8], 0, 3 + i * -3, v[177], v[168 + i * 1])) + k.emit(s_add_u32(NULL, LIT, s[64], 38016 + i * 4224)) + k.emit(buffer_load_dwordx4(v[0:3], v[181 + i * 1], s[12:15], 0, 0, 1, 0, 0, 0, 0, 1)) + k.emit(v_mfma_fp4(v[20 + i * 32:23 + i * 32], v[160 + i * -8:163 + i * -8], v[20 + i * 8:23 + i * 8], 2, 3 + i * -3, v[177], v[168 + i * 1])) + k.emit(v_mfma_fp4(v[24 + i * 32:27 + i * 32], v[164 + i * -8:167 + i * -8], v[16 + i * 8:19 + i * 8], 1, 3 + i * -3, v[177], v[168 + i * 1])) + k.emit(v_mfma_fp4(v[60:63], v[156:159], v[28:31], 3, 0, v[177], v[169])) + k.emit(v_mfma_fp4(v[48:51], v[160:163], v[32:35], 0, 3, v[177], v[169])) + k.emit(s_add_u32(NULL, LIT, s[64], 46464)) + k.emit(buffer_load_dwordx4(v[0:3], v[183], s[12:15], 0, 0, 1, 0, 0, 0, 0, 1)) + k.emit(v_mfma_fp4(v[52:55], v[160:163], v[36:39], 2, 3, v[177], v[169])) + k.emit(s_add_u32(s[62], LIT, s[60], 768)) + k.emit(v_mfma_fp4(v[56:59], v[164:167], v[32:35], 1, 3, v[177], v[169])) + k.emit(s_cmp_lt_u32(s[62], s[61])) + k.emit(v_mfma_fp4(v[60:63], v[164:167], v[36:39], 3, 3, v[177], v[169])) + k.emit(s_cselect_b32(s[66], s[66], 0)) + k.emit(v_mfma_fp4(v[80:83], v[152:155], v[40:43], 0, 0, v[177], v[170])) + k.emit(s_cselect_b32(s[68], s[68], 0)) + k.emit(buffer_load_dwordx4(v[136:139], v[189], s[16:19], 0, 0, 1)) + k.emit(v_mfma_fp4(v[84:87], v[152:155], v[44:47], 2, 0, v[177], v[170])) + k.emit(s_add_u32(s[12], s[12], s[66])) + k.emit(v_mfma_fp4(v[88:91], v[156:159], v[40:43], 1, 0, v[177], v[170])) + k.emit(s_addc_u32(s[13], 0, s[13])) + k.emit(v_mfma_fp4(v[92:95], v[156:159], v[44:47], 3, 0, v[177], v[170])) + k.emit(s_sub_u32(s[14], s[14], s[66])) + k.emit(v_mfma_fp4(v[80:83], v[160:163], v[48:51], 0, 3, v[177], v[170])) + k.emit(s_add_u32(s[20], s[20], s[68])) + k.emit(buffer_load_dwordx4(v[140:143], v[190], s[16:19], 0, 0, 1)) + k.emit(v_mfma_fp4(v[84:87], v[160:163], v[52:55], 2, 3, v[177], v[170])) + k.emit(s_addc_u32(s[21], 0, s[21])) + k.emit(v_mfma_fp4(v[88:91], v[164:167], v[48:51], 1, 3, v[177], v[170])) + k.emit(s_sub_u32(s[22], s[22], s[68])) + k.emit(v_mfma_fp4(v[92:95], v[164:167], v[52:55], 3, 3, v[177], v[170])) + k.emit(s_addk_i32(s[60], 256)) + k.emit(v_mfma_fp4(v[112:115], v[152:155], v[56:59], 0, 0, v[177], v[171])) + k.emit(s_cmp_lt_i32(s[60], s[61])) + for i in range(2): + k.emit(buffer_load_dwordx4(v[144 + i * 4:147 + i * 4], v[189 + i * 1], s[16:19], 0, 1024, 1)) + k.emit(v_mfma_fp4(v[116:119], v[152 + i * 8:155 + i * 8], v[60 + i * 8:63 + i * 8], 2, 0 + i * 3, v[177], v[171])) + k.emit(v_mfma_fp4(v[120:123], v[156 + i * 8:159 + i * 8], v[56 + i * 8:59 + i * 8], 1, 0 + i * 3, v[177], v[171])) + k.emit(v_mfma_fp4(v[124:127], v[156 + i * 8:159 + i * 8], v[60 + i * 8:63 + i * 8], 3, 0 + i * 3, v[177], v[171])) + k.emit(v_mfma_fp4(v[112 + i * 32:115 + i * 32], v[160 + i * -8:163 + i * -8], v[64 + i * 8:67 + i * 8], 0, 3 + i * -3, v[177], v[171 + i * 1])) + k.emit(buffer_load_dword(v[176], v[193], s[24:27], 0, 0, 1)) + k.emit(v_mfma_fp4(v[148:151], v[152:155], v[76:79], 2, 0, v[177], v[172])) + k.emit(ds_read_b128(v[8:11], v[184])) + k.emit(v_mfma_fp4(v[152:155], v[156:159], v[72:75], 1, 0, v[177], v[172])) + k.emit(v_mfma_fp4(v[156:159], v[156:159], v[76:79], 3, 0, v[177], v[172])) + k.emit(ds_read_b128(v[16:19], v[184], v[0], v[0], 0, 64)) + for i in range(2): + k.emit(v_mfma_fp4(v[144 + i * 8:147 + i * 8], v[160 + i * 4:163 + i * 4], v[80:83], 0 + i * 1, 3, v[177], v[172])) + k.emit(v_mfma_fp4(v[148 + i * 8:151 + i * 8], v[160 + i * 4:163 + i * 4], v[84:87], 2 + i * 1, 3, v[177], v[172])) + k.emit(ds_read_b128(v[12 + i * 8:15 + i * 8], v[184], v[0], v[0], 0, 0 + i * 64, 2)) + k.emit(ds_read_b32(v[168], v[188])) + for i in range(2): + for j61 in range(2): + k.emit(v_mfma_fp4(v[176 + j61 * 8:179 + j61 * 8], v[152 + j61 * 4 + i * 8:155 + j61 * 4 + i * 8], v[88 + i * 8:91 + i * 8], 0 + j61 * 1, 0 + i * 3, v[177], v[173])) + k.emit(v_mfma_fp4(v[180 + j61 * 8:183 + j61 * 8], v[152 + j61 * 4 + i * 8:155 + j61 * 4 + i * 8], v[92 + i * 8:95 + i * 8], 2 + j61 * 1, 0 + i * 3, v[177], v[173])) + k.emit(ds_read_b128(v[24 + j61 * 8 + i * 4:27 + j61 * 8 + i * 4], v[184], v[0], v[0], 0, 128 + j61 * 64, 16 + i * 2)) + k.emit(ds_read_b32(v[169], v[188], v[0], v[0], 0, 0, 1)) + k.emit(s_cbranch_scc0(1), target='L1_2EC8') + k.emit(s_branch(64473), target='L1_1E2C') + k.label('L1_2EC8') + k.emit(s_waitcnt()) + k.emit(s_barrier()) + k.emit(v_lshrrev_b32_e32(v[4], 5)) + k.emit(v_mul_i32_i24_e32(v[4], 16, v[4])) + k.emit(v_lshrrev_b32_e32(v[5], 4)) + k.emit(v_and_b32_e32(v[5], 1, v[5])) + k.emit(v_mul_i32_i24_e32(v[5], 32, v[5])) + k.emit(v_add_u32_e32(v[4], v[4], v[5])) + k.emit(v_and_b32_e32(v[5], 15)) + k.emit(v_mul_i32_i24_e32(v[5], LIT, v[5], 128)) + k.emit(v_add_u32_e32(v[4], v[4], v[5])) + k.emit(s_mul_i32(s[62], s[46], LIT, 24576)) + k.emit(s_add_i32(s[62], s[62], 0)) + k.emit(v_add_i32(v[4], v[4], s[62])) + for i in range(2): + for j62 in range(4): + k.emit(v_accvgpr_read(v[8 + j62 * 1 + i * 4], v[0 + j62 * 1 + i * 8])) + for i in range(4): + k.emit(v_cvt_pk_bf16_f32(v[16 + i * 1], v[8 + i * 2], v[9 + i * 2])) + k.emit(s_nop(1)) + k.emit(v_permlane16_swap_b32_e32(v[16], v[18])) + k.emit(s_nop(1)) + k.emit(v_permlane16_swap_b32_e32(v[17], v[19])) + k.emit(s_nop(1)) + k.emit(ds_write_b128(v[0], v[4], v[16:19])) + for i in range(2): + for j63 in range(4): + k.emit(v_accvgpr_read(v[8 + j63 * 1 + i * 4], v[16 + j63 * 1 + i * 8])) + for i in range(4): + k.emit(v_cvt_pk_bf16_f32(v[16 + i * 1], v[8 + i * 2], v[9 + i * 2])) + k.emit(s_nop(1)) + k.emit(v_permlane16_swap_b32_e32(v[16], v[18])) + k.emit(s_nop(1)) + k.emit(v_permlane16_swap_b32_e32(v[17], v[19])) + k.emit(s_nop(1)) + k.emit(ds_write_b128(v[0], v[4], v[16:19], v[0], 0, 64)) + for i in range(5): + for j64 in range(2): + k.emit(v_accvgpr_read(v[8], v[4 + j64 * 28 + i * 32])) + k.emit(v_accvgpr_read(v[9], v[5 + j64 * 28 + i * 32])) + k.emit(v_accvgpr_read(v[10], v[6 + j64 * 28 + i * 32])) + k.emit(v_accvgpr_read(v[11], v[7 + j64 * 28 + i * 32])) + k.emit(v_accvgpr_read(v[12], v[12 + j64 * 28 + i * 32])) + k.emit(v_accvgpr_read(v[13], v[13 + j64 * 28 + i * 32])) + k.emit(v_accvgpr_read(v[14], v[14 + j64 * 28 + i * 32])) + k.emit(v_accvgpr_read(v[15], v[15 + j64 * 28 + i * 32])) + k.emit(v_cvt_pk_bf16_f32(v[16], v[8], v[9])) + k.emit(v_cvt_pk_bf16_f32(v[17], v[10], v[11])) + k.emit(v_cvt_pk_bf16_f32(v[18], v[12], v[13])) + k.emit(v_cvt_pk_bf16_f32(v[19], v[14], v[15])) + k.emit(s_nop(1)) + k.emit(v_permlane16_swap_b32_e32(v[16], v[18])) + k.emit(s_nop(1)) + k.emit(v_permlane16_swap_b32_e32(v[17], v[19])) + k.emit(s_nop(1)) + k.emit(ds_write_b128(v[0], v[4], v[16:19], v[0], 0, 0, 8 + j64 * 8 + i * 16)) + k.emit(v_accvgpr_read(v[8], v[20 + j64 * 28 + i * 32])) + k.emit(v_accvgpr_read(v[9], v[21 + j64 * 28 + i * 32])) + k.emit(v_accvgpr_read(v[10], v[22 + j64 * 28 + i * 32])) + k.emit(v_accvgpr_read(v[11], v[23 + j64 * 28 + i * 32])) + k.emit(v_accvgpr_read(v[12], v[28 + j64 * 28 + i * 32])) + k.emit(v_accvgpr_read(v[13], v[29 + j64 * 28 + i * 32])) + k.emit(v_accvgpr_read(v[14], v[30 + j64 * 28 + i * 32])) + k.emit(v_accvgpr_read(v[15], v[31 + j64 * 28 + i * 32])) + k.emit(v_cvt_pk_bf16_f32(v[16], v[8], v[9])) + k.emit(v_cvt_pk_bf16_f32(v[17], v[10], v[11])) + k.emit(v_cvt_pk_bf16_f32(v[18], v[12], v[13])) + k.emit(v_cvt_pk_bf16_f32(v[19], v[14], v[15])) + k.emit(s_nop(1)) + k.emit(v_permlane16_swap_b32_e32(v[16], v[18])) + k.emit(s_nop(1)) + k.emit(v_permlane16_swap_b32_e32(v[17], v[19])) + k.emit(s_nop(1)) + k.emit(ds_write_b128(v[0], v[4], v[16:19], v[0], 0, 64, 8 + j64 * 8 + i * 16)) + for i in range(2): + for j65 in range(2): + k.emit(v_accvgpr_read(v[8 + j65 * 4], v[164 + j65 * 8 + i * 16])) + k.emit(v_accvgpr_read(v[9 + j65 * 4], v[165 + j65 * 8 + i * 16])) + k.emit(v_accvgpr_read(v[10 + j65 * 4], v[166 + j65 * 8 + i * 16])) + k.emit(v_accvgpr_read(v[11 + j65 * 4], v[167 + j65 * 8 + i * 16])) + for j66 in range(4): + k.emit(v_cvt_pk_bf16_f32(v[16 + j66 * 1], v[8 + j66 * 2], v[9 + j66 * 2])) + k.emit(s_nop(1)) + k.emit(v_permlane16_swap_b32_e32(v[16], v[18])) + k.emit(s_nop(1)) + k.emit(v_permlane16_swap_b32_e32(v[17], v[19])) + k.emit(s_nop(1)) + k.emit(ds_write_b128(v[0], v[4], v[16:19], v[0], 0, 0 + i * 64, 88)) + k.emit(s_waitcnt(49279)) + k.emit(v_mul_i32_i24_e64(v[4], v[0], 16)) + k.emit(v_add_i32(v[4], v[4], s[62])) + k.emit(ds_read_b128(v[16:19], v[4])) + for i in range(23): + k.emit(s_waitcnt(49279)) + k.emit(buffer_store_dwordx4(v[16:19], v[195 + i * 1], s[4:7], 0, 0, 1)) + k.emit(ds_read_b128(v[16:19], v[4], v[0], v[0], 0, 0, 4 + i * 4)) + k.emit(s_waitcnt(49279)) + k.emit(buffer_store_dwordx4(v[16:19], v[218], s[4:7], 0, 0, 1)) + k.emit(s_waitcnt()) + k.emit(s_endpgm()) + elif (tile_m, tile_n) == (256, 256): + k.emit(s_and_b32(s[1], s[1], LIT, 65535)) + k.emit(s_mov_b32(s[56], s[4])) + k.emit(s_load_dwordx2(s[4:5], s[0:1], s[0], 0, 0, 0, 0, 1)) + k.emit(s_mov_b32(s[8], 0)) + k.emit(s_mov_b32(s[9], 0)) + k.emit(s_load_dwordx2(s[12:13], s[0:1], s[0], 8, 0, 0, 0, 1)) + k.emit(s_load_dwordx2(s[16:17], s[0:1], s[0], 16, 0, 0, 0, 1)) + k.emit(s_mov_b32(s[38], 1.0)) + k.emit(s_mov_b32(s[39], 0)) + k.emit(s_mov_b32(s[40], N)) + k.emit(s_mov_b32(s[41], K)) + k.emit(s_mov_b32(s[42], K)) + k.emit(s_mov_b32(s[43], M)) + k.emit(s_mov_b32(s[44], N)) + k.emit(s_mov_b32(s[45], K)) + k.emit(s_load_dwordx2(s[20:21], s[0:1], s[0], 24, 0, 0, 0, 1)) + k.emit(s_load_dwordx2(s[24:25], s[0:1], s[0], 32, 0, 0, 0, 1)) + k.emit(s_mov_b32(s[36], scale_k)) + k.emit(s_mov_b32(s[37], scale_k)) + k.emit(s_mov_b32(s[57], 0)) + k.emit(v_lshrrev_b32_e32(v[1], 10)) + k.emit(v_lshrrev_b32_e32(v[2], 10, v[1])) + k.emit(v_and_b32_e32(v[2], LIT, v[2], 1023)) + k.emit(v_and_b32_e32(v[1], LIT, v[1], 1023)) + k.emit(v_and_b32_e32(v[0], LIT, v[0], 1023)) + k.emit(v_lshrrev_b32_e32(v[3], 6)) + k.emit(v_and_b32_e32(v[0], 63)) + k.emit(s_mov_b32(s[46], s[2])) + k.emit(s_mov_b32(s[47], s[3])) + k.emit(v_readfirstlane_b32_e32(v[49], v[3])) + k.emit(s_waitcnt(49279)) + k.emit(s_add_u32(s[55], s[44], LIT, 255)) + k.emit(s_lshr_b32(s[54], s[55], 8)) + k.emit(s_mul_i32(s[48], s[54], s[47])) + k.emit(s_add_i32(s[48], s[48], s[46])) + k.emit(s_add_u32(s[55], s[43], LIT, 255)) + k.emit(s_lshr_b32(s[52], s[55], 8)) + k.emit(s_lshl_b32(s[52], s[52], 5)) + k.emit(s_mov_b32(s[46], 0)) + k.label('L2_00E8') + k.emit(s_cmp_lt_i32(s[48], s[52])) + k.emit(s_cbranch_scc1(3), target='L2_00FC') + k.emit(s_sub_i32(s[48], s[48], s[52])) + k.emit(s_add_i32(s[46], s[46], 32)) + k.emit(s_branch(65531), target='L2_00E8') + k.label('L2_00FC') + k.emit(s_sub_i32(s[54], s[54], s[46])) + k.emit(s_cmp_lt_i32(s[54], 32)) + k.emit(s_cbranch_scc1(3), target='L2_0114') + k.emit(s_lshr_b32(s[47], s[48], 5)) + k.emit(s_and_b32(s[52], s[48], 31)) + k.emit(s_branch(32), target='L2_0194') + k.label('L2_0114') + k.emit(v_cvt_f32_u32_e32(v[4], s[54])) + k.emit(s_sub_i32(s[47], 0, s[54])) + k.emit(v_rcp_iflag_f32_e32(v[4], v[4])) + k.emit(s_nop()) + k.emit(v_mul_f32_e32(v[4], LIT, v[4], 1333788670)) + k.emit(v_cvt_u32_f32_e32(v[4], v[4])) + k.emit(v_mul_lo_u32(v[5], s[47], v[4])) + k.emit(v_mul_hi_u32(v[5], v[4], v[5])) + k.emit(v_add_u32_e32(v[4], v[4], v[5])) + k.emit(v_mul_hi_u32(v[4], s[48], v[4])) + k.emit(v_mul_lo_u32(v[5], v[4], s[54])) + k.emit(v_sub_u32_e32(v[7], s[48], v[5])) + k.emit(v_add_u32_e32(v[6], 1, v[4])) + k.emit(v_cmp_le_u32_e32(s[54], v[7])) + k.emit(v_subrev_u32_e32(v[5], s[54], v[7])) + k.emit(s_nop()) + k.emit(v_cndmask_b32_e32(v[4], v[4], v[6])) + k.emit(v_cndmask_b32_e32(v[7], v[7], v[5])) + k.emit(v_add_u32_e32(v[5], 1, v[4])) + k.emit(v_cmp_le_u32_e32(s[54], v[7])) + k.emit(s_nop(1)) + k.emit(v_cndmask_b32_e32(v[7], v[4], v[5])) + k.emit(s_nop(3)) + k.emit(v_readfirstlane_b32_e32(v[47], v[7])) + k.emit(s_nop(3)) + k.emit(s_mul_i32(s[52], s[54], s[47])) + k.emit(s_sub_i32(s[52], s[48], s[52])) + k.label('L2_0194') + k.emit(s_add_i32(s[46], s[52], s[46])) + k.emit(s_mov_b32(s[6], -16)) + k.emit(s_mov_b32(s[10], -16)) + k.emit(s_mov_b32(s[18], -16)) + k.emit(s_mov_b32(s[14], -16)) + k.emit(s_mov_b32(s[7], LIT, 131072)) + k.emit(s_mov_b32(s[11], LIT, 131072)) + k.emit(s_mov_b32(s[19], LIT, 131072)) + k.emit(s_mov_b32(s[15], LIT, 131072)) + k.emit(s_and_b32(s[5], s[5], LIT, 65535)) + k.emit(s_and_b32(s[9], s[9], LIT, 65535)) + k.emit(s_and_b32(s[17], s[17], LIT, 65535)) + k.emit(s_and_b32(s[13], s[13], LIT, 65535)) + k.emit(s_or_b32(s[5], s[5], LIT, 262144)) + k.emit(s_or_b32(s[9], s[9], LIT, 262144)) + k.emit(s_or_b32(s[17], s[17], LIT, 262144)) + k.emit(s_or_b32(s[13], s[13], LIT, 262144)) + k.emit(s_cmp_gt_u32(s[57], 0)) + k.emit(s_cbranch_scc0(9), target='L2_0234') + k.emit(s_lshr_b32(s[58], s[45], s[57])) + k.emit(s_add_u32(s[58], s[58], LIT, 255)) + k.emit(s_lshr_b32(s[58], s[58], 8)) + k.emit(s_lshl_b32(s[58], s[58], 8)) + k.emit(s_mul_i32(s[53], s[58], s[56])) + k.emit(s_sub_i32(s[52], s[45], s[53])) + k.emit(s_cmp_lt_i32(s[52], s[58])) + k.emit(s_cselect_b32(s[45], s[52], s[58])) + k.label('L2_0234') + k.emit(s_lshr_b32(s[41], s[41], 1)) + k.emit(s_mul_i32(s[52], s[41], s[43])) + k.emit(s_mov_b32(s[14], s[52])) + k.emit(s_cmp_gt_u32(s[57], 0)) + k.emit(s_cbranch_scc0(5), target='L2_025C') + k.emit(s_mul_i32(s[53], s[58], s[56])) + k.emit(s_lshr_b32(s[52], s[53], 1)) + k.emit(s_add_u32(s[12], s[12], s[52])) + k.emit(s_addc_u32(s[13], s[13], 0)) + k.emit(s_sub_u32(s[14], s[14], s[52])) + k.label('L2_025C') + k.emit(s_lshr_b32(s[42], s[42], 1)) + k.emit(s_mul_i32(s[52], s[42], s[44])) + k.emit(s_mov_b32(s[18], s[52])) + k.emit(s_add_u32(s[52], s[43], 31)) + k.emit(s_lshr_b32(s[52], s[52], 5)) + k.emit(s_lshl_b32(s[52], s[52], 5)) + k.emit(s_mul_i32(s[53], s[52], s[36])) + k.emit(s_mov_b32(s[22], s[53])) + k.emit(s_mul_i32(s[53], s[44], s[37])) + k.emit(s_mov_b32(s[26], s[53])) + k.emit(s_mov_b32(s[23], LIT, 131072)) + k.emit(s_mov_b32(s[27], LIT, 131072)) + k.emit(s_and_b32(s[21], s[21], LIT, 65535)) + k.emit(s_and_b32(s[25], s[25], LIT, 65535)) + k.emit(s_or_b32(s[21], s[21], LIT, 262144)) + k.emit(s_or_b32(s[25], s[25], LIT, 262144)) + k.emit(v_lshrrev_b32_e32(v[4], 3)) + k.emit(v_lshrrev_b32_e32(v[5], 2, v[4])) + k.emit(v_lshlrev_b32_e32(v[5], 4, v[5])) + k.emit(v_and_b32_e32(v[4], 3, v[4])) + k.emit(v_lshrrev_b32_e32(v[6], 1, v[4])) + k.emit(v_lshlrev_b32_e32(v[6], 2, v[6])) + k.emit(v_add_u32_e32(v[5], v[5], v[6])) + k.emit(v_and_b32_e32(v[4], 1, v[4])) + k.emit(v_add_u32_e32(v[5], v[5], v[4])) + k.emit(v_mul_lo_u32(v[212], s[41], v[5])) + k.emit(v_and_b32_e32(v[4], 7)) + k.emit(v_lshlrev_b32_e32(v[4], 4, v[4])) + k.emit(v_add_u32_e32(v[212], v[212], v[4])) + k.emit(s_lshr_b32(s[52], s[49], 1)) + k.emit(s_mul_i32(s[52], s[52], 8)) + k.emit(s_and_b32(s[53], s[49], 1)) + k.emit(s_mul_i32(s[53], s[53], 2)) + k.emit(s_add_u32(s[52], s[52], s[53])) + k.emit(s_mul_i32(s[53], s[47], LIT, 256)) + k.emit(s_add_u32(s[52], s[52], s[53])) + k.emit(s_mul_i32(s[52], s[41], s[52])) + k.emit(v_add_u32_e32(v[212], s[52], v[212])) + k.emit(s_mul_i32(s[52], s[41], 32)) + for i in range(7): + k.emit(v_add_u32_e32(v[213 + i * 1], s[52], v[212 + i * 1])) + k.emit(s_mul_i32(s[59], LIT, s[49], 1056)) + k.emit(s_add_u32(s[59], LIT, s[59], 4096)) + k.emit(v_and_b32_e32(v[4], 15)) + k.emit(v_lshrrev_b32_e32(v[5], 3, v[4])) + k.emit(v_mul_i32_i24_e32(v[5], 2, v[5])) + k.emit(v_and_b32_e32(v[4], 3)) + k.emit(v_lshrrev_b32_e32(v[6], 1, v[4])) + k.emit(v_add_u32_e32(v[4], v[5], v[6])) + k.emit(v_mul_i32_i24_e32(v[220], LIT, v[4], 1056)) + k.emit(v_and_b32_e32(v[4], 7)) + k.emit(v_lshrrev_b32_e32(v[5], 2, v[4])) + k.emit(v_mul_i32_i24_e32(v[5], LIT, v[5], 256)) + k.emit(v_add_u32_e32(v[220], v[5], v[220])) + k.emit(v_and_b32_e32(v[4], 1)) + k.emit(v_mul_i32_i24_e32(v[6], LIT, v[4], 128)) + k.emit(v_add_u32_e32(v[220], v[6], v[220])) + k.emit(v_lshrrev_b32_e32(v[4], 4)) + k.emit(v_mul_i32_i24_e32(v[4], 16, v[4])) + k.emit(v_add_u32_e32(v[220], v[4], v[220])) + k.emit(s_mov_b32(s[52], LIT, 4096)) + k.emit(v_add_u32_e64(v[220], v[220], s[52])) + k.emit(v_add_u32_e32(v[221], LIT, v[220], 33792)) + k.emit(s_cmp_gt_u32(s[57], 0)) + k.emit(s_cbranch_scc0(4), target='L2_03C4') + k.emit(s_mul_i32(s[53], s[58], s[56])) + k.emit(s_add_u32(s[20], s[20], s[53])) + k.emit(s_addc_u32(s[21], s[21], 0)) + k.emit(s_sub_u32(s[22], s[22], s[53])) + k.label('L2_03C4') + k.emit(v_lshlrev_b32_e32(v[222], 2)) + k.emit(s_mul_i32(s[52], s[47], LIT, 256)) + k.emit(s_mul_i32(s[53], s[49], 32)) + k.emit(s_add_i32(s[52], s[53], s[52])) + k.emit(s_mul_i32(s[53], s[52], s[36])) + k.emit(v_add_u32_e32(v[222], s[53], v[222])) + k.emit(s_mul_i32(s[53], LIT, s[36], 128)) + k.emit(v_add_u32_e32(v[223], s[53], v[222])) + k.emit(s_mul_i32(s[60], s[49], LIT, 256)) + k.emit(s_add_i32(s[60], s[60], 0)) + k.emit(v_lshlrev_b32_e32(v[224], 2)) + k.emit(v_add_u32_e32(v[224], 0, v[224])) + k.emit(s_cmp_gt_u32(s[57], 0)) + k.emit(s_cbranch_scc0(6), target='L2_0420') + k.emit(s_mul_i32(s[53], s[58], s[56])) + k.emit(s_lshr_b32(s[52], s[53], 1)) + k.emit(s_mul_i32(s[52], s[52], 16)) + k.emit(s_add_u32(s[16], s[16], s[52])) + k.emit(s_addc_u32(s[17], s[17], 0)) + k.emit(s_sub_u32(s[18], s[18], s[52])) + k.label('L2_0420') + k.emit(v_lshlrev_b32_e32(v[225], 4)) + k.emit(s_mul_i32(s[52], s[46], LIT, 256)) + k.emit(s_mul_i32(s[53], s[49], 64)) + k.emit(s_add_u32(s[52], s[52], s[53])) + k.emit(s_mul_i32(s[52], s[52], s[42])) + k.emit(v_add_u32_e32(v[225], s[52], v[225])) + k.emit(s_mul_i32(s[52], 16, s[42])) + k.emit(v_add_u32_e32(v[226], s[52], v[225])) + k.emit(v_add_u32_e32(v[227], s[52], v[226])) + k.emit(v_add_u32_e32(v[228], s[52], v[227])) + for i in range(4): + k.emit(v_add_u32_e32(v[229 + i * 1], LIT, v[225 + i * 1], 1024)) + k.emit(s_cmp_gt_u32(s[57], 0)) + k.emit(s_cbranch_scc0(4), target='L2_0484') + k.emit(s_mul_i32(s[53], s[58], s[56])) + k.emit(s_add_u32(s[24], s[24], s[53])) + k.emit(s_addc_u32(s[25], s[25], 0)) + k.emit(s_sub_u32(s[26], s[26], s[53])) + k.label('L2_0484') + k.emit(v_lshlrev_b32_e32(v[233], 2)) + k.emit(s_mul_i32(s[52], s[46], LIT, 256)) + k.emit(s_mul_i32(s[53], s[49], 64)) + k.emit(s_add_i32(s[52], s[53], s[52])) + k.emit(s_mul_i32(s[53], s[52], s[37])) + k.emit(v_add_u32_e32(v[233], s[53], v[233])) + k.emit(s_mul_i32(s[52], 32, s[37])) + k.emit(v_add_u32_e32(v[234], s[52], v[233])) + k.emit(s_mov_b32(s[61], LIT, 128)) + k.emit(s_mov_b32(s[62], LIT, 2048)) + k.emit(s_mov_b32(s[63], LIT, 256)) + k.emit(s_mov_b32(s[64], LIT, 256)) + k.emit(s_add_u32(NULL, 0, s[59])) + for i in range(3): + k.emit(buffer_load_dwordx4(v[0:3], v[212 + i * 1], s[12:15], 0, 0, 1, 0, 0, 0, 0, 1)) + for j67 in range(8): + k.emit(v_accvgpr_write(v[0 + j67 * 1 + i * 8], 0)) + k.emit(s_add_u32(NULL, LIT, s[59], 4224 + i * 4224)) + k.emit(buffer_load_dwordx4(v[0:3], v[215], s[12:15], 0, 0, 1, 0, 0, 0, 0, 1)) + for i in range(8): + k.emit(v_accvgpr_write(v[24 + i * 1], 0)) + k.emit(s_add_u32(NULL, 0, s[60])) + k.emit(buffer_load_dword(v[0], v[222], s[20:23], 0, 0, 1, 0, 0, 0, 0, 1)) + for i in range(4): + for j68 in range(8): + k.emit(v_accvgpr_write(v[32 + j68 * 1 + i * 8], 0)) + k.emit(s_add_u32(NULL, LIT, s[59], 16896 + i * 4224)) + k.emit(buffer_load_dwordx4(v[0:3], v[216 + i * 1], s[12:15], 0, 0, 1, 0, 0, 0, 0, 1)) + for i in range(8): + k.emit(v_accvgpr_write(v[64 + i * 1], 0)) + k.emit(s_add_u32(NULL, LIT, s[60], 1024)) + k.emit(buffer_load_dword(v[0], v[223], s[20:23], 0, 0, 1, 0, 0, 0, 0, 1)) + for i in range(8): + k.emit(v_accvgpr_write(v[72 + i * 1], 0)) + for i in range(2): + k.emit(s_add_u32(s[12 + i * 8], s[61 + i * 2], s[12 + i * 8])) + k.emit(s_addc_u32(s[13 + i * 8], 0, s[13 + i * 8])) + k.emit(s_sub_u32(s[14 + i * 8], s[14 + i * 8], s[61 + i * 2])) + for i in range(8): + for j69 in range(8): + k.emit(v_accvgpr_write(v[80 + j69 * 1 + i * 8], 0)) + k.emit(buffer_load_dwordx4(v[136 + i * 4:139 + i * 4], v[225 + i * 1], s[16:19], 0, 0, 1)) + for i in range(8): + k.emit(v_accvgpr_write(v[144 + i * 1], 0)) + k.emit(s_add_u32(s[16], s[62], s[16])) + k.emit(s_addc_u32(s[17], 0, s[17])) + k.emit(s_sub_u32(s[18], s[18], s[62])) + for i in range(2): + k.emit(buffer_load_dword(v[208 + i * 1], v[233 + i * 1], s[24:27], 0, 0, 1)) + for j70 in range(8): + k.emit(v_accvgpr_write(v[152 + j70 * 1 + i * 8], 0)) + k.emit(s_add_u32(s[24], s[64], s[24])) + k.emit(s_addc_u32(s[25], 0, s[25])) + k.emit(s_sub_u32(s[26], s[26], s[64])) + for i in range(2): + for j71 in range(4): + k.emit(s_add_u32(NULL, LIT, s[59], 33792 + j71 * 4224 + i * 16896)) + k.emit(buffer_load_dwordx4(v[0:3], v[212 + j71 * 1 + i * 4], s[12:15], 0, 0, 1, 0, 0, 0, 0, 1)) + k.emit(v_accvgpr_write(v[168 + j71 * 8 + i * 40], 0)) + k.emit(v_accvgpr_write(v[169 + j71 * 8 + i * 40], 0)) + k.emit(v_accvgpr_write(v[170 + j71 * 8 + i * 40], 0)) + k.emit(v_accvgpr_write(v[171 + j71 * 8 + i * 40], 0)) + k.emit(v_accvgpr_write(v[172 + j71 * 8 + i * 40], 0)) + k.emit(v_accvgpr_write(v[173 + j71 * 8 + i * 40], 0)) + k.emit(v_accvgpr_write(v[174 + j71 * 8 + i * 40], 0)) + k.emit(v_accvgpr_write(v[175 + j71 * 8 + i * 40], 0)) + k.emit(s_add_u32(NULL, LIT, s[60], 2048 + i * 1024)) + k.emit(buffer_load_dword(v[0], v[222 + i * 1], s[20:23], 0, 0, 1, 0, 0, 0, 0, 1)) + for j72 in range(8): + k.emit(v_accvgpr_write(v[200 + j72 * 1 + i * 40], 0)) + for i in range(2): + k.emit(s_add_u32(s[12 + i * 8], s[61 + i * 2], s[12 + i * 8])) + k.emit(s_addc_u32(s[13 + i * 8], 0, s[13 + i * 8])) + k.emit(s_sub_u32(s[14 + i * 8], s[14 + i * 8], s[61 + i * 2])) + for i in range(8): + k.emit(v_accvgpr_write(v[248 + i * 1], 0)) + k.emit(s_waitcnt(20345)) + k.emit(s_barrier()) + k.emit(ds_read_b128(v[8:11], v[220])) + k.emit(ds_read_b128(v[40:43], v[220], v[0], v[0], 0, 64)) + for i in range(2): + k.emit(ds_read_b128(v[12 + i * 44:15 + i * 44], v[220], v[0], v[0], 0, 0 + i * 64, 2 + i * 31)) + k.emit(ds_read_b128(v[44 + i * -16:47 + i * -16], v[220], v[0], v[0], 0, 64 + i * -64, 2 + i * 33)) + k.emit(ds_read_b128(v[16 + i * 44:19 + i * 44], v[220], v[0], v[0], 0, 128 + i * -64, 16 + i * 19)) + k.emit(ds_read_b128(v[48 + i * -16:51 + i * -16], v[220], v[0], v[0], 0, 192 + i * -64, 16 + i * 33)) + k.emit(ds_read_b128(v[20 + i * 44:23 + i * 44], v[220], v[0], v[0], 0, 128 + i * 64, 18 + i * 31)) + k.emit(ds_read_b128(v[52 + i * -16:55 + i * -16], v[220], v[0], v[0], 0, 192 + i * -64, 18 + i * 33)) + k.emit(ds_read_b128(v[24 + i * 44:27 + i * 44], v[220], v[0], v[0], 0, 0 + i * 192, 33 + i * 18)) + k.emit(ds_read_b32(v[200], v[224])) + k.emit(ds_read_b32(v[201], v[224], v[0], v[0], 0, 0, 1)) + k.emit(ds_read_b32(v[202], v[224], v[0], v[0], 0, 0, 2)) + k.emit(ds_read_b32(v[203], v[224], v[0], v[0], 0, 0, 3)) + k.emit(s_lshl_b32(s[40], s[40], 1)) + k.emit(s_mul_i32(s[52], s[47], LIT, 256)) + k.emit(s_mul_hi_u32(s[53], s[52], s[40])) + k.emit(s_add_u32(s[5], s[5], s[53])) + k.emit(s_mul_i32(s[53], s[52], s[40])) + k.emit(s_add_u32(s[4], s[4], s[53])) + k.emit(s_addc_u32(s[5], 0, s[5])) + k.emit(s_sub_i32(s[52], s[43], s[52])) + k.emit(s_mul_i32(s[52], s[52], s[40])) + k.emit(s_mov_b32(s[6], s[52])) + k.emit(v_and_b32_e64(v[235], v[0], 15)) + k.emit(v_mul_lo_u32(v[235], v[235], s[40])) + k.emit(v_lshrrev_b32_e32(v[4], 5)) + k.emit(v_mul_i32_i24_e32(v[4], 16, v[4])) + k.emit(v_add_u32_e32(v[235], v[4], v[235])) + k.emit(v_lshrrev_b32_e32(v[4], 4)) + k.emit(v_and_b32_e32(v[4], 1, v[4])) + k.emit(v_mul_i32_i24_e32(v[4], 32, v[4])) + k.emit(v_add_u32_e32(v[235], v[4], v[235])) + k.emit(s_mul_i32(s[52], s[46], LIT, 256)) + k.emit(s_mul_i32(s[53], s[49], 64)) + k.emit(s_add_i32(s[52], s[52], s[53])) + k.emit(s_lshl_b32(s[52], s[52], 1)) + k.emit(v_add_u32_e32(v[235], s[52], v[235])) + k.emit(s_mul_i32(s[53], s[40], 16)) + for i in range(15): + k.emit(v_add_u32_e64(v[236 + i * 1], v[235 + i * 1], s[53])) + k.emit(s_mov_b32(s[50], 0)) + k.emit(s_mov_b32(s[51], s[45])) + for i in range(2): + k.emit(s_cmp_lt_u32(LIT, s[51], 512 + i * -256)) + k.emit(s_cselect_b32(s[61 + i * 1], s[61 + i * 1], 0)) + k.emit(s_cselect_b32(s[63 + i * 1], s[63 + i * 1], 0)) + k.emit(s_cmp_lt_i32(s[49], 2)) + k.emit(s_cbranch_scc0(1367), target='L2_25B8') + k.label('L2_105C') + k.emit(s_waitcnt(122)) + k.emit(v_mfma_fp4(v[0:3], v[136:139], v[8:11], 0, 0, v[208], v[200])) + k.emit(s_barrier()) + k.emit(s_nop()) + k.emit(s_nop()) + for i in range(2): + k.emit(v_mfma_fp4(v[4 + i * 64:7 + i * 64], v[136 + i * 8:139 + i * 8], v[12:15], 2, 0, v[208 + i * 1], v[200])) + k.emit(buffer_load_dwordx4(v[168 + i * 16:171 + i * 16], v[225 + i * 4], s[16:19], 0, 0, 1)) + k.emit(v_mfma_fp4(v[32 + i * 64:35 + i * 64], v[140 + i * 8:143 + i * 8], v[8:11], 1, 0, v[208 + i * 1], v[200])) + k.emit(ds_read_b128(v[72 + i * 8:75 + i * 8], v[220], v[0], v[0], 0, 0 + i * 128, 66 + i * 16)) + k.emit(v_mfma_fp4(v[36 + i * 64:39 + i * 64], v[140 + i * 8:143 + i * 8], v[12:15], 3, 0, v[208 + i * 1], v[200])) + k.emit(v_mfma_fp4(v[8 + i * 64:11 + i * 64], v[136 + i * 8:139 + i * 8], v[16:19], 0, 0, v[208 + i * 1], v[201])) + for j73 in range(2): + k.emit(v_mfma_fp4(v[12 + j73 * 8 + i * 64:15 + j73 * 8 + i * 64], v[136 + i * 8:139 + i * 8], v[20 + j73 * 8:23 + j73 * 8], 2, 0, v[208 + i * 1], v[201 + j73 * 1])) + k.emit(buffer_load_dwordx4(v[172 + j73 * 4 + i * 16:175 + j73 * 4 + i * 16], v[226 + j73 * 1 + i * 4], s[16:19], 0, 0, 1)) + k.emit(v_mfma_fp4(v[40 + j73 * 8 + i * 64:43 + j73 * 8 + i * 64], v[140 + i * 8:143 + i * 8], v[16 + j73 * 8:19 + j73 * 8], 1, 0, v[208 + i * 1], v[201 + j73 * 1])) + k.emit(ds_read_b128(v[104 + j73 * -28 + i * 8:107 + j73 * -28 + i * 8], v[220], v[0], v[0], 0, 64 + j73 * -64 + i * 128, 66 + j73 * 2 + i * 16)) + k.emit(v_mfma_fp4(v[44 + j73 * 8 + i * 64:47 + j73 * 8 + i * 64], v[140 + i * 8:143 + i * 8], v[20 + j73 * 8:23 + j73 * 8], 3, 0, v[208 + i * 1], v[201 + j73 * 1])) + k.emit(v_mfma_fp4(v[16 + j73 * 8 + i * 64:19 + j73 * 8 + i * 64], v[136 + i * 8:139 + i * 8], v[24 + j73 * 8:27 + j73 * 8], 0, 0, v[208 + i * 1], v[202 + j73 * 1])) + k.emit(v_mfma_fp4(v[28 + i * 64:31 + i * 64], v[136 + i * 8:139 + i * 8], v[36:39], 2, 0, v[208 + i * 1], v[203])) + k.emit(buffer_load_dwordx4(v[180 + i * 16:183 + i * 16], v[228 + i * 4], s[16:19], 0, 0, 1)) + k.emit(v_mfma_fp4(v[56 + i * 64:59 + i * 64], v[140 + i * 8:143 + i * 8], v[32:35], 1, 0, v[208 + i * 1], v[203])) + k.emit(ds_read_b128(v[108 + i * 8:111 + i * 8], v[220], v[0], v[0], 0, 64 + i * 128, 68 + i * 16)) + k.emit(v_mfma_fp4(v[60 + i * 64:63 + i * 64], v[140 + i * 8:143 + i * 8], v[36:39], 3, 0, v[208 + i * 1], v[203])) + k.emit(v_mfma_fp4(v[64 + i * -64:67 + i * -64], v[144 + i * 8:147 + i * 8], v[8 + i * 32:11 + i * 32], 0, 0 + i * 3, v[209 + i * -1], v[200])) + for i in range(2): + k.emit(v_mfma_fp4(v[4 + i * 8:7 + i * 8], v[152:155], v[44 + i * 8:47 + i * 8], 2, 3, v[208], v[200 + i * 1])) + k.emit(buffer_load_dword(v[210 + i * 1], v[233 + i * 1], s[24:27], 0, 0, 1)) + k.emit(v_mfma_fp4(v[32 + i * 8:35 + i * 8], v[156:159], v[40 + i * 8:43 + i * 8], 1, 3, v[208], v[200 + i * 1])) + k.emit(ds_read_b128(v[88 + i * 32:91 + i * 32], v[220], v[0], v[0], 0, 0 + i * 64, 99)) + k.emit(v_mfma_fp4(v[36 + i * 8:39 + i * 8], v[156:159], v[44 + i * 8:47 + i * 8], 3, 3, v[208], v[200 + i * 1])) + k.emit(v_mfma_fp4(v[8 + i * 8:11 + i * 8], v[152:155], v[48 + i * 8:51 + i * 8], 0, 3, v[208], v[201 + i * 1])) + k.emit(s_add_u32(s[53], LIT, s[50], 512)) + k.emit(v_mfma_fp4(v[20:23], v[152:155], v[60:63], 2, 3, v[208], v[202])) + k.emit(ds_read_b128(v[92:95], v[220], v[0], v[0], 0, 0, 101)) + k.emit(v_mfma_fp4(v[48:51], v[156:159], v[56:59], 1, 3, v[208], v[202])) + k.emit(s_cmp_lt_u32(s[53], s[51])) + for i in range(2): + k.emit(v_mfma_fp4(v[52 + i * -24:55 + i * -24], v[156 + i * -4:159 + i * -4], v[60 + i * 8:63 + i * 8], 3 + i * -1, 3, v[208], v[202 + i * 1])) + k.emit(ds_read_b128(v[124 + i * -28:127 + i * -28], v[220], v[0], v[0], 0, 64 + i * 64, 101 + i * 14)) + k.emit(v_mfma_fp4(v[24 + i * 32:27 + i * 32], v[152 + i * 4:155 + i * 4], v[64:67], 0 + i * 1, 3, v[208], v[203])) + k.emit(s_cselect_b32(s[62 + i * 2], s[62 + i * 2], 0)) + k.emit(v_mfma_fp4(v[60:63], v[156:159], v[68:71], 3, 3, v[208], v[203])) + k.emit(ds_read_b128(v[128:131], v[220], v[0], v[0], 0, 192, 115)) + k.emit(v_mfma_fp4(v[64:67], v[160:163], v[40:43], 0, 3, v[209], v[200])) + k.emit(s_add_u32(s[16], s[62], s[16])) + k.emit(v_mfma_fp4(v[68:71], v[160:163], v[44:47], 2, 3, v[209], v[200])) + k.emit(ds_read_b128(v[100:103], v[220], v[0], v[0], 0, 128, 117)) + k.emit(v_mfma_fp4(v[96:99], v[164:167], v[40:43], 1, 3, v[209], v[200])) + k.emit(s_addc_u32(s[17], 0, s[17])) + k.emit(v_mfma_fp4(v[100:103], v[164:167], v[44:47], 3, 3, v[209], v[200])) + k.emit(ds_read_b128(v[132:135], v[220], v[0], v[0], 0, 192, 117)) + k.emit(v_mfma_fp4(v[72:75], v[160:163], v[48:51], 0, 3, v[209], v[201])) + k.emit(s_sub_u32(s[18], s[18], s[62])) + k.emit(v_mfma_fp4(v[76:79], v[160:163], v[52:55], 2, 3, v[209], v[201])) + k.emit(ds_read_b32(v[204], v[224], v[0], v[0], 0, 0, 4)) + k.emit(v_mfma_fp4(v[104:107], v[164:167], v[48:51], 1, 3, v[209], v[201])) + k.emit(s_add_u32(s[24], s[64], s[24])) + k.emit(v_mfma_fp4(v[108:111], v[164:167], v[52:55], 3, 3, v[209], v[201])) + k.emit(ds_read_b32(v[205], v[224], v[0], v[0], 0, 0, 5)) + k.emit(v_mfma_fp4(v[80:83], v[160:163], v[56:59], 0, 3, v[209], v[202])) + k.emit(s_addc_u32(s[25], 0, s[25])) + k.emit(v_mfma_fp4(v[84:87], v[160:163], v[60:63], 2, 3, v[209], v[202])) + k.emit(ds_read_b32(v[206], v[224], v[0], v[0], 0, 0, 6)) + k.emit(v_mfma_fp4(v[112:115], v[164:167], v[56:59], 1, 3, v[209], v[202])) + k.emit(s_sub_u32(s[26], s[26], s[64])) + k.emit(v_mfma_fp4(v[116:119], v[164:167], v[60:63], 3, 3, v[209], v[202])) + k.emit(ds_read_b32(v[207], v[224], v[0], v[0], 0, 0, 7)) + k.emit(v_mfma_fp4(v[88:91], v[160:163], v[64:67], 0, 3, v[209], v[203])) + k.emit(v_mfma_fp4(v[92:95], v[160:163], v[68:71], 2, 3, v[209], v[203])) + k.emit(v_mfma_fp4(v[120:123], v[164:167], v[64:67], 1, 3, v[209], v[203])) + k.emit(v_mfma_fp4(v[124:127], v[164:167], v[68:71], 3, 3, v[209], v[203])) + k.emit(s_waitcnt(127)) + k.emit(v_mfma_fp4(v[128:131], v[136:139], v[72:75], 0, 0, v[208], v[204])) + k.emit(s_barrier()) + k.emit(s_nop()) + k.emit(v_mfma_fp4(v[132:135], v[136:139], v[76:79], 2, 0, v[208], v[204])) + k.emit(s_add_u32(NULL, 0, s[59])) + k.emit(v_mfma_fp4(v[160:163], v[140:143], v[72:75], 1, 0, v[208], v[204])) + k.emit(ds_read_b128(v[8:11], v[221])) + k.emit(v_mfma_fp4(v[164:167], v[140:143], v[76:79], 3, 0, v[208], v[204])) + k.emit(buffer_load_dwordx4(v[0:3], v[212], s[12:15], 0, 0, 1, 0, 0, 0, 0, 1)) + k.emit(v_mfma_fp4(v[136:139], v[136:139], v[80:83], 0, 0, v[208], v[205])) + k.emit(ds_read_b128(v[40:43], v[221], v[0], v[0], 0, 64)) + k.emit(v_mfma_fp4(v[140:143], v[136:139], v[84:87], 2, 0, v[208], v[205])) + k.emit(s_add_u32(NULL, LIT, s[59], 4224)) + for i in range(2): + k.emit(v_mfma_fp4(v[168 + i * 8:171 + i * 8], v[140:143], v[80 + i * 8:83 + i * 8], 1, 0, v[208], v[205 + i * 1])) + k.emit(ds_read_b128(v[12 + i * 4:15 + i * 4], v[221], v[0], v[0], 0, 0 + i * 128, 2 + i * 14)) + k.emit(v_mfma_fp4(v[172 + i * 8:175 + i * 8], v[140:143], v[84 + i * 8:87 + i * 8], 3, 0, v[208], v[205 + i * 1])) + k.emit(buffer_load_dwordx4(v[0:3], v[213 + i * 1], s[12:15], 0, 0, 1, 0, 0, 0, 0, 1)) + k.emit(v_mfma_fp4(v[144 + i * 8:147 + i * 8], v[136:139], v[88 + i * 8:91 + i * 8], 0, 0, v[208], v[206 + i * 1])) + k.emit(ds_read_b128(v[44 + i * 4:47 + i * 4], v[221], v[0], v[0], 0, 64 + i * 128, 2 + i * 14)) + k.emit(v_mfma_fp4(v[148 + i * 8:151 + i * 8], v[136:139], v[92 + i * 8:95 + i * 8], 2, 0, v[208], v[206 + i * 1])) + k.emit(s_add_u32(NULL, LIT, s[59], 8448 + i * 4224)) + k.emit(v_mfma_fp4(v[184:187], v[140:143], v[96:99], 1, 0, v[208], v[207])) + k.emit(ds_read_b128(v[20:23], v[221], v[0], v[0], 0, 128, 18)) + k.emit(v_mfma_fp4(v[188:191], v[140:143], v[100:103], 3, 0, v[208], v[207])) + k.emit(buffer_load_dwordx4(v[0:3], v[215], s[12:15], 0, 0, 1, 0, 0, 0, 0, 1)) + k.emit(v_mfma_fp4(v[192:195], v[144:147], v[72:75], 0, 0, v[209], v[204])) + k.emit(ds_read_b128(v[52:55], v[221], v[0], v[0], 0, 192, 18)) + k.emit(v_mfma_fp4(v[196:199], v[144:147], v[76:79], 2, 0, v[209], v[204])) + k.emit(s_add_u32(NULL, 0, s[60])) + k.emit(v_mfma_fp4(v[224:227], v[148:151], v[72:75], 1, 0, v[209], v[204])) + k.emit(ds_read_b128(v[24:27], v[221], v[0], v[0], 0, 0, 33)) + k.emit(v_mfma_fp4(v[228:231], v[148:151], v[76:79], 3, 0, v[209], v[204])) + k.emit(buffer_load_dword(v[0], v[222], s[20:23], 0, 0, 1, 0, 0, 0, 0, 1)) + k.emit(v_mfma_fp4(v[200:203], v[144:147], v[80:83], 0, 0, v[209], v[205])) + k.emit(ds_read_b128(v[56:59], v[221], v[0], v[0], 0, 64, 33)) + k.emit(v_mfma_fp4(v[204:207], v[144:147], v[84:87], 2, 0, v[209], v[205])) + k.emit(s_add_u32(NULL, LIT, s[59], 16896)) + k.emit(v_mfma_fp4(v[232:235], v[148:151], v[80:83], 1, 0, v[209], v[205])) + k.emit(ds_read_b128(v[28:31], v[221], v[0], v[0], 0, 0, 35)) + for i in range(2): + k.emit(v_mfma_fp4(v[236 + i * 8:239 + i * 8], v[148:151], v[84 + i * 8:87 + i * 8], 3, 0, v[209], v[205 + i * 1])) + k.emit(buffer_load_dwordx4(v[0:3], v[216 + i * 1], s[12:15], 0, 0, 1, 0, 0, 0, 0, 1)) + k.emit(v_mfma_fp4(v[208 + i * 8:211 + i * 8], v[144:147], v[88 + i * 8:91 + i * 8], 0, 0, v[209], v[206 + i * 1])) + k.emit(ds_read_b128(v[60 + i * 4:63 + i * 4], v[221], v[0], v[0], 0, 64 + i * 128, 35 + i * 14)) + k.emit(v_mfma_fp4(v[212 + i * 8:215 + i * 8], v[144:147], v[92 + i * 8:95 + i * 8], 2, 0, v[209], v[206 + i * 1])) + k.emit(s_add_u32(NULL, LIT, s[59], 21120 + i * 4224)) + k.emit(v_mfma_fp4(v[240 + i * 8:243 + i * 8], v[148:151], v[88 + i * 8:91 + i * 8], 1, 0, v[209], v[206 + i * 1])) + k.emit(ds_read_b128(v[32 + i * 4:35 + i * 4], v[221], v[0], v[0], 0, 128, 49 + i * 2)) + k.emit(v_mfma_fp4(v[252:255], v[148:151], v[100:103], 3, 0, v[209], v[207])) + k.emit(buffer_load_dwordx4(v[0:3], v[218], s[12:15], 0, 0, 1, 0, 0, 0, 0, 1)) + k.emit(v_mfma_fp4(v[128:131], v[152:155], v[104:107], 0, 3, v[208], v[204])) + k.emit(ds_read_b128(v[68:71], v[221], v[0], v[0], 0, 192, 51)) + k.emit(v_mfma_fp4(v[132:135], v[152:155], v[108:111], 2, 3, v[208], v[204])) + k.emit(s_add_u32(NULL, LIT, s[59], 29568)) + k.emit(v_mfma_fp4(v[160:163], v[156:159], v[104:107], 1, 3, v[208], v[204])) + k.emit(ds_read_b32(v[200], v[224], v[0], v[0], 0, 0, 8)) + k.emit(v_mfma_fp4(v[164:167], v[156:159], v[108:111], 3, 3, v[208], v[204])) + k.emit(buffer_load_dwordx4(v[0:3], v[219], s[12:15], 0, 0, 1, 0, 0, 0, 0, 1)) + k.emit(v_mfma_fp4(v[136:139], v[152:155], v[112:115], 0, 3, v[208], v[205])) + k.emit(ds_read_b32(v[201], v[224], v[0], v[0], 0, 0, 9)) + k.emit(v_mfma_fp4(v[140:143], v[152:155], v[116:119], 2, 3, v[208], v[205])) + k.emit(s_add_u32(NULL, LIT, s[60], 1024)) + k.emit(v_mfma_fp4(v[168:171], v[156:159], v[112:115], 1, 3, v[208], v[205])) + k.emit(ds_read_b32(v[202], v[224], v[0], v[0], 0, 0, 10)) + k.emit(v_mfma_fp4(v[172:175], v[156:159], v[116:119], 3, 3, v[208], v[205])) + k.emit(buffer_load_dword(v[0], v[223], s[20:23], 0, 0, 1, 0, 0, 0, 0, 1)) + k.emit(v_mfma_fp4(v[144:147], v[152:155], v[120:123], 0, 3, v[208], v[206])) + k.emit(ds_read_b32(v[203], v[224], v[0], v[0], 0, 0, 11)) + k.emit(v_mfma_fp4(v[148:151], v[152:155], v[124:127], 2, 3, v[208], v[206])) + k.emit(s_add_u32(s[52], LIT, s[50], 768)) + k.emit(v_mfma_fp4(v[176:179], v[156:159], v[120:123], 1, 3, v[208], v[206])) + k.emit(v_mfma_fp4(v[180:183], v[156:159], v[124:127], 3, 3, v[208], v[206])) + k.emit(s_cmp_lt_u32(s[52], s[51])) + for i in range(2): + k.emit(v_mfma_fp4(v[152 + i * 32:155 + i * 32], v[152 + i * 4:155 + i * 4], v[128:131], 0 + i * 1, 3, v[208], v[207])) + k.emit(v_mfma_fp4(v[156 + i * 32:159 + i * 32], v[152 + i * 4:155 + i * 4], v[132:135], 2 + i * 1, 3, v[208], v[207])) + k.emit(s_cselect_b32(s[61 + i * 2], s[61 + i * 2], 0)) + for i in range(2): + k.emit(v_mfma_fp4(v[192 + i * 40:195 + i * 40], v[160 + i * 4:163 + i * 4], v[104 + i * 8:107 + i * 8], 0 + i * 1, 3, v[209], v[204 + i * 1])) + k.emit(v_mfma_fp4(v[196 + i * 40:199 + i * 40], v[160 + i * 4:163 + i * 4], v[108 + i * 8:111 + i * 8], 2 + i * 1, 3, v[209], v[204 + i * 1])) + k.emit(s_add_u32(s[12 + i * 8], s[61 + i * 2], s[12 + i * 8])) + k.emit(v_mfma_fp4(v[224 + i * -16:227 + i * -16], v[164 + i * -4:167 + i * -4], v[104 + i * 16:107 + i * 16], 1 + i * -1, 3, v[209], v[204 + i * 2])) + k.emit(v_mfma_fp4(v[228 + i * -16:231 + i * -16], v[164 + i * -4:167 + i * -4], v[108 + i * 16:111 + i * 16], 3 + i * -1, 3, v[209], v[204 + i * 2])) + k.emit(s_addc_u32(s[13 + i * 8], 0, s[13 + i * 8])) + k.emit(v_mfma_fp4(v[200 + i * 40:203 + i * 40], v[160 + i * 4:163 + i * 4], v[112 + i * 8:115 + i * 8], 0 + i * 1, 3, v[209], v[205 + i * 1])) + k.emit(v_mfma_fp4(v[204 + i * 40:207 + i * 40], v[160 + i * 4:163 + i * 4], v[116 + i * 8:119 + i * 8], 2 + i * 1, 3, v[209], v[205 + i * 1])) + k.emit(s_sub_u32(s[14 + i * 8], s[14 + i * 8], s[61 + i * 2])) + k.emit(v_mfma_fp4(v[216:219], v[160:163], v[128:131], 0, 3, v[209], v[207])) + k.emit(s_addk_i32(s[50], 256)) + k.emit(v_mfma_fp4(v[220:223], v[160:163], v[132:135], 2, 3, v[209], v[207])) + k.emit(s_cmp_lt_i32(s[50], s[51])) + k.emit(v_mfma_fp4(v[248:251], v[164:167], v[128:131], 1, 3, v[209], v[207])) + k.emit(v_mfma_fp4(v[252:255], v[164:167], v[132:135], 3, 3, v[209], v[207])) + k.emit(s_cbranch_scc0(2051), target='L2_3B10') + k.emit(s_waitcnt(122)) + k.emit(v_mfma_fp4(v[0:3], v[168:171], v[8:11], 0, 0, v[210], v[200])) + k.emit(s_barrier()) + k.emit(s_nop()) + k.emit(s_nop()) + for i in range(2): + k.emit(v_mfma_fp4(v[4 + i * 64:7 + i * 64], v[168 + i * 8:171 + i * 8], v[12:15], 2, 0, v[210 + i * 1], v[200])) + k.emit(buffer_load_dwordx4(v[136 + i * 16:139 + i * 16], v[225 + i * 4], s[16:19], 0, 0, 1)) + k.emit(v_mfma_fp4(v[32 + i * 64:35 + i * 64], v[172 + i * 8:175 + i * 8], v[8:11], 1, 0, v[210 + i * 1], v[200])) + k.emit(ds_read_b128(v[72 + i * 8:75 + i * 8], v[221], v[0], v[0], 0, 0 + i * 128, 66 + i * 16)) + k.emit(v_mfma_fp4(v[36 + i * 64:39 + i * 64], v[172 + i * 8:175 + i * 8], v[12:15], 3, 0, v[210 + i * 1], v[200])) + k.emit(v_mfma_fp4(v[8 + i * 64:11 + i * 64], v[168 + i * 8:171 + i * 8], v[16:19], 0, 0, v[210 + i * 1], v[201])) + for j74 in range(2): + k.emit(v_mfma_fp4(v[12 + j74 * 8 + i * 64:15 + j74 * 8 + i * 64], v[168 + i * 8:171 + i * 8], v[20 + j74 * 8:23 + j74 * 8], 2, 0, v[210 + i * 1], v[201 + j74 * 1])) + k.emit(buffer_load_dwordx4(v[140 + j74 * 4 + i * 16:143 + j74 * 4 + i * 16], v[226 + j74 * 1 + i * 4], s[16:19], 0, 0, 1)) + k.emit(v_mfma_fp4(v[40 + j74 * 8 + i * 64:43 + j74 * 8 + i * 64], v[172 + i * 8:175 + i * 8], v[16 + j74 * 8:19 + j74 * 8], 1, 0, v[210 + i * 1], v[201 + j74 * 1])) + k.emit(ds_read_b128(v[104 + j74 * -28 + i * 8:107 + j74 * -28 + i * 8], v[221], v[0], v[0], 0, 64 + j74 * -64 + i * 128, 66 + j74 * 2 + i * 16)) + k.emit(v_mfma_fp4(v[44 + j74 * 8 + i * 64:47 + j74 * 8 + i * 64], v[172 + i * 8:175 + i * 8], v[20 + j74 * 8:23 + j74 * 8], 3, 0, v[210 + i * 1], v[201 + j74 * 1])) + k.emit(v_mfma_fp4(v[16 + j74 * 8 + i * 64:19 + j74 * 8 + i * 64], v[168 + i * 8:171 + i * 8], v[24 + j74 * 8:27 + j74 * 8], 0, 0, v[210 + i * 1], v[202 + j74 * 1])) + k.emit(v_mfma_fp4(v[28 + i * 64:31 + i * 64], v[168 + i * 8:171 + i * 8], v[36:39], 2, 0, v[210 + i * 1], v[203])) + k.emit(buffer_load_dwordx4(v[148 + i * 16:151 + i * 16], v[228 + i * 4], s[16:19], 0, 0, 1)) + k.emit(v_mfma_fp4(v[56 + i * 64:59 + i * 64], v[172 + i * 8:175 + i * 8], v[32:35], 1, 0, v[210 + i * 1], v[203])) + k.emit(ds_read_b128(v[108 + i * 8:111 + i * 8], v[221], v[0], v[0], 0, 64 + i * 128, 68 + i * 16)) + k.emit(v_mfma_fp4(v[60 + i * 64:63 + i * 64], v[172 + i * 8:175 + i * 8], v[36:39], 3, 0, v[210 + i * 1], v[203])) + k.emit(v_mfma_fp4(v[64 + i * -64:67 + i * -64], v[176 + i * 8:179 + i * 8], v[8 + i * 32:11 + i * 32], 0, 0 + i * 3, v[211 + i * -1], v[200])) + for i in range(2): + k.emit(v_mfma_fp4(v[4 + i * 8:7 + i * 8], v[184:187], v[44 + i * 8:47 + i * 8], 2, 3, v[210], v[200 + i * 1])) + k.emit(buffer_load_dword(v[208 + i * 1], v[233 + i * 1], s[24:27], 0, 0, 1)) + k.emit(v_mfma_fp4(v[32 + i * 8:35 + i * 8], v[188:191], v[40 + i * 8:43 + i * 8], 1, 3, v[210], v[200 + i * 1])) + k.emit(ds_read_b128(v[88 + i * 32:91 + i * 32], v[221], v[0], v[0], 0, 0 + i * 64, 99)) + k.emit(v_mfma_fp4(v[36 + i * 8:39 + i * 8], v[188:191], v[44 + i * 8:47 + i * 8], 3, 3, v[210], v[200 + i * 1])) + k.emit(v_mfma_fp4(v[8 + i * 8:11 + i * 8], v[184:187], v[48 + i * 8:51 + i * 8], 0, 3, v[210], v[201 + i * 1])) + k.emit(s_add_u32(s[53], LIT, s[50], 512)) + k.emit(v_mfma_fp4(v[20:23], v[184:187], v[60:63], 2, 3, v[210], v[202])) + k.emit(ds_read_b128(v[92:95], v[221], v[0], v[0], 0, 0, 101)) + k.emit(v_mfma_fp4(v[48:51], v[188:191], v[56:59], 1, 3, v[210], v[202])) + k.emit(s_cmp_lt_u32(s[53], s[51])) + for i in range(2): + k.emit(v_mfma_fp4(v[52 + i * -24:55 + i * -24], v[188 + i * -4:191 + i * -4], v[60 + i * 8:63 + i * 8], 3 + i * -1, 3, v[210], v[202 + i * 1])) + k.emit(ds_read_b128(v[124 + i * -28:127 + i * -28], v[221], v[0], v[0], 0, 64 + i * 64, 101 + i * 14)) + k.emit(v_mfma_fp4(v[24 + i * 32:27 + i * 32], v[184 + i * 4:187 + i * 4], v[64:67], 0 + i * 1, 3, v[210], v[203])) + k.emit(s_cselect_b32(s[62 + i * 2], s[62 + i * 2], 0)) + k.emit(v_mfma_fp4(v[60:63], v[188:191], v[68:71], 3, 3, v[210], v[203])) + k.emit(ds_read_b128(v[128:131], v[221], v[0], v[0], 0, 192, 115)) + k.emit(v_mfma_fp4(v[64:67], v[192:195], v[40:43], 0, 3, v[211], v[200])) + k.emit(s_add_u32(s[16], s[62], s[16])) + k.emit(v_mfma_fp4(v[68:71], v[192:195], v[44:47], 2, 3, v[211], v[200])) + k.emit(ds_read_b128(v[100:103], v[221], v[0], v[0], 0, 128, 117)) + k.emit(v_mfma_fp4(v[96:99], v[196:199], v[40:43], 1, 3, v[211], v[200])) + k.emit(s_addc_u32(s[17], 0, s[17])) + k.emit(v_mfma_fp4(v[100:103], v[196:199], v[44:47], 3, 3, v[211], v[200])) + k.emit(ds_read_b128(v[132:135], v[221], v[0], v[0], 0, 192, 117)) + k.emit(v_mfma_fp4(v[72:75], v[192:195], v[48:51], 0, 3, v[211], v[201])) + k.emit(s_sub_u32(s[18], s[18], s[62])) + k.emit(v_mfma_fp4(v[76:79], v[192:195], v[52:55], 2, 3, v[211], v[201])) + k.emit(ds_read_b32(v[204], v[224], v[0], v[0], 0, 0, 12)) + k.emit(v_mfma_fp4(v[104:107], v[196:199], v[48:51], 1, 3, v[211], v[201])) + k.emit(s_add_u32(s[24], s[64], s[24])) + k.emit(v_mfma_fp4(v[108:111], v[196:199], v[52:55], 3, 3, v[211], v[201])) + k.emit(ds_read_b32(v[205], v[224], v[0], v[0], 0, 0, 13)) + k.emit(v_mfma_fp4(v[80:83], v[192:195], v[56:59], 0, 3, v[211], v[202])) + k.emit(s_addc_u32(s[25], 0, s[25])) + k.emit(v_mfma_fp4(v[84:87], v[192:195], v[60:63], 2, 3, v[211], v[202])) + k.emit(ds_read_b32(v[206], v[224], v[0], v[0], 0, 0, 14)) + k.emit(v_mfma_fp4(v[112:115], v[196:199], v[56:59], 1, 3, v[211], v[202])) + k.emit(s_sub_u32(s[26], s[26], s[64])) + k.emit(v_mfma_fp4(v[116:119], v[196:199], v[60:63], 3, 3, v[211], v[202])) + k.emit(ds_read_b32(v[207], v[224], v[0], v[0], 0, 0, 15)) + k.emit(v_mfma_fp4(v[88:91], v[192:195], v[64:67], 0, 3, v[211], v[203])) + k.emit(v_mfma_fp4(v[92:95], v[192:195], v[68:71], 2, 3, v[211], v[203])) + k.emit(v_mfma_fp4(v[120:123], v[196:199], v[64:67], 1, 3, v[211], v[203])) + k.emit(v_mfma_fp4(v[124:127], v[196:199], v[68:71], 3, 3, v[211], v[203])) + k.emit(s_waitcnt(127)) + k.emit(v_mfma_fp4(v[128:131], v[168:171], v[72:75], 0, 0, v[210], v[204])) + k.emit(s_barrier()) + k.emit(s_nop()) + k.emit(v_mfma_fp4(v[132:135], v[168:171], v[76:79], 2, 0, v[210], v[204])) + k.emit(s_add_u32(NULL, LIT, s[59], 33792)) + k.emit(v_mfma_fp4(v[160:163], v[172:175], v[72:75], 1, 0, v[210], v[204])) + k.emit(ds_read_b128(v[8:11], v[220])) + k.emit(v_mfma_fp4(v[164:167], v[172:175], v[76:79], 3, 0, v[210], v[204])) + k.emit(buffer_load_dwordx4(v[0:3], v[212], s[12:15], 0, 0, 1, 0, 0, 0, 0, 1)) + k.emit(v_mfma_fp4(v[136:139], v[168:171], v[80:83], 0, 0, v[210], v[205])) + k.emit(ds_read_b128(v[40:43], v[220], v[0], v[0], 0, 64)) + k.emit(v_mfma_fp4(v[140:143], v[168:171], v[84:87], 2, 0, v[210], v[205])) + k.emit(s_add_u32(NULL, LIT, s[59], 38016)) + k.emit(v_mfma_fp4(v[168:171], v[172:175], v[80:83], 1, 0, v[210], v[205])) + k.emit(ds_read_b128(v[12:15], v[220], v[0], v[0], 0, 0, 2)) + k.emit(v_mfma_fp4(v[172:175], v[172:175], v[84:87], 3, 0, v[210], v[205])) + k.emit(buffer_load_dwordx4(v[0:3], v[213], s[12:15], 0, 0, 1, 0, 0, 0, 0, 1)) + for i in range(2): + k.emit(v_mfma_fp4(v[144 + i * 8:147 + i * 8], v[168:171], v[88 + i * 8:91 + i * 8], 0, 0, v[210], v[206 + i * 1])) + k.emit(ds_read_b128(v[44 + i * 4:47 + i * 4], v[220], v[0], v[0], 0, 64 + i * 128, 2 + i * 14)) + k.emit(v_mfma_fp4(v[148 + i * 8:151 + i * 8], v[168:171], v[92 + i * 8:95 + i * 8], 2, 0, v[210], v[206 + i * 1])) + k.emit(s_add_u32(NULL, LIT, s[59], 42240 + i * 4224)) + k.emit(v_mfma_fp4(v[176 + i * 8:179 + i * 8], v[172:175], v[88 + i * 8:91 + i * 8], 1, 0, v[210], v[206 + i * 1])) + k.emit(ds_read_b128(v[16 + i * 4:19 + i * 4], v[220], v[0], v[0], 0, 128, 16 + i * 2)) + k.emit(v_mfma_fp4(v[180 + i * 8:183 + i * 8], v[172:175], v[92 + i * 8:95 + i * 8], 3, 0, v[210], v[206 + i * 1])) + k.emit(buffer_load_dwordx4(v[0:3], v[214 + i * 1], s[12:15], 0, 0, 1, 0, 0, 0, 0, 1)) + k.emit(v_mfma_fp4(v[192:195], v[176:179], v[72:75], 0, 0, v[211], v[204])) + k.emit(ds_read_b128(v[52:55], v[220], v[0], v[0], 0, 192, 18)) + k.emit(v_mfma_fp4(v[196:199], v[176:179], v[76:79], 2, 0, v[211], v[204])) + k.emit(s_add_u32(NULL, LIT, s[60], 2048)) + k.emit(v_mfma_fp4(v[224:227], v[180:183], v[72:75], 1, 0, v[211], v[204])) + k.emit(ds_read_b128(v[24:27], v[220], v[0], v[0], 0, 0, 33)) + k.emit(v_mfma_fp4(v[228:231], v[180:183], v[76:79], 3, 0, v[211], v[204])) + k.emit(buffer_load_dword(v[0], v[222], s[20:23], 0, 0, 1, 0, 0, 0, 0, 1)) + k.emit(v_mfma_fp4(v[200:203], v[176:179], v[80:83], 0, 0, v[211], v[205])) + k.emit(ds_read_b128(v[56:59], v[220], v[0], v[0], 0, 64, 33)) + k.emit(v_mfma_fp4(v[204:207], v[176:179], v[84:87], 2, 0, v[211], v[205])) + k.emit(s_add_u32(NULL, LIT, s[59], 50688)) + k.emit(v_mfma_fp4(v[232:235], v[180:183], v[80:83], 1, 0, v[211], v[205])) + k.emit(ds_read_b128(v[28:31], v[220], v[0], v[0], 0, 0, 35)) + for i in range(2): + k.emit(v_mfma_fp4(v[236 + i * 8:239 + i * 8], v[180:183], v[84 + i * 8:87 + i * 8], 3, 0, v[211], v[205 + i * 1])) + k.emit(buffer_load_dwordx4(v[0:3], v[216 + i * 1], s[12:15], 0, 0, 1, 0, 0, 0, 0, 1)) + k.emit(v_mfma_fp4(v[208 + i * 8:211 + i * 8], v[176:179], v[88 + i * 8:91 + i * 8], 0, 0, v[211], v[206 + i * 1])) + k.emit(ds_read_b128(v[60 + i * 4:63 + i * 4], v[220], v[0], v[0], 0, 64 + i * 128, 35 + i * 14)) + k.emit(v_mfma_fp4(v[212 + i * 8:215 + i * 8], v[176:179], v[92 + i * 8:95 + i * 8], 2, 0, v[211], v[206 + i * 1])) + k.emit(s_add_u32(NULL, LIT, s[59], 54912 + i * 4224)) + k.emit(v_mfma_fp4(v[240 + i * 8:243 + i * 8], v[180:183], v[88 + i * 8:91 + i * 8], 1, 0, v[211], v[206 + i * 1])) + k.emit(ds_read_b128(v[32 + i * 4:35 + i * 4], v[220], v[0], v[0], 0, 128, 49 + i * 2)) + k.emit(v_mfma_fp4(v[252:255], v[180:183], v[100:103], 3, 0, v[211], v[207])) + k.emit(buffer_load_dwordx4(v[0:3], v[218], s[12:15], 0, 0, 1, 0, 0, 0, 0, 1)) + k.emit(v_mfma_fp4(v[128:131], v[184:187], v[104:107], 0, 3, v[210], v[204])) + k.emit(ds_read_b128(v[68:71], v[220], v[0], v[0], 0, 192, 51)) + k.emit(v_mfma_fp4(v[132:135], v[184:187], v[108:111], 2, 3, v[210], v[204])) + k.emit(s_add_u32(NULL, LIT, s[59], 63360)) + k.emit(v_mfma_fp4(v[160:163], v[188:191], v[104:107], 1, 3, v[210], v[204])) + k.emit(ds_read_b32(v[200], v[224])) + k.emit(v_mfma_fp4(v[164:167], v[188:191], v[108:111], 3, 3, v[210], v[204])) + k.emit(buffer_load_dwordx4(v[0:3], v[219], s[12:15], 0, 0, 1, 0, 0, 0, 0, 1)) + k.emit(v_mfma_fp4(v[136:139], v[184:187], v[112:115], 0, 3, v[210], v[205])) + k.emit(ds_read_b32(v[201], v[224], v[0], v[0], 0, 0, 1)) + k.emit(v_mfma_fp4(v[140:143], v[184:187], v[116:119], 2, 3, v[210], v[205])) + k.emit(s_add_u32(NULL, LIT, s[60], 3072)) + k.emit(v_mfma_fp4(v[168:171], v[188:191], v[112:115], 1, 3, v[210], v[205])) + k.emit(ds_read_b32(v[202], v[224], v[0], v[0], 0, 0, 2)) + k.emit(v_mfma_fp4(v[172:175], v[188:191], v[116:119], 3, 3, v[210], v[205])) + k.emit(buffer_load_dword(v[0], v[223], s[20:23], 0, 0, 1, 0, 0, 0, 0, 1)) + k.emit(v_mfma_fp4(v[144:147], v[184:187], v[120:123], 0, 3, v[210], v[206])) + k.emit(ds_read_b32(v[203], v[224], v[0], v[0], 0, 0, 3)) + k.emit(v_mfma_fp4(v[148:151], v[184:187], v[124:127], 2, 3, v[210], v[206])) + k.emit(s_add_u32(s[52], LIT, s[50], 768)) + k.emit(v_mfma_fp4(v[176:179], v[188:191], v[120:123], 1, 3, v[210], v[206])) + k.emit(v_mfma_fp4(v[180:183], v[188:191], v[124:127], 3, 3, v[210], v[206])) + k.emit(s_cmp_lt_u32(s[52], s[51])) + for i in range(2): + k.emit(v_mfma_fp4(v[152 + i * 32:155 + i * 32], v[184 + i * 4:187 + i * 4], v[128:131], 0 + i * 1, 3, v[210], v[207])) + k.emit(v_mfma_fp4(v[156 + i * 32:159 + i * 32], v[184 + i * 4:187 + i * 4], v[132:135], 2 + i * 1, 3, v[210], v[207])) + k.emit(s_cselect_b32(s[61 + i * 2], s[61 + i * 2], 0)) + for i in range(2): + k.emit(v_mfma_fp4(v[192 + i * 40:195 + i * 40], v[192 + i * 4:195 + i * 4], v[104 + i * 8:107 + i * 8], 0 + i * 1, 3, v[211], v[204 + i * 1])) + k.emit(v_mfma_fp4(v[196 + i * 40:199 + i * 40], v[192 + i * 4:195 + i * 4], v[108 + i * 8:111 + i * 8], 2 + i * 1, 3, v[211], v[204 + i * 1])) + k.emit(s_add_u32(s[12 + i * 8], s[61 + i * 2], s[12 + i * 8])) + k.emit(v_mfma_fp4(v[224 + i * -16:227 + i * -16], v[196 + i * -4:199 + i * -4], v[104 + i * 16:107 + i * 16], 1 + i * -1, 3, v[211], v[204 + i * 2])) + k.emit(v_mfma_fp4(v[228 + i * -16:231 + i * -16], v[196 + i * -4:199 + i * -4], v[108 + i * 16:111 + i * 16], 3 + i * -1, 3, v[211], v[204 + i * 2])) + k.emit(s_addc_u32(s[13 + i * 8], 0, s[13 + i * 8])) + k.emit(v_mfma_fp4(v[200 + i * 40:203 + i * 40], v[192 + i * 4:195 + i * 4], v[112 + i * 8:115 + i * 8], 0 + i * 1, 3, v[211], v[205 + i * 1])) + k.emit(v_mfma_fp4(v[204 + i * 40:207 + i * 40], v[192 + i * 4:195 + i * 4], v[116 + i * 8:119 + i * 8], 2 + i * 1, 3, v[211], v[205 + i * 1])) + k.emit(s_sub_u32(s[14 + i * 8], s[14 + i * 8], s[61 + i * 2])) + k.emit(v_mfma_fp4(v[216:219], v[192:195], v[128:131], 0, 3, v[211], v[207])) + k.emit(s_addk_i32(s[50], 256)) + k.emit(v_mfma_fp4(v[220:223], v[192:195], v[132:135], 2, 3, v[211], v[207])) + k.emit(s_cmp_lt_i32(s[50], s[51])) + k.emit(v_mfma_fp4(v[248:251], v[196:199], v[128:131], 1, 3, v[211], v[207])) + k.emit(v_mfma_fp4(v[252:255], v[196:199], v[132:135], 3, 3, v[211], v[207])) + k.emit(s_cbranch_scc0(1367), target='L2_3B10') + k.emit(s_branch(64169), target='L2_105C') + k.label('L2_25B8') + k.emit(s_nop()) + k.label('L2_25BC') + k.emit(s_waitcnt(122)) + k.emit(v_mfma_fp4(v[0:3], v[136:139], v[8:11], 0, 0, v[208], v[200])) + k.emit(s_barrier()) + k.emit(s_nop()) + k.emit(s_nop()) + for i in range(2): + k.emit(v_mfma_fp4(v[4 + i * 64:7 + i * 64], v[136 + i * 8:139 + i * 8], v[12:15], 2, 0, v[208 + i * 1], v[200])) + k.emit(ds_read_b128(v[72 + i * 8:75 + i * 8], v[220], v[0], v[0], 0, 0 + i * 128, 66 + i * 16)) + k.emit(v_mfma_fp4(v[32 + i * 64:35 + i * 64], v[140 + i * 8:143 + i * 8], v[8:11], 1, 0, v[208 + i * 1], v[200])) + k.emit(buffer_load_dwordx4(v[168 + i * 16:171 + i * 16], v[225 + i * 4], s[16:19], 0, 0, 1)) + k.emit(v_mfma_fp4(v[36 + i * 64:39 + i * 64], v[140 + i * 8:143 + i * 8], v[12:15], 3, 0, v[208 + i * 1], v[200])) + k.emit(v_mfma_fp4(v[8 + i * 64:11 + i * 64], v[136 + i * 8:139 + i * 8], v[16:19], 0, 0, v[208 + i * 1], v[201])) + for j75 in range(2): + k.emit(v_mfma_fp4(v[12 + j75 * 8 + i * 64:15 + j75 * 8 + i * 64], v[136 + i * 8:139 + i * 8], v[20 + j75 * 8:23 + j75 * 8], 2, 0, v[208 + i * 1], v[201 + j75 * 1])) + k.emit(ds_read_b128(v[104 + j75 * -28 + i * 8:107 + j75 * -28 + i * 8], v[220], v[0], v[0], 0, 64 + j75 * -64 + i * 128, 66 + j75 * 2 + i * 16)) + k.emit(v_mfma_fp4(v[40 + j75 * 8 + i * 64:43 + j75 * 8 + i * 64], v[140 + i * 8:143 + i * 8], v[16 + j75 * 8:19 + j75 * 8], 1, 0, v[208 + i * 1], v[201 + j75 * 1])) + k.emit(buffer_load_dwordx4(v[172 + j75 * 4 + i * 16:175 + j75 * 4 + i * 16], v[226 + j75 * 1 + i * 4], s[16:19], 0, 0, 1)) + k.emit(v_mfma_fp4(v[44 + j75 * 8 + i * 64:47 + j75 * 8 + i * 64], v[140 + i * 8:143 + i * 8], v[20 + j75 * 8:23 + j75 * 8], 3, 0, v[208 + i * 1], v[201 + j75 * 1])) + k.emit(v_mfma_fp4(v[16 + j75 * 8 + i * 64:19 + j75 * 8 + i * 64], v[136 + i * 8:139 + i * 8], v[24 + j75 * 8:27 + j75 * 8], 0, 0, v[208 + i * 1], v[202 + j75 * 1])) + k.emit(v_mfma_fp4(v[28 + i * 64:31 + i * 64], v[136 + i * 8:139 + i * 8], v[36:39], 2, 0, v[208 + i * 1], v[203])) + k.emit(ds_read_b128(v[108 + i * 8:111 + i * 8], v[220], v[0], v[0], 0, 64 + i * 128, 68 + i * 16)) + k.emit(v_mfma_fp4(v[56 + i * 64:59 + i * 64], v[140 + i * 8:143 + i * 8], v[32:35], 1, 0, v[208 + i * 1], v[203])) + k.emit(buffer_load_dwordx4(v[180 + i * 16:183 + i * 16], v[228 + i * 4], s[16:19], 0, 0, 1)) + k.emit(v_mfma_fp4(v[60 + i * 64:63 + i * 64], v[140 + i * 8:143 + i * 8], v[36:39], 3, 0, v[208 + i * 1], v[203])) + k.emit(v_mfma_fp4(v[64 + i * -64:67 + i * -64], v[144 + i * 8:147 + i * 8], v[8 + i * 32:11 + i * 32], 0, 0 + i * 3, v[209 + i * -1], v[200])) + for i in range(2): + k.emit(v_mfma_fp4(v[4 + i * 8:7 + i * 8], v[152:155], v[44 + i * 8:47 + i * 8], 2, 3, v[208], v[200 + i * 1])) + k.emit(ds_read_b128(v[88 + i * 32:91 + i * 32], v[220], v[0], v[0], 0, 0 + i * 64, 99)) + k.emit(v_mfma_fp4(v[32 + i * 8:35 + i * 8], v[156:159], v[40 + i * 8:43 + i * 8], 1, 3, v[208], v[200 + i * 1])) + k.emit(buffer_load_dword(v[210 + i * 1], v[233 + i * 1], s[24:27], 0, 0, 1)) + k.emit(v_mfma_fp4(v[36 + i * 8:39 + i * 8], v[156:159], v[44 + i * 8:47 + i * 8], 3, 3, v[208], v[200 + i * 1])) + k.emit(v_mfma_fp4(v[8 + i * 8:11 + i * 8], v[152:155], v[48 + i * 8:51 + i * 8], 0, 3, v[208], v[201 + i * 1])) + k.emit(ds_read_b128(v[92:95], v[220], v[0], v[0], 0, 0, 101)) + k.emit(v_mfma_fp4(v[20:23], v[152:155], v[60:63], 2, 3, v[208], v[202])) + k.emit(s_add_u32(s[53], LIT, s[50], 512)) + k.emit(v_mfma_fp4(v[48:51], v[156:159], v[56:59], 1, 3, v[208], v[202])) + k.emit(ds_read_b128(v[124:127], v[220], v[0], v[0], 0, 64, 101)) + k.emit(v_mfma_fp4(v[52:55], v[156:159], v[60:63], 3, 3, v[208], v[202])) + k.emit(s_cmp_lt_u32(s[53], s[51])) + for i in range(2): + k.emit(v_mfma_fp4(v[24 + i * 32:27 + i * 32], v[152 + i * 4:155 + i * 4], v[64:67], 0 + i * 1, 3, v[208], v[203])) + k.emit(ds_read_b128(v[96 + i * 32:99 + i * 32], v[220], v[0], v[0], 0, 128 + i * 64, 115)) + k.emit(v_mfma_fp4(v[28 + i * 32:31 + i * 32], v[152 + i * 4:155 + i * 4], v[68:71], 2 + i * 1, 3, v[208], v[203])) + k.emit(s_cselect_b32(s[62 + i * 2], s[62 + i * 2], 0)) + k.emit(v_mfma_fp4(v[64:67], v[160:163], v[40:43], 0, 3, v[209], v[200])) + k.emit(ds_read_b128(v[100:103], v[220], v[0], v[0], 0, 128, 117)) + k.emit(v_mfma_fp4(v[68:71], v[160:163], v[44:47], 2, 3, v[209], v[200])) + k.emit(s_add_u32(s[16], s[62], s[16])) + k.emit(v_mfma_fp4(v[96:99], v[164:167], v[40:43], 1, 3, v[209], v[200])) + k.emit(ds_read_b128(v[132:135], v[220], v[0], v[0], 0, 192, 117)) + k.emit(v_mfma_fp4(v[100:103], v[164:167], v[44:47], 3, 3, v[209], v[200])) + k.emit(s_addc_u32(s[17], 0, s[17])) + k.emit(v_mfma_fp4(v[72:75], v[160:163], v[48:51], 0, 3, v[209], v[201])) + k.emit(ds_read_b32(v[204], v[224], v[0], v[0], 0, 0, 4)) + k.emit(v_mfma_fp4(v[76:79], v[160:163], v[52:55], 2, 3, v[209], v[201])) + k.emit(s_sub_u32(s[18], s[18], s[62])) + k.emit(v_mfma_fp4(v[104:107], v[164:167], v[48:51], 1, 3, v[209], v[201])) + k.emit(ds_read_b32(v[205], v[224], v[0], v[0], 0, 0, 5)) + k.emit(v_mfma_fp4(v[108:111], v[164:167], v[52:55], 3, 3, v[209], v[201])) + k.emit(s_add_u32(s[24], s[64], s[24])) + k.emit(v_mfma_fp4(v[80:83], v[160:163], v[56:59], 0, 3, v[209], v[202])) + k.emit(ds_read_b32(v[206], v[224], v[0], v[0], 0, 0, 6)) + k.emit(v_mfma_fp4(v[84:87], v[160:163], v[60:63], 2, 3, v[209], v[202])) + k.emit(s_addc_u32(s[25], 0, s[25])) + k.emit(v_mfma_fp4(v[112:115], v[164:167], v[56:59], 1, 3, v[209], v[202])) + k.emit(ds_read_b32(v[207], v[224], v[0], v[0], 0, 0, 7)) + k.emit(v_mfma_fp4(v[116:119], v[164:167], v[60:63], 3, 3, v[209], v[202])) + k.emit(s_sub_u32(s[26], s[26], s[64])) + k.emit(v_mfma_fp4(v[88:91], v[160:163], v[64:67], 0, 3, v[209], v[203])) + k.emit(v_mfma_fp4(v[92:95], v[160:163], v[68:71], 2, 3, v[209], v[203])) + k.emit(v_mfma_fp4(v[120:123], v[164:167], v[64:67], 1, 3, v[209], v[203])) + k.emit(v_mfma_fp4(v[124:127], v[164:167], v[68:71], 3, 3, v[209], v[203])) + k.emit(s_waitcnt(127)) + k.emit(v_mfma_fp4(v[128:131], v[136:139], v[72:75], 0, 0, v[208], v[204])) + k.emit(s_barrier()) + k.emit(s_nop()) + k.emit(v_mfma_fp4(v[132:135], v[136:139], v[76:79], 2, 0, v[208], v[204])) + k.emit(ds_read_b128(v[8:11], v[221])) + k.emit(v_mfma_fp4(v[160:163], v[140:143], v[72:75], 1, 0, v[208], v[204])) + k.emit(s_add_u32(NULL, 0, s[59])) + k.emit(v_mfma_fp4(v[164:167], v[140:143], v[76:79], 3, 0, v[208], v[204])) + k.emit(ds_read_b128(v[40:43], v[221], v[0], v[0], 0, 64)) + k.emit(v_mfma_fp4(v[136:139], v[136:139], v[80:83], 0, 0, v[208], v[205])) + k.emit(buffer_load_dwordx4(v[0:3], v[212], s[12:15], 0, 0, 1, 0, 0, 0, 0, 1)) + k.emit(v_mfma_fp4(v[140:143], v[136:139], v[84:87], 2, 0, v[208], v[205])) + k.emit(ds_read_b128(v[12:15], v[221], v[0], v[0], 0, 0, 2)) + k.emit(v_mfma_fp4(v[168:171], v[140:143], v[80:83], 1, 0, v[208], v[205])) + k.emit(s_add_u32(NULL, LIT, s[59], 4224)) + for i in range(2): + k.emit(v_mfma_fp4(v[172 + i * 8:175 + i * 8], v[140:143], v[84 + i * 8:87 + i * 8], 3, 0, v[208], v[205 + i * 1])) + k.emit(ds_read_b128(v[44 + i * 4:47 + i * 4], v[221], v[0], v[0], 0, 64 + i * 128, 2 + i * 14)) + k.emit(v_mfma_fp4(v[144 + i * 8:147 + i * 8], v[136:139], v[88 + i * 8:91 + i * 8], 0, 0, v[208], v[206 + i * 1])) + k.emit(buffer_load_dwordx4(v[0:3], v[213 + i * 1], s[12:15], 0, 0, 1, 0, 0, 0, 0, 1)) + k.emit(v_mfma_fp4(v[148 + i * 8:151 + i * 8], v[136:139], v[92 + i * 8:95 + i * 8], 2, 0, v[208], v[206 + i * 1])) + k.emit(ds_read_b128(v[16 + i * 4:19 + i * 4], v[221], v[0], v[0], 0, 128, 16 + i * 2)) + k.emit(v_mfma_fp4(v[176 + i * 8:179 + i * 8], v[140:143], v[88 + i * 8:91 + i * 8], 1, 0, v[208], v[206 + i * 1])) + k.emit(s_add_u32(NULL, LIT, s[59], 8448 + i * 4224)) + k.emit(v_mfma_fp4(v[188:191], v[140:143], v[100:103], 3, 0, v[208], v[207])) + k.emit(ds_read_b128(v[52:55], v[221], v[0], v[0], 0, 192, 18)) + k.emit(v_mfma_fp4(v[192:195], v[144:147], v[72:75], 0, 0, v[209], v[204])) + k.emit(buffer_load_dwordx4(v[0:3], v[215], s[12:15], 0, 0, 1, 0, 0, 0, 0, 1)) + k.emit(v_mfma_fp4(v[196:199], v[144:147], v[76:79], 2, 0, v[209], v[204])) + k.emit(ds_read_b128(v[24:27], v[221], v[0], v[0], 0, 0, 33)) + k.emit(v_mfma_fp4(v[224:227], v[148:151], v[72:75], 1, 0, v[209], v[204])) + k.emit(s_add_u32(NULL, 0, s[60])) + k.emit(v_mfma_fp4(v[228:231], v[148:151], v[76:79], 3, 0, v[209], v[204])) + k.emit(ds_read_b128(v[56:59], v[221], v[0], v[0], 0, 64, 33)) + k.emit(v_mfma_fp4(v[200:203], v[144:147], v[80:83], 0, 0, v[209], v[205])) + k.emit(buffer_load_dword(v[0], v[222], s[20:23], 0, 0, 1, 0, 0, 0, 0, 1)) + k.emit(v_mfma_fp4(v[204:207], v[144:147], v[84:87], 2, 0, v[209], v[205])) + k.emit(ds_read_b128(v[28:31], v[221], v[0], v[0], 0, 0, 35)) + for i in range(2): + k.emit(v_mfma_fp4(v[232 + i * 8:235 + i * 8], v[148:151], v[80 + i * 8:83 + i * 8], 1, 0, v[209], v[205 + i * 1])) + k.emit(s_add_u32(NULL, LIT, s[59], 16896 + i * 4224)) + k.emit(v_mfma_fp4(v[236 + i * 8:239 + i * 8], v[148:151], v[84 + i * 8:87 + i * 8], 3, 0, v[209], v[205 + i * 1])) + k.emit(ds_read_b128(v[60 + i * 4:63 + i * 4], v[221], v[0], v[0], 0, 64 + i * 128, 35 + i * 14)) + k.emit(v_mfma_fp4(v[208 + i * 8:211 + i * 8], v[144:147], v[88 + i * 8:91 + i * 8], 0, 0, v[209], v[206 + i * 1])) + k.emit(buffer_load_dwordx4(v[0:3], v[216 + i * 1], s[12:15], 0, 0, 1, 0, 0, 0, 0, 1)) + k.emit(v_mfma_fp4(v[212 + i * 8:215 + i * 8], v[144:147], v[92 + i * 8:95 + i * 8], 2, 0, v[209], v[206 + i * 1])) + k.emit(ds_read_b128(v[32 + i * 4:35 + i * 4], v[221], v[0], v[0], 0, 128, 49 + i * 2)) + k.emit(v_mfma_fp4(v[248:251], v[148:151], v[96:99], 1, 0, v[209], v[207])) + k.emit(s_add_u32(NULL, LIT, s[59], 25344)) + k.emit(v_mfma_fp4(v[252:255], v[148:151], v[100:103], 3, 0, v[209], v[207])) + k.emit(ds_read_b128(v[68:71], v[221], v[0], v[0], 0, 192, 51)) + for i in range(2): + k.emit(v_mfma_fp4(v[128 + i * 8:131 + i * 8], v[152:155], v[104 + i * 8:107 + i * 8], 0, 3, v[208], v[204 + i * 1])) + k.emit(buffer_load_dwordx4(v[0:3], v[218 + i * 1], s[12:15], 0, 0, 1, 0, 0, 0, 0, 1)) + k.emit(v_mfma_fp4(v[132 + i * 8:135 + i * 8], v[152:155], v[108 + i * 8:111 + i * 8], 2, 3, v[208], v[204 + i * 1])) + k.emit(ds_read_b32(v[200 + i * 2], v[224], v[0], v[0], 0, 0, 8 + i * 2)) + k.emit(v_mfma_fp4(v[160 + i * 8:163 + i * 8], v[156:159], v[104 + i * 8:107 + i * 8], 1, 3, v[208], v[204 + i * 1])) + k.emit(s_add_u32(NULL, LIT, s[59 + i * 1], 29568 + i * -28544)) + k.emit(v_mfma_fp4(v[164 + i * 8:167 + i * 8], v[156:159], v[108 + i * 8:111 + i * 8], 3, 3, v[208], v[204 + i * 1])) + k.emit(ds_read_b32(v[201 + i * 2], v[224], v[0], v[0], 0, 0, 9 + i * 2)) + k.emit(v_mfma_fp4(v[144:147], v[152:155], v[120:123], 0, 3, v[208], v[206])) + k.emit(buffer_load_dword(v[0], v[223], s[20:23], 0, 0, 1, 0, 0, 0, 0, 1)) + k.emit(v_mfma_fp4(v[148:151], v[152:155], v[124:127], 2, 3, v[208], v[206])) + k.emit(v_mfma_fp4(v[176:179], v[156:159], v[120:123], 1, 3, v[208], v[206])) + k.emit(s_add_u32(s[52], LIT, s[50], 768)) + k.emit(v_mfma_fp4(v[180:183], v[156:159], v[124:127], 3, 3, v[208], v[206])) + k.emit(v_mfma_fp4(v[152:155], v[152:155], v[128:131], 0, 3, v[208], v[207])) + k.emit(s_cmp_lt_u32(s[52], s[51])) + for i in range(2): + k.emit(v_mfma_fp4(v[156 + i * 32:159 + i * 32], v[152 + i * 4:155 + i * 4], v[132:135], 2 + i * 1, 3, v[208], v[207])) + k.emit(v_mfma_fp4(v[184 + i * 8:187 + i * 8], v[156 + i * 4:159 + i * 4], v[128 + i * -24:131 + i * -24], 1 + i * -1, 3, v[208 + i * 1], v[207 + i * -3])) + k.emit(s_cselect_b32(s[61 + i * 2], s[61 + i * 2], 0)) + k.emit(v_mfma_fp4(v[196:199], v[160:163], v[108:111], 2, 3, v[209], v[204])) + k.emit(v_mfma_fp4(v[224:227], v[164:167], v[104:107], 1, 3, v[209], v[204])) + k.emit(s_add_u32(s[12], s[61], s[12])) + k.emit(v_mfma_fp4(v[228:231], v[164:167], v[108:111], 3, 3, v[209], v[204])) + k.emit(v_mfma_fp4(v[200:203], v[160:163], v[112:115], 0, 3, v[209], v[205])) + k.emit(s_addc_u32(s[13], 0, s[13])) + k.emit(v_mfma_fp4(v[204:207], v[160:163], v[116:119], 2, 3, v[209], v[205])) + k.emit(v_mfma_fp4(v[232:235], v[164:167], v[112:115], 1, 3, v[209], v[205])) + k.emit(s_sub_u32(s[14], s[14], s[61])) + k.emit(v_mfma_fp4(v[236:239], v[164:167], v[116:119], 3, 3, v[209], v[205])) + k.emit(v_mfma_fp4(v[208:211], v[160:163], v[120:123], 0, 3, v[209], v[206])) + k.emit(s_add_u32(s[20], s[63], s[20])) + k.emit(v_mfma_fp4(v[212:215], v[160:163], v[124:127], 2, 3, v[209], v[206])) + k.emit(v_mfma_fp4(v[240:243], v[164:167], v[120:123], 1, 3, v[209], v[206])) + k.emit(s_addc_u32(s[21], 0, s[21])) + k.emit(v_mfma_fp4(v[244:247], v[164:167], v[124:127], 3, 3, v[209], v[206])) + k.emit(v_mfma_fp4(v[216:219], v[160:163], v[128:131], 0, 3, v[209], v[207])) + k.emit(s_addk_i32(s[50], 256)) + k.emit(v_mfma_fp4(v[220:223], v[160:163], v[132:135], 2, 3, v[209], v[207])) + k.emit(s_cmp_lt_i32(s[50], s[51])) + k.emit(v_mfma_fp4(v[248:251], v[164:167], v[128:131], 1, 3, v[209], v[207])) + k.emit(v_mfma_fp4(v[252:255], v[164:167], v[132:135], 3, 3, v[209], v[207])) + k.emit(s_cbranch_scc0(684), target='L2_3B10') + k.emit(s_waitcnt(122)) + k.emit(v_mfma_fp4(v[0:3], v[168:171], v[8:11], 0, 0, v[210], v[200])) + k.emit(s_barrier()) + k.emit(s_nop()) + k.emit(s_nop()) + for i in range(2): + k.emit(v_mfma_fp4(v[4 + i * 64:7 + i * 64], v[168 + i * 8:171 + i * 8], v[12:15], 2, 0, v[210 + i * 1], v[200])) + k.emit(ds_read_b128(v[72 + i * 8:75 + i * 8], v[221], v[0], v[0], 0, 0 + i * 128, 66 + i * 16)) + k.emit(v_mfma_fp4(v[32 + i * 64:35 + i * 64], v[172 + i * 8:175 + i * 8], v[8:11], 1, 0, v[210 + i * 1], v[200])) + k.emit(buffer_load_dwordx4(v[136 + i * 16:139 + i * 16], v[225 + i * 4], s[16:19], 0, 0, 1)) + k.emit(v_mfma_fp4(v[36 + i * 64:39 + i * 64], v[172 + i * 8:175 + i * 8], v[12:15], 3, 0, v[210 + i * 1], v[200])) + k.emit(v_mfma_fp4(v[8 + i * 64:11 + i * 64], v[168 + i * 8:171 + i * 8], v[16:19], 0, 0, v[210 + i * 1], v[201])) + for j76 in range(2): + k.emit(v_mfma_fp4(v[12 + j76 * 8 + i * 64:15 + j76 * 8 + i * 64], v[168 + i * 8:171 + i * 8], v[20 + j76 * 8:23 + j76 * 8], 2, 0, v[210 + i * 1], v[201 + j76 * 1])) + k.emit(ds_read_b128(v[104 + j76 * -28 + i * 8:107 + j76 * -28 + i * 8], v[221], v[0], v[0], 0, 64 + j76 * -64 + i * 128, 66 + j76 * 2 + i * 16)) + k.emit(v_mfma_fp4(v[40 + j76 * 8 + i * 64:43 + j76 * 8 + i * 64], v[172 + i * 8:175 + i * 8], v[16 + j76 * 8:19 + j76 * 8], 1, 0, v[210 + i * 1], v[201 + j76 * 1])) + k.emit(buffer_load_dwordx4(v[140 + j76 * 4 + i * 16:143 + j76 * 4 + i * 16], v[226 + j76 * 1 + i * 4], s[16:19], 0, 0, 1)) + k.emit(v_mfma_fp4(v[44 + j76 * 8 + i * 64:47 + j76 * 8 + i * 64], v[172 + i * 8:175 + i * 8], v[20 + j76 * 8:23 + j76 * 8], 3, 0, v[210 + i * 1], v[201 + j76 * 1])) + k.emit(v_mfma_fp4(v[16 + j76 * 8 + i * 64:19 + j76 * 8 + i * 64], v[168 + i * 8:171 + i * 8], v[24 + j76 * 8:27 + j76 * 8], 0, 0, v[210 + i * 1], v[202 + j76 * 1])) + k.emit(v_mfma_fp4(v[28 + i * 64:31 + i * 64], v[168 + i * 8:171 + i * 8], v[36:39], 2, 0, v[210 + i * 1], v[203])) + k.emit(ds_read_b128(v[108 + i * 8:111 + i * 8], v[221], v[0], v[0], 0, 64 + i * 128, 68 + i * 16)) + k.emit(v_mfma_fp4(v[56 + i * 64:59 + i * 64], v[172 + i * 8:175 + i * 8], v[32:35], 1, 0, v[210 + i * 1], v[203])) + k.emit(buffer_load_dwordx4(v[148 + i * 16:151 + i * 16], v[228 + i * 4], s[16:19], 0, 0, 1)) + k.emit(v_mfma_fp4(v[60 + i * 64:63 + i * 64], v[172 + i * 8:175 + i * 8], v[36:39], 3, 0, v[210 + i * 1], v[203])) + k.emit(v_mfma_fp4(v[64 + i * -64:67 + i * -64], v[176 + i * 8:179 + i * 8], v[8 + i * 32:11 + i * 32], 0, 0 + i * 3, v[211 + i * -1], v[200])) + for i in range(2): + k.emit(v_mfma_fp4(v[4 + i * 8:7 + i * 8], v[184:187], v[44 + i * 8:47 + i * 8], 2, 3, v[210], v[200 + i * 1])) + k.emit(ds_read_b128(v[88 + i * 32:91 + i * 32], v[221], v[0], v[0], 0, 0 + i * 64, 99)) + k.emit(v_mfma_fp4(v[32 + i * 8:35 + i * 8], v[188:191], v[40 + i * 8:43 + i * 8], 1, 3, v[210], v[200 + i * 1])) + k.emit(buffer_load_dword(v[208 + i * 1], v[233 + i * 1], s[24:27], 0, 0, 1)) + k.emit(v_mfma_fp4(v[36 + i * 8:39 + i * 8], v[188:191], v[44 + i * 8:47 + i * 8], 3, 3, v[210], v[200 + i * 1])) + k.emit(v_mfma_fp4(v[8 + i * 8:11 + i * 8], v[184:187], v[48 + i * 8:51 + i * 8], 0, 3, v[210], v[201 + i * 1])) + k.emit(ds_read_b128(v[92:95], v[221], v[0], v[0], 0, 0, 101)) + k.emit(v_mfma_fp4(v[20:23], v[184:187], v[60:63], 2, 3, v[210], v[202])) + k.emit(s_add_u32(s[53], LIT, s[50], 512)) + k.emit(v_mfma_fp4(v[48:51], v[188:191], v[56:59], 1, 3, v[210], v[202])) + k.emit(ds_read_b128(v[124:127], v[221], v[0], v[0], 0, 64, 101)) + k.emit(v_mfma_fp4(v[52:55], v[188:191], v[60:63], 3, 3, v[210], v[202])) + k.emit(s_cmp_lt_u32(s[53], s[51])) + for i in range(2): + k.emit(v_mfma_fp4(v[24 + i * 32:27 + i * 32], v[184 + i * 4:187 + i * 4], v[64:67], 0 + i * 1, 3, v[210], v[203])) + k.emit(ds_read_b128(v[96 + i * 32:99 + i * 32], v[221], v[0], v[0], 0, 128 + i * 64, 115)) + k.emit(v_mfma_fp4(v[28 + i * 32:31 + i * 32], v[184 + i * 4:187 + i * 4], v[68:71], 2 + i * 1, 3, v[210], v[203])) + k.emit(s_cselect_b32(s[62 + i * 2], s[62 + i * 2], 0)) + k.emit(v_mfma_fp4(v[64:67], v[192:195], v[40:43], 0, 3, v[211], v[200])) + k.emit(ds_read_b128(v[100:103], v[221], v[0], v[0], 0, 128, 117)) + k.emit(v_mfma_fp4(v[68:71], v[192:195], v[44:47], 2, 3, v[211], v[200])) + k.emit(s_add_u32(s[16], s[62], s[16])) + k.emit(v_mfma_fp4(v[96:99], v[196:199], v[40:43], 1, 3, v[211], v[200])) + k.emit(ds_read_b128(v[132:135], v[221], v[0], v[0], 0, 192, 117)) + k.emit(v_mfma_fp4(v[100:103], v[196:199], v[44:47], 3, 3, v[211], v[200])) + k.emit(s_addc_u32(s[17], 0, s[17])) + k.emit(v_mfma_fp4(v[72:75], v[192:195], v[48:51], 0, 3, v[211], v[201])) + k.emit(ds_read_b32(v[204], v[224], v[0], v[0], 0, 0, 12)) + k.emit(v_mfma_fp4(v[76:79], v[192:195], v[52:55], 2, 3, v[211], v[201])) + k.emit(s_sub_u32(s[18], s[18], s[62])) + k.emit(v_mfma_fp4(v[104:107], v[196:199], v[48:51], 1, 3, v[211], v[201])) + k.emit(ds_read_b32(v[205], v[224], v[0], v[0], 0, 0, 13)) + k.emit(v_mfma_fp4(v[108:111], v[196:199], v[52:55], 3, 3, v[211], v[201])) + k.emit(s_add_u32(s[24], s[64], s[24])) + k.emit(v_mfma_fp4(v[80:83], v[192:195], v[56:59], 0, 3, v[211], v[202])) + k.emit(ds_read_b32(v[206], v[224], v[0], v[0], 0, 0, 14)) + k.emit(v_mfma_fp4(v[84:87], v[192:195], v[60:63], 2, 3, v[211], v[202])) + k.emit(s_addc_u32(s[25], 0, s[25])) + k.emit(v_mfma_fp4(v[112:115], v[196:199], v[56:59], 1, 3, v[211], v[202])) + k.emit(ds_read_b32(v[207], v[224], v[0], v[0], 0, 0, 15)) + k.emit(v_mfma_fp4(v[116:119], v[196:199], v[60:63], 3, 3, v[211], v[202])) + k.emit(s_sub_u32(s[26], s[26], s[64])) + k.emit(v_mfma_fp4(v[88:91], v[192:195], v[64:67], 0, 3, v[211], v[203])) + k.emit(v_mfma_fp4(v[92:95], v[192:195], v[68:71], 2, 3, v[211], v[203])) + k.emit(v_mfma_fp4(v[120:123], v[196:199], v[64:67], 1, 3, v[211], v[203])) + k.emit(v_mfma_fp4(v[124:127], v[196:199], v[68:71], 3, 3, v[211], v[203])) + k.emit(s_waitcnt(127)) + k.emit(v_mfma_fp4(v[128:131], v[168:171], v[72:75], 0, 0, v[210], v[204])) + k.emit(s_barrier()) + k.emit(s_nop()) + k.emit(v_mfma_fp4(v[132:135], v[168:171], v[76:79], 2, 0, v[210], v[204])) + k.emit(ds_read_b128(v[8:11], v[220])) + k.emit(v_mfma_fp4(v[160:163], v[172:175], v[72:75], 1, 0, v[210], v[204])) + k.emit(s_add_u32(NULL, LIT, s[59], 33792)) + k.emit(v_mfma_fp4(v[164:167], v[172:175], v[76:79], 3, 0, v[210], v[204])) + k.emit(ds_read_b128(v[40:43], v[220], v[0], v[0], 0, 64)) + for i in range(2): + k.emit(v_mfma_fp4(v[136 + i * 16:139 + i * 16], v[168:171], v[80 + i * 16:83 + i * 16], 0, 0, v[210], v[205 + i * 2])) + k.emit(buffer_load_dwordx4(v[0:3], v[212 + i * 2], s[12:15], 0, 0, 1, 0, 0, 0, 0, 1)) + k.emit(v_mfma_fp4(v[140 + i * 16:143 + i * 16], v[168:171], v[84 + i * 16:87 + i * 16], 2, 0, v[210], v[205 + i * 2])) + k.emit(ds_read_b128(v[12 + i * 8:15 + i * 8], v[220], v[0], v[0], 0, 0 + i * 128, 2 + i * 16)) + k.emit(v_mfma_fp4(v[168 + i * 16:171 + i * 16], v[172:175], v[80 + i * 16:83 + i * 16], 1, 0, v[210], v[205 + i * 2])) + k.emit(s_add_u32(NULL, LIT, s[59], 38016 + i * 8448)) + k.emit(v_mfma_fp4(v[172 + i * 16:175 + i * 16], v[172:175], v[84 + i * 16:87 + i * 16], 3, 0, v[210], v[205 + i * 2])) + k.emit(ds_read_b128(v[44 + i * 8:47 + i * 8], v[220], v[0], v[0], 0, 64 + i * 128, 2 + i * 16)) + k.emit(v_mfma_fp4(v[144 + i * 48:147 + i * 48], v[168 + i * 8:171 + i * 8], v[88 + i * -16:91 + i * -16], 0, 0, v[210 + i * 1], v[206 + i * -2])) + k.emit(buffer_load_dwordx4(v[0:3], v[213 + i * 2], s[12:15], 0, 0, 1, 0, 0, 0, 0, 1)) + k.emit(v_mfma_fp4(v[148 + i * 48:151 + i * 48], v[168 + i * 8:171 + i * 8], v[92 + i * -16:95 + i * -16], 2, 0, v[210 + i * 1], v[206 + i * -2])) + k.emit(ds_read_b128(v[16 + i * 8:19 + i * 8], v[220], v[0], v[0], 0, 128 + i * -128, 16 + i * 17)) + k.emit(v_mfma_fp4(v[176 + i * 48:179 + i * 48], v[172 + i * 8:175 + i * 8], v[88 + i * -16:91 + i * -16], 1, 0, v[210 + i * 1], v[206 + i * -2])) + k.emit(s_add_u32(NULL, LIT, s[59 + i * 1], 42240 + i * -40192)) + k.emit(v_mfma_fp4(v[180 + i * 48:183 + i * 48], v[172 + i * 8:175 + i * 8], v[92 + i * -16:95 + i * -16], 3, 0, v[210 + i * 1], v[206 + i * -2])) + k.emit(ds_read_b128(v[48 + i * 8:51 + i * 8], v[220], v[0], v[0], 0, 192 + i * -128, 16 + i * 17)) + k.emit(v_mfma_fp4(v[200:203], v[176:179], v[80:83], 0, 0, v[211], v[205])) + k.emit(buffer_load_dword(v[0], v[222], s[20:23], 0, 0, 1, 0, 0, 0, 0, 1)) + k.emit(v_mfma_fp4(v[204:207], v[176:179], v[84:87], 2, 0, v[211], v[205])) + k.emit(ds_read_b128(v[28:31], v[220], v[0], v[0], 0, 0, 35)) + for i in range(2): + k.emit(v_mfma_fp4(v[232 + i * 8:235 + i * 8], v[180:183], v[80 + i * 8:83 + i * 8], 1, 0, v[211], v[205 + i * 1])) + k.emit(s_add_u32(NULL, LIT, s[59], 50688 + i * 4224)) + k.emit(v_mfma_fp4(v[236 + i * 8:239 + i * 8], v[180:183], v[84 + i * 8:87 + i * 8], 3, 0, v[211], v[205 + i * 1])) + k.emit(ds_read_b128(v[60 + i * 4:63 + i * 4], v[220], v[0], v[0], 0, 64 + i * 128, 35 + i * 14)) + k.emit(v_mfma_fp4(v[208 + i * 8:211 + i * 8], v[176:179], v[88 + i * 8:91 + i * 8], 0, 0, v[211], v[206 + i * 1])) + k.emit(buffer_load_dwordx4(v[0:3], v[216 + i * 1], s[12:15], 0, 0, 1, 0, 0, 0, 0, 1)) + k.emit(v_mfma_fp4(v[212 + i * 8:215 + i * 8], v[176:179], v[92 + i * 8:95 + i * 8], 2, 0, v[211], v[206 + i * 1])) + k.emit(ds_read_b128(v[32 + i * 4:35 + i * 4], v[220], v[0], v[0], 0, 128, 49 + i * 2)) + k.emit(v_mfma_fp4(v[248:251], v[180:183], v[96:99], 1, 0, v[211], v[207])) + k.emit(s_add_u32(NULL, LIT, s[59], 59136)) + k.emit(v_mfma_fp4(v[252:255], v[180:183], v[100:103], 3, 0, v[211], v[207])) + k.emit(ds_read_b128(v[68:71], v[220], v[0], v[0], 0, 192, 51)) + k.emit(v_mfma_fp4(v[128:131], v[184:187], v[104:107], 0, 3, v[210], v[204])) + k.emit(buffer_load_dwordx4(v[0:3], v[218], s[12:15], 0, 0, 1, 0, 0, 0, 0, 1)) + k.emit(v_mfma_fp4(v[132:135], v[184:187], v[108:111], 2, 3, v[210], v[204])) + k.emit(ds_read_b32(v[200], v[224])) + k.emit(v_mfma_fp4(v[160:163], v[188:191], v[104:107], 1, 3, v[210], v[204])) + k.emit(s_add_u32(NULL, LIT, s[59], 63360)) + k.emit(v_mfma_fp4(v[164:167], v[188:191], v[108:111], 3, 3, v[210], v[204])) + k.emit(ds_read_b32(v[201], v[224], v[0], v[0], 0, 0, 1)) + k.emit(v_mfma_fp4(v[136:139], v[184:187], v[112:115], 0, 3, v[210], v[205])) + k.emit(buffer_load_dwordx4(v[0:3], v[219], s[12:15], 0, 0, 1, 0, 0, 0, 0, 1)) + k.emit(v_mfma_fp4(v[140:143], v[184:187], v[116:119], 2, 3, v[210], v[205])) + k.emit(ds_read_b32(v[202], v[224], v[0], v[0], 0, 0, 2)) + k.emit(v_mfma_fp4(v[168:171], v[188:191], v[112:115], 1, 3, v[210], v[205])) + k.emit(s_add_u32(NULL, LIT, s[60], 3072)) + k.emit(v_mfma_fp4(v[172:175], v[188:191], v[116:119], 3, 3, v[210], v[205])) + k.emit(ds_read_b32(v[203], v[224], v[0], v[0], 0, 0, 3)) + k.emit(v_mfma_fp4(v[144:147], v[184:187], v[120:123], 0, 3, v[210], v[206])) + k.emit(buffer_load_dword(v[0], v[223], s[20:23], 0, 0, 1, 0, 0, 0, 0, 1)) + k.emit(v_mfma_fp4(v[148:151], v[184:187], v[124:127], 2, 3, v[210], v[206])) + k.emit(v_mfma_fp4(v[176:179], v[188:191], v[120:123], 1, 3, v[210], v[206])) + k.emit(s_add_u32(s[52], LIT, s[50], 768)) + k.emit(v_mfma_fp4(v[180:183], v[188:191], v[124:127], 3, 3, v[210], v[206])) + k.emit(v_mfma_fp4(v[152:155], v[184:187], v[128:131], 0, 3, v[210], v[207])) + k.emit(s_cmp_lt_u32(s[52], s[51])) + for i in range(2): + k.emit(v_mfma_fp4(v[156 + i * 32:159 + i * 32], v[184 + i * 4:187 + i * 4], v[132:135], 2 + i * 1, 3, v[210], v[207])) + k.emit(v_mfma_fp4(v[184 + i * 8:187 + i * 8], v[188 + i * 4:191 + i * 4], v[128 + i * -24:131 + i * -24], 1 + i * -1, 3, v[210 + i * 1], v[207 + i * -3])) + k.emit(s_cselect_b32(s[61 + i * 2], s[61 + i * 2], 0)) + k.emit(v_mfma_fp4(v[196:199], v[192:195], v[108:111], 2, 3, v[211], v[204])) + k.emit(v_mfma_fp4(v[224:227], v[196:199], v[104:107], 1, 3, v[211], v[204])) + k.emit(s_add_u32(s[12], s[61], s[12])) + k.emit(v_mfma_fp4(v[228:231], v[196:199], v[108:111], 3, 3, v[211], v[204])) + k.emit(v_mfma_fp4(v[200:203], v[192:195], v[112:115], 0, 3, v[211], v[205])) + k.emit(s_addc_u32(s[13], 0, s[13])) + k.emit(v_mfma_fp4(v[204:207], v[192:195], v[116:119], 2, 3, v[211], v[205])) + k.emit(v_mfma_fp4(v[232:235], v[196:199], v[112:115], 1, 3, v[211], v[205])) + k.emit(s_sub_u32(s[14], s[14], s[61])) + k.emit(v_mfma_fp4(v[236:239], v[196:199], v[116:119], 3, 3, v[211], v[205])) + k.emit(v_mfma_fp4(v[208:211], v[192:195], v[120:123], 0, 3, v[211], v[206])) + k.emit(s_add_u32(s[20], s[63], s[20])) + k.emit(v_mfma_fp4(v[212:215], v[192:195], v[124:127], 2, 3, v[211], v[206])) + k.emit(v_mfma_fp4(v[240:243], v[196:199], v[120:123], 1, 3, v[211], v[206])) + k.emit(s_addc_u32(s[21], 0, s[21])) + k.emit(v_mfma_fp4(v[244:247], v[196:199], v[124:127], 3, 3, v[211], v[206])) + k.emit(v_mfma_fp4(v[216:219], v[192:195], v[128:131], 0, 3, v[211], v[207])) + k.emit(s_addk_i32(s[50], 256)) + k.emit(v_mfma_fp4(v[220:223], v[192:195], v[132:135], 2, 3, v[211], v[207])) + k.emit(s_cmp_lt_i32(s[50], s[51])) + k.emit(v_mfma_fp4(v[248:251], v[196:199], v[128:131], 1, 3, v[211], v[207])) + k.emit(v_mfma_fp4(v[252:255], v[196:199], v[132:135], 3, 3, v[211], v[207])) + k.emit(s_cbranch_scc0(1), target='L2_3B10') + k.emit(s_branch(64171), target='L2_25BC') + k.label('L2_3B10') + k.emit(s_waitcnt(112)) + k.emit(s_barrier()) + k.emit(s_cmp_eq_u32(s[57], 0)) + k.emit(s_cbranch_scc1(1505), target='L2_52A4') + for i in range(2): + for j77 in range(4): + k.emit(v_accvgpr_read(v[8 + j77 * 1 + i * 4], v[0 + j77 * 1 + i * 32])) + k.emit(v_mul_f32_e32(v[8 + j77 * 1 + i * 4], s[38], v[8 + j77 * 1 + i * 4])) + for i in range(4): + k.emit(v_cvt_pk_bf16_f32(v[16 + i * 1], v[8 + i * 2], v[9 + i * 2])) + k.emit(s_nop(1)) + k.emit(v_permlane16_swap_b32_e32(v[16], v[18])) + k.emit(s_nop(1)) + k.emit(v_permlane16_swap_b32_e32(v[17], v[19])) + k.emit(s_nop(1)) + for i in range(4): + k.emit(buffer_atomic_pk_add_bf16(v[16 + i * 1], v[235], s[4:7], 0, 0 + i * 4, 1)) + k.emit(v_add_i32(v[235], v[235], 64)) + for i in range(2): + for j78 in range(4): + k.emit(v_accvgpr_read(v[8 + j78 * 1 + i * 4], v[64 + j78 * 1 + i * 32])) + k.emit(v_mul_f32_e32(v[8 + j78 * 1 + i * 4], s[38], v[8 + j78 * 1 + i * 4])) + for i in range(4): + k.emit(v_cvt_pk_bf16_f32(v[16 + i * 1], v[8 + i * 2], v[9 + i * 2])) + k.emit(s_nop(1)) + k.emit(v_permlane16_swap_b32_e32(v[16], v[18])) + k.emit(s_nop(1)) + k.emit(v_permlane16_swap_b32_e32(v[17], v[19])) + k.emit(s_nop(1)) + for i in range(4): + k.emit(buffer_atomic_pk_add_bf16(v[16 + i * 1], v[235], s[4:7], 0, 0 + i * 4, 1)) + k.emit(v_add_i32(v[235], v[235], 64)) + for i in range(2): + for j79 in range(4): + k.emit(v_accvgpr_read(v[8 + j79 * 1 + i * 4], v[4 + j79 * 1 + i * 32])) + k.emit(v_mul_f32_e32(v[8 + j79 * 1 + i * 4], s[38], v[8 + j79 * 1 + i * 4])) + for i in range(4): + k.emit(v_cvt_pk_bf16_f32(v[16 + i * 1], v[8 + i * 2], v[9 + i * 2])) + k.emit(s_nop(1)) + k.emit(v_permlane16_swap_b32_e32(v[16], v[18])) + k.emit(s_nop(1)) + k.emit(v_permlane16_swap_b32_e32(v[17], v[19])) + k.emit(s_nop(1)) + for i in range(4): + k.emit(buffer_atomic_pk_add_bf16(v[16 + i * 1], v[236], s[4:7], 0, 0 + i * 4, 1)) + k.emit(v_add_i32(v[236], v[236], 64)) + for i in range(2): + for j80 in range(4): + k.emit(v_accvgpr_read(v[8 + j80 * 1 + i * 4], v[68 + j80 * 1 + i * 32])) + k.emit(v_mul_f32_e32(v[8 + j80 * 1 + i * 4], s[38], v[8 + j80 * 1 + i * 4])) + for i in range(4): + k.emit(v_cvt_pk_bf16_f32(v[16 + i * 1], v[8 + i * 2], v[9 + i * 2])) + k.emit(s_nop(1)) + k.emit(v_permlane16_swap_b32_e32(v[16], v[18])) + k.emit(s_nop(1)) + k.emit(v_permlane16_swap_b32_e32(v[17], v[19])) + k.emit(s_nop(1)) + for i in range(4): + k.emit(buffer_atomic_pk_add_bf16(v[16 + i * 1], v[236], s[4:7], 0, 0 + i * 4, 1)) + k.emit(v_add_i32(v[236], v[236], 64)) + for i in range(2): + for j81 in range(4): + k.emit(v_accvgpr_read(v[8 + j81 * 1 + i * 4], v[8 + j81 * 1 + i * 32])) + k.emit(v_mul_f32_e32(v[8 + j81 * 1 + i * 4], s[38], v[8 + j81 * 1 + i * 4])) + for i in range(4): + k.emit(v_cvt_pk_bf16_f32(v[16 + i * 1], v[8 + i * 2], v[9 + i * 2])) + k.emit(s_nop(1)) + k.emit(v_permlane16_swap_b32_e32(v[16], v[18])) + k.emit(s_nop(1)) + k.emit(v_permlane16_swap_b32_e32(v[17], v[19])) + k.emit(s_nop(1)) + for i in range(4): + k.emit(buffer_atomic_pk_add_bf16(v[16 + i * 1], v[237], s[4:7], 0, 0 + i * 4, 1)) + k.emit(v_add_i32(v[237], v[237], 64)) + for i in range(2): + for j82 in range(4): + k.emit(v_accvgpr_read(v[8 + j82 * 1 + i * 4], v[72 + j82 * 1 + i * 32])) + k.emit(v_mul_f32_e32(v[8 + j82 * 1 + i * 4], s[38], v[8 + j82 * 1 + i * 4])) + for i in range(4): + k.emit(v_cvt_pk_bf16_f32(v[16 + i * 1], v[8 + i * 2], v[9 + i * 2])) + k.emit(s_nop(1)) + k.emit(v_permlane16_swap_b32_e32(v[16], v[18])) + k.emit(s_nop(1)) + k.emit(v_permlane16_swap_b32_e32(v[17], v[19])) + k.emit(s_nop(1)) + for i in range(4): + k.emit(buffer_atomic_pk_add_bf16(v[16 + i * 1], v[237], s[4:7], 0, 0 + i * 4, 1)) + k.emit(v_add_i32(v[237], v[237], 64)) + for i in range(2): + for j83 in range(4): + k.emit(v_accvgpr_read(v[8 + j83 * 1 + i * 4], v[12 + j83 * 1 + i * 32])) + k.emit(v_mul_f32_e32(v[8 + j83 * 1 + i * 4], s[38], v[8 + j83 * 1 + i * 4])) + for i in range(4): + k.emit(v_cvt_pk_bf16_f32(v[16 + i * 1], v[8 + i * 2], v[9 + i * 2])) + k.emit(s_nop(1)) + k.emit(v_permlane16_swap_b32_e32(v[16], v[18])) + k.emit(s_nop(1)) + k.emit(v_permlane16_swap_b32_e32(v[17], v[19])) + k.emit(s_nop(1)) + for i in range(4): + k.emit(buffer_atomic_pk_add_bf16(v[16 + i * 1], v[238], s[4:7], 0, 0 + i * 4, 1)) + k.emit(v_add_i32(v[238], v[238], 64)) + for i in range(2): + for j84 in range(4): + k.emit(v_accvgpr_read(v[8 + j84 * 1 + i * 4], v[76 + j84 * 1 + i * 32])) + k.emit(v_mul_f32_e32(v[8 + j84 * 1 + i * 4], s[38], v[8 + j84 * 1 + i * 4])) + for i in range(4): + k.emit(v_cvt_pk_bf16_f32(v[16 + i * 1], v[8 + i * 2], v[9 + i * 2])) + k.emit(s_nop(1)) + k.emit(v_permlane16_swap_b32_e32(v[16], v[18])) + k.emit(s_nop(1)) + k.emit(v_permlane16_swap_b32_e32(v[17], v[19])) + k.emit(s_nop(1)) + for i in range(4): + k.emit(buffer_atomic_pk_add_bf16(v[16 + i * 1], v[238], s[4:7], 0, 0 + i * 4, 1)) + k.emit(v_add_i32(v[238], v[238], 64)) + for i in range(2): + for j85 in range(4): + k.emit(v_accvgpr_read(v[8 + j85 * 1 + i * 4], v[16 + j85 * 1 + i * 32])) + k.emit(v_mul_f32_e32(v[8 + j85 * 1 + i * 4], s[38], v[8 + j85 * 1 + i * 4])) + for i in range(4): + k.emit(v_cvt_pk_bf16_f32(v[16 + i * 1], v[8 + i * 2], v[9 + i * 2])) + k.emit(s_nop(1)) + k.emit(v_permlane16_swap_b32_e32(v[16], v[18])) + k.emit(s_nop(1)) + k.emit(v_permlane16_swap_b32_e32(v[17], v[19])) + k.emit(s_nop(1)) + for i in range(4): + k.emit(buffer_atomic_pk_add_bf16(v[16 + i * 1], v[239], s[4:7], 0, 0 + i * 4, 1)) + k.emit(v_add_i32(v[239], v[239], 64)) + for i in range(2): + for j86 in range(4): + k.emit(v_accvgpr_read(v[8 + j86 * 1 + i * 4], v[80 + j86 * 1 + i * 32])) + k.emit(v_mul_f32_e32(v[8 + j86 * 1 + i * 4], s[38], v[8 + j86 * 1 + i * 4])) + for i in range(4): + k.emit(v_cvt_pk_bf16_f32(v[16 + i * 1], v[8 + i * 2], v[9 + i * 2])) + k.emit(s_nop(1)) + k.emit(v_permlane16_swap_b32_e32(v[16], v[18])) + k.emit(s_nop(1)) + k.emit(v_permlane16_swap_b32_e32(v[17], v[19])) + k.emit(s_nop(1)) + for i in range(4): + k.emit(buffer_atomic_pk_add_bf16(v[16 + i * 1], v[239], s[4:7], 0, 0 + i * 4, 1)) + k.emit(v_add_i32(v[239], v[239], 64)) + for i in range(2): + for j87 in range(4): + k.emit(v_accvgpr_read(v[8 + j87 * 1 + i * 4], v[20 + j87 * 1 + i * 32])) + k.emit(v_mul_f32_e32(v[8 + j87 * 1 + i * 4], s[38], v[8 + j87 * 1 + i * 4])) + for i in range(4): + k.emit(v_cvt_pk_bf16_f32(v[16 + i * 1], v[8 + i * 2], v[9 + i * 2])) + k.emit(s_nop(1)) + k.emit(v_permlane16_swap_b32_e32(v[16], v[18])) + k.emit(s_nop(1)) + k.emit(v_permlane16_swap_b32_e32(v[17], v[19])) + k.emit(s_nop(1)) + for i in range(4): + k.emit(buffer_atomic_pk_add_bf16(v[16 + i * 1], v[240], s[4:7], 0, 0 + i * 4, 1)) + k.emit(v_add_i32(v[240], v[240], 64)) + for i in range(2): + for j88 in range(4): + k.emit(v_accvgpr_read(v[8 + j88 * 1 + i * 4], v[84 + j88 * 1 + i * 32])) + k.emit(v_mul_f32_e32(v[8 + j88 * 1 + i * 4], s[38], v[8 + j88 * 1 + i * 4])) + for i in range(4): + k.emit(v_cvt_pk_bf16_f32(v[16 + i * 1], v[8 + i * 2], v[9 + i * 2])) + k.emit(s_nop(1)) + k.emit(v_permlane16_swap_b32_e32(v[16], v[18])) + k.emit(s_nop(1)) + k.emit(v_permlane16_swap_b32_e32(v[17], v[19])) + k.emit(s_nop(1)) + for i in range(4): + k.emit(buffer_atomic_pk_add_bf16(v[16 + i * 1], v[240], s[4:7], 0, 0 + i * 4, 1)) + k.emit(v_add_i32(v[240], v[240], 64)) + for i in range(2): + for j89 in range(4): + k.emit(v_accvgpr_read(v[8 + j89 * 1 + i * 4], v[24 + j89 * 1 + i * 32])) + k.emit(v_mul_f32_e32(v[8 + j89 * 1 + i * 4], s[38], v[8 + j89 * 1 + i * 4])) + for i in range(4): + k.emit(v_cvt_pk_bf16_f32(v[16 + i * 1], v[8 + i * 2], v[9 + i * 2])) + k.emit(s_nop(1)) + k.emit(v_permlane16_swap_b32_e32(v[16], v[18])) + k.emit(s_nop(1)) + k.emit(v_permlane16_swap_b32_e32(v[17], v[19])) + k.emit(s_nop(1)) + for i in range(4): + k.emit(buffer_atomic_pk_add_bf16(v[16 + i * 1], v[241], s[4:7], 0, 0 + i * 4, 1)) + k.emit(v_add_i32(v[241], v[241], 64)) + for i in range(2): + for j90 in range(4): + k.emit(v_accvgpr_read(v[8 + j90 * 1 + i * 4], v[88 + j90 * 1 + i * 32])) + k.emit(v_mul_f32_e32(v[8 + j90 * 1 + i * 4], s[38], v[8 + j90 * 1 + i * 4])) + for i in range(4): + k.emit(v_cvt_pk_bf16_f32(v[16 + i * 1], v[8 + i * 2], v[9 + i * 2])) + k.emit(s_nop(1)) + k.emit(v_permlane16_swap_b32_e32(v[16], v[18])) + k.emit(s_nop(1)) + k.emit(v_permlane16_swap_b32_e32(v[17], v[19])) + k.emit(s_nop(1)) + for i in range(4): + k.emit(buffer_atomic_pk_add_bf16(v[16 + i * 1], v[241], s[4:7], 0, 0 + i * 4, 1)) + k.emit(v_add_i32(v[241], v[241], 64)) + for i in range(2): + for j91 in range(4): + k.emit(v_accvgpr_read(v[8 + j91 * 1 + i * 4], v[28 + j91 * 1 + i * 32])) + k.emit(v_mul_f32_e32(v[8 + j91 * 1 + i * 4], s[38], v[8 + j91 * 1 + i * 4])) + for i in range(4): + k.emit(v_cvt_pk_bf16_f32(v[16 + i * 1], v[8 + i * 2], v[9 + i * 2])) + k.emit(s_nop(1)) + k.emit(v_permlane16_swap_b32_e32(v[16], v[18])) + k.emit(s_nop(1)) + k.emit(v_permlane16_swap_b32_e32(v[17], v[19])) + k.emit(s_nop(1)) + for i in range(4): + k.emit(buffer_atomic_pk_add_bf16(v[16 + i * 1], v[242], s[4:7], 0, 0 + i * 4, 1)) + k.emit(v_add_i32(v[242], v[242], 64)) + for i in range(2): + for j92 in range(4): + k.emit(v_accvgpr_read(v[8 + j92 * 1 + i * 4], v[92 + j92 * 1 + i * 32])) + k.emit(v_mul_f32_e32(v[8 + j92 * 1 + i * 4], s[38], v[8 + j92 * 1 + i * 4])) + for i in range(8): + for j93 in range(2): + k.emit(v_cvt_pk_bf16_f32(v[16], v[8], v[9])) + k.emit(v_cvt_pk_bf16_f32(v[17], v[10], v[11])) + k.emit(v_cvt_pk_bf16_f32(v[18], v[12], v[13])) + k.emit(v_cvt_pk_bf16_f32(v[19], v[14], v[15])) + k.emit(s_nop(1)) + k.emit(v_permlane16_swap_b32_e32(v[16], v[18])) + k.emit(s_nop(1)) + k.emit(v_permlane16_swap_b32_e32(v[17], v[19])) + k.emit(s_nop(1)) + k.emit(buffer_atomic_pk_add_bf16(v[16], v[242 + j93 * 1 + i * 1], s[4:7], 0, 0, 1)) + k.emit(buffer_atomic_pk_add_bf16(v[17], v[242 + j93 * 1 + i * 1], s[4:7], 0, 4, 1)) + k.emit(buffer_atomic_pk_add_bf16(v[18], v[242 + j93 * 1 + i * 1], s[4:7], 0, 8, 1)) + k.emit(buffer_atomic_pk_add_bf16(v[19], v[242 + j93 * 1 + i * 1], s[4:7], 0, 12, 1)) + k.emit(v_add_i32(v[242 + j93 * 1 + i * 1], v[242 + j93 * 1 + i * 1], 64)) + k.emit(v_accvgpr_read(v[8], v[128 + j93 * 64 + i * 4])) + k.emit(v_mul_f32_e32(v[8], s[38], v[8])) + k.emit(v_accvgpr_read(v[9], v[129 + j93 * 64 + i * 4])) + k.emit(v_mul_f32_e32(v[9], s[38], v[9])) + k.emit(v_accvgpr_read(v[10], v[130 + j93 * 64 + i * 4])) + k.emit(v_mul_f32_e32(v[10], s[38], v[10])) + k.emit(v_accvgpr_read(v[11], v[131 + j93 * 64 + i * 4])) + k.emit(v_mul_f32_e32(v[11], s[38], v[11])) + k.emit(v_accvgpr_read(v[12], v[160 + j93 * 64 + i * 4])) + k.emit(v_mul_f32_e32(v[12], s[38], v[12])) + k.emit(v_accvgpr_read(v[13], v[161 + j93 * 64 + i * 4])) + k.emit(v_mul_f32_e32(v[13], s[38], v[13])) + k.emit(v_accvgpr_read(v[14], v[162 + j93 * 64 + i * 4])) + k.emit(v_mul_f32_e32(v[14], s[38], v[14])) + k.emit(v_accvgpr_read(v[15], v[163 + j93 * 64 + i * 4])) + k.emit(v_mul_f32_e32(v[15], s[38], v[15])) + for i in range(4): + k.emit(v_cvt_pk_bf16_f32(v[16 + i * 1], v[8 + i * 2], v[9 + i * 2])) + k.emit(s_nop(1)) + k.emit(v_permlane16_swap_b32_e32(v[16], v[18])) + k.emit(s_nop(1)) + k.emit(v_permlane16_swap_b32_e32(v[17], v[19])) + k.emit(s_nop(1)) + for i in range(4): + k.emit(buffer_atomic_pk_add_bf16(v[16 + i * 1], v[250], s[4:7], 0, 0 + i * 4, 1)) + k.emit(v_add_i32(v[250], v[250], 64)) + k.emit(s_branch(1312), target='L2_6724') + k.label('L2_52A4') + for i in range(2): + for j94 in range(4): + k.emit(v_accvgpr_read(v[8 + j94 * 1 + i * 4], v[0 + j94 * 1 + i * 32])) + k.emit(v_mul_f32_e32(v[8 + j94 * 1 + i * 4], s[38], v[8 + j94 * 1 + i * 4])) + for i in range(4): + k.emit(v_cvt_pk_bf16_f32(v[16 + i * 1], v[8 + i * 2], v[9 + i * 2])) + k.emit(s_nop(1)) + k.emit(v_permlane16_swap_b32_e32(v[16], v[18])) + k.emit(s_nop(1)) + k.emit(v_permlane16_swap_b32_e32(v[17], v[19])) + k.emit(s_nop(1)) + k.emit(buffer_store_dwordx4(v[16:19], v[235], s[4:7], 0, 0, 1)) + k.emit(v_add_i32(v[235], v[235], 64)) + for i in range(2): + for j95 in range(4): + k.emit(v_accvgpr_read(v[8 + j95 * 1 + i * 4], v[64 + j95 * 1 + i * 32])) + k.emit(v_mul_f32_e32(v[8 + j95 * 1 + i * 4], s[38], v[8 + j95 * 1 + i * 4])) + for i in range(4): + k.emit(v_cvt_pk_bf16_f32(v[16 + i * 1], v[8 + i * 2], v[9 + i * 2])) + k.emit(s_nop(1)) + k.emit(v_permlane16_swap_b32_e32(v[16], v[18])) + k.emit(s_nop(1)) + k.emit(v_permlane16_swap_b32_e32(v[17], v[19])) + k.emit(s_nop(1)) + k.emit(buffer_store_dwordx4(v[16:19], v[235], s[4:7], 0, 0, 1)) + k.emit(v_add_i32(v[235], v[235], 64)) + for i in range(2): + for j96 in range(4): + k.emit(v_accvgpr_read(v[8 + j96 * 1 + i * 4], v[4 + j96 * 1 + i * 32])) + k.emit(v_mul_f32_e32(v[8 + j96 * 1 + i * 4], s[38], v[8 + j96 * 1 + i * 4])) + for i in range(4): + k.emit(v_cvt_pk_bf16_f32(v[16 + i * 1], v[8 + i * 2], v[9 + i * 2])) + k.emit(s_nop(1)) + k.emit(v_permlane16_swap_b32_e32(v[16], v[18])) + k.emit(s_nop(1)) + k.emit(v_permlane16_swap_b32_e32(v[17], v[19])) + k.emit(s_nop(1)) + k.emit(buffer_store_dwordx4(v[16:19], v[236], s[4:7], 0, 0, 1)) + k.emit(v_add_i32(v[236], v[236], 64)) + for i in range(2): + for j97 in range(4): + k.emit(v_accvgpr_read(v[8 + j97 * 1 + i * 4], v[68 + j97 * 1 + i * 32])) + k.emit(v_mul_f32_e32(v[8 + j97 * 1 + i * 4], s[38], v[8 + j97 * 1 + i * 4])) + for i in range(4): + k.emit(v_cvt_pk_bf16_f32(v[16 + i * 1], v[8 + i * 2], v[9 + i * 2])) + k.emit(s_nop(1)) + k.emit(v_permlane16_swap_b32_e32(v[16], v[18])) + k.emit(s_nop(1)) + k.emit(v_permlane16_swap_b32_e32(v[17], v[19])) + k.emit(s_nop(1)) + k.emit(buffer_store_dwordx4(v[16:19], v[236], s[4:7], 0, 0, 1)) + k.emit(v_add_i32(v[236], v[236], 64)) + for i in range(2): + for j98 in range(4): + k.emit(v_accvgpr_read(v[8 + j98 * 1 + i * 4], v[8 + j98 * 1 + i * 32])) + k.emit(v_mul_f32_e32(v[8 + j98 * 1 + i * 4], s[38], v[8 + j98 * 1 + i * 4])) + for i in range(4): + k.emit(v_cvt_pk_bf16_f32(v[16 + i * 1], v[8 + i * 2], v[9 + i * 2])) + k.emit(s_nop(1)) + k.emit(v_permlane16_swap_b32_e32(v[16], v[18])) + k.emit(s_nop(1)) + k.emit(v_permlane16_swap_b32_e32(v[17], v[19])) + k.emit(s_nop(1)) + k.emit(buffer_store_dwordx4(v[16:19], v[237], s[4:7], 0, 0, 1)) + k.emit(v_add_i32(v[237], v[237], 64)) + for i in range(2): + for j99 in range(4): + k.emit(v_accvgpr_read(v[8 + j99 * 1 + i * 4], v[72 + j99 * 1 + i * 32])) + k.emit(v_mul_f32_e32(v[8 + j99 * 1 + i * 4], s[38], v[8 + j99 * 1 + i * 4])) + for i in range(4): + k.emit(v_cvt_pk_bf16_f32(v[16 + i * 1], v[8 + i * 2], v[9 + i * 2])) + k.emit(s_nop(1)) + k.emit(v_permlane16_swap_b32_e32(v[16], v[18])) + k.emit(s_nop(1)) + k.emit(v_permlane16_swap_b32_e32(v[17], v[19])) + k.emit(s_nop(1)) + k.emit(buffer_store_dwordx4(v[16:19], v[237], s[4:7], 0, 0, 1)) + k.emit(v_add_i32(v[237], v[237], 64)) + for i in range(2): + for j100 in range(4): + k.emit(v_accvgpr_read(v[8 + j100 * 1 + i * 4], v[12 + j100 * 1 + i * 32])) + k.emit(v_mul_f32_e32(v[8 + j100 * 1 + i * 4], s[38], v[8 + j100 * 1 + i * 4])) + for i in range(4): + k.emit(v_cvt_pk_bf16_f32(v[16 + i * 1], v[8 + i * 2], v[9 + i * 2])) + k.emit(s_nop(1)) + k.emit(v_permlane16_swap_b32_e32(v[16], v[18])) + k.emit(s_nop(1)) + k.emit(v_permlane16_swap_b32_e32(v[17], v[19])) + k.emit(s_nop(1)) + k.emit(buffer_store_dwordx4(v[16:19], v[238], s[4:7], 0, 0, 1)) + k.emit(v_add_i32(v[238], v[238], 64)) + for i in range(2): + for j101 in range(4): + k.emit(v_accvgpr_read(v[8 + j101 * 1 + i * 4], v[76 + j101 * 1 + i * 32])) + k.emit(v_mul_f32_e32(v[8 + j101 * 1 + i * 4], s[38], v[8 + j101 * 1 + i * 4])) + for i in range(4): + k.emit(v_cvt_pk_bf16_f32(v[16 + i * 1], v[8 + i * 2], v[9 + i * 2])) + k.emit(s_nop(1)) + k.emit(v_permlane16_swap_b32_e32(v[16], v[18])) + k.emit(s_nop(1)) + k.emit(v_permlane16_swap_b32_e32(v[17], v[19])) + k.emit(s_nop(1)) + k.emit(buffer_store_dwordx4(v[16:19], v[238], s[4:7], 0, 0, 1)) + k.emit(v_add_i32(v[238], v[238], 64)) + for i in range(2): + for j102 in range(4): + k.emit(v_accvgpr_read(v[8 + j102 * 1 + i * 4], v[16 + j102 * 1 + i * 32])) + k.emit(v_mul_f32_e32(v[8 + j102 * 1 + i * 4], s[38], v[8 + j102 * 1 + i * 4])) + for i in range(4): + k.emit(v_cvt_pk_bf16_f32(v[16 + i * 1], v[8 + i * 2], v[9 + i * 2])) + k.emit(s_nop(1)) + k.emit(v_permlane16_swap_b32_e32(v[16], v[18])) + k.emit(s_nop(1)) + k.emit(v_permlane16_swap_b32_e32(v[17], v[19])) + k.emit(s_nop(1)) + k.emit(buffer_store_dwordx4(v[16:19], v[239], s[4:7], 0, 0, 1)) + k.emit(v_add_i32(v[239], v[239], 64)) + for i in range(2): + for j103 in range(4): + k.emit(v_accvgpr_read(v[8 + j103 * 1 + i * 4], v[80 + j103 * 1 + i * 32])) + k.emit(v_mul_f32_e32(v[8 + j103 * 1 + i * 4], s[38], v[8 + j103 * 1 + i * 4])) + for i in range(4): + k.emit(v_cvt_pk_bf16_f32(v[16 + i * 1], v[8 + i * 2], v[9 + i * 2])) + k.emit(s_nop(1)) + k.emit(v_permlane16_swap_b32_e32(v[16], v[18])) + k.emit(s_nop(1)) + k.emit(v_permlane16_swap_b32_e32(v[17], v[19])) + k.emit(s_nop(1)) + k.emit(buffer_store_dwordx4(v[16:19], v[239], s[4:7], 0, 0, 1)) + k.emit(v_add_i32(v[239], v[239], 64)) + for i in range(2): + for j104 in range(4): + k.emit(v_accvgpr_read(v[8 + j104 * 1 + i * 4], v[20 + j104 * 1 + i * 32])) + k.emit(v_mul_f32_e32(v[8 + j104 * 1 + i * 4], s[38], v[8 + j104 * 1 + i * 4])) + for i in range(4): + k.emit(v_cvt_pk_bf16_f32(v[16 + i * 1], v[8 + i * 2], v[9 + i * 2])) + k.emit(s_nop(1)) + k.emit(v_permlane16_swap_b32_e32(v[16], v[18])) + k.emit(s_nop(1)) + k.emit(v_permlane16_swap_b32_e32(v[17], v[19])) + k.emit(s_nop(1)) + k.emit(buffer_store_dwordx4(v[16:19], v[240], s[4:7], 0, 0, 1)) + k.emit(v_add_i32(v[240], v[240], 64)) + for i in range(2): + for j105 in range(4): + k.emit(v_accvgpr_read(v[8 + j105 * 1 + i * 4], v[84 + j105 * 1 + i * 32])) + k.emit(v_mul_f32_e32(v[8 + j105 * 1 + i * 4], s[38], v[8 + j105 * 1 + i * 4])) + for i in range(4): + k.emit(v_cvt_pk_bf16_f32(v[16 + i * 1], v[8 + i * 2], v[9 + i * 2])) + k.emit(s_nop(1)) + k.emit(v_permlane16_swap_b32_e32(v[16], v[18])) + k.emit(s_nop(1)) + k.emit(v_permlane16_swap_b32_e32(v[17], v[19])) + k.emit(s_nop(1)) + k.emit(buffer_store_dwordx4(v[16:19], v[240], s[4:7], 0, 0, 1)) + k.emit(v_add_i32(v[240], v[240], 64)) + for i in range(2): + for j106 in range(4): + k.emit(v_accvgpr_read(v[8 + j106 * 1 + i * 4], v[24 + j106 * 1 + i * 32])) + k.emit(v_mul_f32_e32(v[8 + j106 * 1 + i * 4], s[38], v[8 + j106 * 1 + i * 4])) + for i in range(4): + k.emit(v_cvt_pk_bf16_f32(v[16 + i * 1], v[8 + i * 2], v[9 + i * 2])) + k.emit(s_nop(1)) + k.emit(v_permlane16_swap_b32_e32(v[16], v[18])) + k.emit(s_nop(1)) + k.emit(v_permlane16_swap_b32_e32(v[17], v[19])) + k.emit(s_nop(1)) + k.emit(buffer_store_dwordx4(v[16:19], v[241], s[4:7], 0, 0, 1)) + k.emit(v_add_i32(v[241], v[241], 64)) + for i in range(2): + for j107 in range(4): + k.emit(v_accvgpr_read(v[8 + j107 * 1 + i * 4], v[88 + j107 * 1 + i * 32])) + k.emit(v_mul_f32_e32(v[8 + j107 * 1 + i * 4], s[38], v[8 + j107 * 1 + i * 4])) + for i in range(4): + k.emit(v_cvt_pk_bf16_f32(v[16 + i * 1], v[8 + i * 2], v[9 + i * 2])) + k.emit(s_nop(1)) + k.emit(v_permlane16_swap_b32_e32(v[16], v[18])) + k.emit(s_nop(1)) + k.emit(v_permlane16_swap_b32_e32(v[17], v[19])) + k.emit(s_nop(1)) + k.emit(buffer_store_dwordx4(v[16:19], v[241], s[4:7], 0, 0, 1)) + k.emit(v_add_i32(v[241], v[241], 64)) + for i in range(2): + for j108 in range(4): + k.emit(v_accvgpr_read(v[8 + j108 * 1 + i * 4], v[28 + j108 * 1 + i * 32])) + k.emit(v_mul_f32_e32(v[8 + j108 * 1 + i * 4], s[38], v[8 + j108 * 1 + i * 4])) + for i in range(4): + k.emit(v_cvt_pk_bf16_f32(v[16 + i * 1], v[8 + i * 2], v[9 + i * 2])) + k.emit(s_nop(1)) + k.emit(v_permlane16_swap_b32_e32(v[16], v[18])) + k.emit(s_nop(1)) + k.emit(v_permlane16_swap_b32_e32(v[17], v[19])) + k.emit(s_nop(1)) + k.emit(buffer_store_dwordx4(v[16:19], v[242], s[4:7], 0, 0, 1)) + k.emit(v_add_i32(v[242], v[242], 64)) + for i in range(2): + for j109 in range(4): + k.emit(v_accvgpr_read(v[8 + j109 * 1 + i * 4], v[92 + j109 * 1 + i * 32])) + k.emit(v_mul_f32_e32(v[8 + j109 * 1 + i * 4], s[38], v[8 + j109 * 1 + i * 4])) + for i in range(8): + for j110 in range(2): + k.emit(v_cvt_pk_bf16_f32(v[16], v[8], v[9])) + k.emit(v_cvt_pk_bf16_f32(v[17], v[10], v[11])) + k.emit(v_cvt_pk_bf16_f32(v[18], v[12], v[13])) + k.emit(v_cvt_pk_bf16_f32(v[19], v[14], v[15])) + k.emit(s_nop(1)) + k.emit(v_permlane16_swap_b32_e32(v[16], v[18])) + k.emit(s_nop(1)) + k.emit(v_permlane16_swap_b32_e32(v[17], v[19])) + k.emit(s_nop(1)) + k.emit(buffer_store_dwordx4(v[16:19], v[242 + j110 * 1 + i * 1], s[4:7], 0, 0, 1)) + k.emit(v_add_i32(v[242 + j110 * 1 + i * 1], v[242 + j110 * 1 + i * 1], 64)) + k.emit(v_accvgpr_read(v[8], v[128 + j110 * 64 + i * 4])) + k.emit(v_mul_f32_e32(v[8], s[38], v[8])) + k.emit(v_accvgpr_read(v[9], v[129 + j110 * 64 + i * 4])) + k.emit(v_mul_f32_e32(v[9], s[38], v[9])) + k.emit(v_accvgpr_read(v[10], v[130 + j110 * 64 + i * 4])) + k.emit(v_mul_f32_e32(v[10], s[38], v[10])) + k.emit(v_accvgpr_read(v[11], v[131 + j110 * 64 + i * 4])) + k.emit(v_mul_f32_e32(v[11], s[38], v[11])) + k.emit(v_accvgpr_read(v[12], v[160 + j110 * 64 + i * 4])) + k.emit(v_mul_f32_e32(v[12], s[38], v[12])) + k.emit(v_accvgpr_read(v[13], v[161 + j110 * 64 + i * 4])) + k.emit(v_mul_f32_e32(v[13], s[38], v[13])) + k.emit(v_accvgpr_read(v[14], v[162 + j110 * 64 + i * 4])) + k.emit(v_mul_f32_e32(v[14], s[38], v[14])) + k.emit(v_accvgpr_read(v[15], v[163 + j110 * 64 + i * 4])) + k.emit(v_mul_f32_e32(v[15], s[38], v[15])) + for i in range(4): + k.emit(v_cvt_pk_bf16_f32(v[16 + i * 1], v[8 + i * 2], v[9 + i * 2])) + k.emit(s_nop(1)) + k.emit(v_permlane16_swap_b32_e32(v[16], v[18])) + k.emit(s_nop(1)) + k.emit(v_permlane16_swap_b32_e32(v[17], v[19])) + k.emit(s_nop(1)) + k.emit(buffer_store_dwordx4(v[16:19], v[250], s[4:7], 0, 0, 1)) + k.emit(v_add_i32(v[250], v[250], 64)) + k.label('L2_6724') + k.emit(s_waitcnt()) + k.emit(s_endpgm()) + else: + raise AssertionError(f'unsupported tile {(tile_m, tile_n)}') + return k.finalize() diff --git a/test/backend/test_asm_gemm.py b/test/backend/test_asm_gemm.py index db729bc91d..9044f849fc 100644 --- a/test/backend/test_asm_gemm.py +++ b/test/backend/test_asm_gemm.py @@ -150,6 +150,19 @@ class TestAsmGEMM(unittest.TestCase): with self.assertRaisesRegex(AssertionError, "not a multiple"): verify_asm_gemm(1, 256, 1000, 256) +class TestMXFP4(unittest.TestCase): + def setUp(self): + if not is_cdna4() or DEV.interface.startswith("MOCK"): + self.skipTest("requires real amd machine") + + def test_empty(self): + M, N, K = getenv("M", 16384), getenv("N", 4096), getenv("K", 14336) + a = Tensor.empty(M, K // 2, dtype=dtypes.uint8) + b = Tensor.empty(N, K // 2, dtype=dtypes.uint8) + scale_a = Tensor.empty(M, K // 32, dtype=dtypes.uint8) + scale_b = Tensor.empty(N, K // 32, dtype=dtypes.uint8) + asm_gemm(a, b.T, mx_scales=(scale_a, scale_b)).realize() + # test the Asm GEMM with Llama shapes, only run on the real machine for speed @unittest.skipUnless(has_hipcc(), "requires hipcc to compile") From 13452b3775d3d60fc802f2c68841a97578cb3b11 Mon Sep 17 00:00:00 2001 From: Christopher Milan Date: Fri, 31 Jul 2026 01:37:34 -0400 Subject: [PATCH 13/44] benchmark comma big model (#17312) --- .github/workflows/benchmark.yml | 12 ++++----- examples/openpilot/compile3.py | 42 +++++++++++++++++++++++++++---- examples/openpilot/load_pickle.py | 5 ++-- 3 files changed, 46 insertions(+), 13 deletions(-) diff --git a/.github/workflows/benchmark.yml b/.github/workflows/benchmark.yml index 04d314e837..167dced5ae 100644 --- a/.github/workflows/benchmark.yml +++ b/.github/workflows/benchmark.yml @@ -517,12 +517,12 @@ jobs: rm -f /tmp/staging.db /tmp/staging.db-shm /tmp/staging.db-wal - name: reset chestnut run: python3 extra/usbgpu/debug.py -rn - - name: openpilot compile3 0.10.1 driving_vision - run: BENCHMARK_LOG=usbgpu_openpilot_0_10_1_vision PYTHONPATH="." GMMU=0 DEV=USB+AMD:LLVM ASSERT_MIN_STEP_TIME=50 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_vision.onnx - - name: openpilot load_pickle 0.10.1 driving_vision - run: BENCHMARK_LOG=usbgpu_openpilot_0_10_1_vision_load_pickle PYTHONPATH="." GMMU=0 DEV=USB+AMD ASSERT_MIN_LOAD_TIME=15 python3 examples/openpilot/load_pickle.py - - name: openpilot run_pickle 0.10.1 driving_vision - run: BENCHMARK_LOG=usbgpu_openpilot_0_10_1_vision_run_pickle RUN_PICKLE=1 PYTHONPATH="." GMMU=0 DEV=USB+AMD ASSERT_MIN_STEP_TIME=50 python3 examples/openpilot/compile3.py + - name: openpilot compile3 big_driving_supercombo + run: BENCHMARK_LOG=usbgpu_openpilot_big_driving_supercombo PICKLE_OOB=1 PYTHONPATH="." GMMU=0 DEV=USB+AMD:LLVM ASSERT_MIN_STEP_TIME=50 python3 examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/10926f2c0911821ca0e72439c1c3bf3ec11f0a08789aa14b7ee8f25379b2afa4 openpilot.pkl + - name: openpilot load_pickle big_driving_supercombo + run: BENCHMARK_LOG=usbgpu_openpilot_big_driving_supercombo_load_pickle PICKLE_OOB=1 PYTHONPATH="." GMMU=0 DEV=USB+AMD ASSERT_MIN_LOAD_TIME=25 python3 examples/openpilot/load_pickle.py openpilot.pkl + - name: openpilot run_pickle big_driving_supercombo + run: BENCHMARK_LOG=usbgpu_openpilot_big_driving_supercombo_run_pickle RUN_PICKLE=1 PICKLE_OOB=1 PYTHONPATH="." GMMU=0 DEV=USB+AMD ASSERT_MIN_STEP_TIME=50 python3 examples/openpilot/compile3.py - openpilot.pkl - name: Test copy speeds run: SIZE=64e6 PYTHONPATH=. GMMU=0 DEV=USB+AMD python3 test/external/external_test_usb_asm24.py TestDevCopySpeeds diff --git a/examples/openpilot/compile3.py b/examples/openpilot/compile3.py index 3c5a587b78..7bf253fd8d 100644 --- a/examples/openpilot/compile3.py +++ b/examples/openpilot/compile3.py @@ -1,4 +1,4 @@ -import os, sys, pickle, time, re +import os, sys, pickle, time, re, tempfile, struct, shutil, io import numpy as np if "JIT_BATCH_SIZE" not in os.environ: os.environ["JIT_BATCH_SIZE"] = "0" @@ -9,6 +9,39 @@ from tinygrad.nn.onnx import OnnxRunner OPENPILOT_MODEL = sys.argv[1] if len(sys.argv) > 1 else "https://github.com/commaai/openpilot/raw/v0.9.7/selfdrive/modeld/models/supercombo.onnx" OUTPUT = sys.argv[2] if len(sys.argv) > 2 else "/tmp/openpilot.pkl" +PICKLE_OOB = getenv("PICKLE_OOB") + +def dump_pickle(obj, f): + if PICKLE_OOB: + # allows pickling when buffers don't fit in (CPU) RAM + # from openpilot/selfdrive/modeld/helpers.py + with tempfile.TemporaryFile(dir=".") as tmp: + def buffer_callback(pb: pickle.PickleBuffer): + m = pb.raw() + tmp.write(struct.pack(' 1 else "/tmp/openpilot.pkl" @@ -7,7 +8,7 @@ PKL = sys.argv[1] if len(sys.argv) > 1 else "/tmp/openpilot.pkl" load_times = [] for _ in range(10): - with WallTimeEvent(BenchEvent.STEP) as wte: pickle.load(open(PKL, 'rb')) + with WallTimeEvent(BenchEvent.STEP) as wte: load_pickle(open(PKL, 'rb')) load_times.append(wte.time) print(f"pickle load: {wte.time:6.2f} s") From 0a3325f9c260775e688de4be701d82519247a943 Mon Sep 17 00:00:00 2001 From: qazal <77887910+Qazalin@users.noreply.github.com> Date: Fri, 31 Jul 2026 13:50:18 +0800 Subject: [PATCH 14/44] add mxfp4 quantize and layout kernels (#17320) --- extra/gemm/cdna_asm_gemm.py | 52 +++++++++++++++++++++++++++-------- test/backend/test_asm_gemm.py | 38 +++++++++++++++++++++---- 2 files changed, 73 insertions(+), 17 deletions(-) diff --git a/extra/gemm/cdna_asm_gemm.py b/extra/gemm/cdna_asm_gemm.py index 31c75c9513..feca12a051 100644 --- a/extra/gemm/cdna_asm_gemm.py +++ b/extra/gemm/cdna_asm_gemm.py @@ -137,6 +137,38 @@ def quantize_mxfp8(x:Tensor) -> tuple[Tensor, Tensor, Tensor]: packed = mx_pack(e8) if len(batch) == 1 and scale_K % 4 == 0 else None return x_clamped.cast(FP8_DTYPE), e8, packed +def _mxfp4_shuffle_weight(x:Tensor) -> Tensor: + # shuffle_weight(x, layout=(16, 16)) on the packed uint8 buffer. + rows, half_k = x.shape + return x.reshape(rows//16, 16, half_k//32, 2, 16).permute(0, 2, 3, 1, 4).reshape(rows, half_k).contiguous() + +def _mxfp4_shuffle_scales(x:Tensor) -> Tensor: + # e8m0_shuffle: each 256x8 scale tile is arranged for the raw MFMA scale loads. + rows, scale_k = x.shape + return x.reshape(rows//32, 2, 16, scale_k//8, 2, 4).permute(0, 3, 5, 2, 4, 1).reshape(rows, scale_k).contiguous() + +def quantize_mxfp4(x:Tensor) -> tuple[Tensor, Tensor, Tensor]: + # OCP MXFP4: 1x32 blocks, E2M1 values packed low-nibble first, and E8M0 scales. + assert x.ndim == 2 and x.shape[1] % 256 == 0 and x.shape[0] % 32 == 0, \ + f"mxfp4 quantization needs rows%32 and K%256, got {x.shape}" + rows, K = x.shape + xb = x.float().reshape(rows, K//32, 32) + amax = xb.abs().max(axis=-1) + + # even scale rounding: round the fp32 significand before choosing 2^(floor(log2)-2). + amax_rounded = ((amax.bitcast(dtypes.uint32) + 0x200000) & 0xFF800000).bitcast(dtypes.float32) + scale_exp = (amax_rounded.maximum(2**-126).log2().floor() - 2).clamp(-127, 127) + e8 = (scale_exp + 127).cast(dtypes.uint8) + scaled = xb * (-scale_exp).exp2().reshape(rows, K//32, 1) + + mag = scaled.abs() + code = sum(x.cast(dtypes.uint8) for x in + (mag > .25, mag >= .75, mag > 1.25, mag >= 1.75, mag > 2.5, mag >= 3.5, mag > 5.0)) + code = code | ((scaled < 0).cast(dtypes.uint8) << 3) + code = code.reshape(rows, K) + packed = code[:, 0::2] | (code[:, 1::2] << 4) + return packed, e8, _mxfp4_shuffle_scales(e8) + def mx_pack(e8:Tensor) -> Tensor: rows, scale_K = e8.shape return e8.reshape(rows, scale_K // 4, 4).bitcast(dtypes.uint32).reshape(rows, scale_K // 4).permute(1, 0).contiguous() @@ -162,9 +194,7 @@ atexit.register(_asm_gemm_report) def can_use_asm_gemm(a:Tensor, b:Tensor) -> bool: if a.dtype != b.dtype: return todo(f"dtypes must match {a.dtype} != {b.dtype}") - # fp4 encoded as packed uint8 - # TODO: add fp4 dtype? - if a.dtype not in {dtypes.bfloat16, dtypes.float16, FP8_DTYPE, dtypes.uint8}: return todo(f"only bfloat16/float16/fp8/fp4, got {a.dtype}") + if a.dtype not in {dtypes.bfloat16, dtypes.float16, FP8_DTYPE}: return todo(f"only bfloat16/float16/fp8, got {a.dtype}") batch, M, K = (1, *a.shape) if a.ndim == 2 else a.shape N = b.shape[1] if isinstance(a.device, tuple): @@ -365,14 +395,11 @@ def custom_mx_gemm_bw(gradient:UOp, kernel:UOp, has_w_post:bool, w_stored:bool=F def asm_gemm(a:Tensor, b:Tensor, x_scale:Tensor|None=None, w_scale:Tensor|None=None, grad_amax_state:Tensor|None=None, next_grad_amax_state:Tensor|None=None, w_post_scale:Tensor|None=None, mx:bool=False, mx_scales:tuple|None=None, mx_w_stored:bool=False, g_amax:Tensor|None=None, - a_pretranspose:Tensor|None=None) -> Tensor: + a_pretranspose:Tensor|None=None, mxfp4:bool=False) -> Tensor: assert can_use_asm_gemm(a, b), f"{counters['todos'][-1]}" - if (mxfp4:=a.dtype == dtypes.uint8): - assert mx_scales is not None and len(mx_scales) == 2 - scale_a, scale_b = mx_scales - K = a.shape[-1] * 2 - assert scale_a.shape == (*a.shape[:-1], K // 32) and scale_b.shape == (b.shape[1], K // 32) - assert scale_a.dtype == scale_b.dtype == dtypes.uint8 and a.device == b.device == scale_a.device == scale_b.device + if mxfp4: + assert not mx and mx_scales is None, "mxfp4 owns quantization; mx/mx_scales are for mxfp8" + assert a.dtype == dtypes.bfloat16, f"cannot quantize {a.dtype} to mxfp4" counters["used"] += 1 unfold_batch = a.ndim == 3 and isinstance(a.device, tuple) and a.uop.axis == 2 and b.uop.axis == 0 if unfold_batch: @@ -406,7 +433,10 @@ def asm_gemm(a:Tensor, b:Tensor, x_scale:Tensor|None=None, w_scale:Tensor|None=N if mxfp4: tile_m, tile_n = next((tm, tn) for tm, tn in ((256, 256), (192, 256), (128, 512)) if (batch*M) % tm == N % tn == 0) fxn = functools.partial(custom_mxfp4_gemm, tile_m=tile_m, tile_n=tile_n) - out = Tensor.custom_kernel(out, a, b.T, scale_a, scale_b, fxn=fxn)[0] + a_q, _, scale_a = quantize_mxfp4(a.reshape(batch*M, K)) + b_q, _, scale_b = quantize_mxfp4(b.T) + a_q, b_q = a_q.reshape(batch, M, K//2).contiguous(), _mxfp4_shuffle_weight(b_q) + out = Tensor.custom_kernel(out, a_q, b_q, scale_a, scale_b, fxn=fxn)[0] elif mx: # mxfp8 1x32 block scaling if mx_scales is not None: diff --git a/test/backend/test_asm_gemm.py b/test/backend/test_asm_gemm.py index 9044f849fc..53c148c862 100644 --- a/test/backend/test_asm_gemm.py +++ b/test/backend/test_asm_gemm.py @@ -1,7 +1,7 @@ import unittest from tinygrad import Tensor, Device, dtypes, Context from tinygrad.helpers import getenv, system, DEV -from extra.gemm.cdna_asm_gemm import asm_gemm, hk_bf16_atb_gemm +from extra.gemm.cdna_asm_gemm import asm_gemm, hk_bf16_atb_gemm, quantize_mxfp4 from test.helpers import needs_second_gpu from examples.mlperf.models.flat_llama import FP8_DTYPE, quantize_fp8, FP8_MAX @@ -155,13 +155,39 @@ class TestMXFP4(unittest.TestCase): if not is_cdna4() or DEV.interface.startswith("MOCK"): self.skipTest("requires real amd machine") + def test_quantize(self): + import numpy as np + block = np.array([0, .26, .74, .75, 1.26, 1.75, 2.51, 3.5, 5.1, 6, -6] + [0] * 21, dtype=np.float32) + x = Tensor(np.tile(block, (32, 8)), dtype=dtypes.bfloat16) + packed, scale, _ = quantize_mxfp4(x) + p = packed.numpy() + codes = np.stack((p & 0xF, p >> 4), axis=-1).reshape(32, 256) + np.testing.assert_array_equal(codes[0, :11], [0, 1, 1, 2, 3, 4, 5, 6, 7, 7, 15]) + np.testing.assert_array_equal(scale.numpy(), np.full((32, 8), 127, dtype=np.uint8)) + + def test_correctness(self): + import numpy as np + M = N = K = 256 + rng = np.random.default_rng(1) + a = Tensor(rng.standard_normal((M, K), dtype=np.float32), dtype=dtypes.bfloat16) + b = Tensor(rng.standard_normal((N, K), dtype=np.float32), dtype=dtypes.bfloat16) + out = asm_gemm(a, b.T, mxfp4=True).realize() + # reference gemm + a_packed, scale_a, _ = quantize_mxfp4(a) + b_packed, scale_b, _ = quantize_mxfp4(b) + def unpack(x): return np.stack((x & 0xF, x >> 4), axis=-1).reshape(x.shape[0], -1) + code_a, code_b = unpack(a_packed.numpy()), unpack(b_packed.numpy()) + lut = np.array([0, .5, 1, 1.5, 2, 3, 4, 6, -0., -.5, -1, -1.5, -2, -3, -4, -6], dtype=np.float32) + a_dequant = lut[code_a] * np.repeat(np.exp2(scale_a.numpy().astype(np.int16)-127), 32, axis=1) + b_dequant = lut[code_b] * np.repeat(np.exp2(scale_b.numpy().astype(np.int16)-127), 32, axis=1) + ref = Tensor(a_dequant @ b_dequant.T, dtype=dtypes.bfloat16).realize().numpy() + np.testing.assert_array_equal(out.numpy(), ref) + def test_empty(self): M, N, K = getenv("M", 16384), getenv("N", 4096), getenv("K", 14336) - a = Tensor.empty(M, K // 2, dtype=dtypes.uint8) - b = Tensor.empty(N, K // 2, dtype=dtypes.uint8) - scale_a = Tensor.empty(M, K // 32, dtype=dtypes.uint8) - scale_b = Tensor.empty(N, K // 32, dtype=dtypes.uint8) - asm_gemm(a, b.T, mx_scales=(scale_a, scale_b)).realize() + a = Tensor.empty(M, K, dtype=dtypes.bfloat16) + b = Tensor.empty(N, K, dtype=dtypes.bfloat16) + asm_gemm(a, b.T, mxfp4=True).realize() # test the Asm GEMM with Llama shapes, only run on the real machine for speed From f7964acb645b204390e60f28bbba7acf19777009 Mon Sep 17 00:00:00 2001 From: qazal <77887910+Qazalin@users.noreply.github.com> Date: Fri, 31 Jul 2026 17:25:25 +0800 Subject: [PATCH 15/44] llama with MXFP4 (#17321) * mxfp4 in llama * less * name --- examples/mlperf/models/flat_llama.py | 15 ++++-- extra/gemm/cdna_asm_gemm.py | 68 ++++++++++++++++++++-------- 2 files changed, 62 insertions(+), 21 deletions(-) diff --git a/examples/mlperf/models/flat_llama.py b/examples/mlperf/models/flat_llama.py index f93a752ddb..8cbc0e9223 100644 --- a/examples/mlperf/models/flat_llama.py +++ b/examples/mlperf/models/flat_llama.py @@ -25,6 +25,7 @@ FUSED_SILU_W13 = getenv("FUSED_SILU_W13", 0) SPLIT_W13 = getenv("SPLIT_W13", 0) COLUMNWISE_WEIGHT_SCALE = getenv("COLUMNWISE_WEIGHT_SCALE", 0) MXFP8 = getenv("MXFP8", 0) +MXFP4 = getenv("MXFP4", 0) FP8_DTYPE = dtypes.fp8e4m3 FP8_GRAD_DTYPE = dtypes.fp8e5m2 @@ -44,6 +45,11 @@ def matmul(x:Tensor, w:Tensor, fp8:bool=True, amax_x:Tensor|None=None, w_inv_sca from extra.gemm.cdna_asm_gemm import can_use_asm_gemm, asm_gemm if can_use_asm_gemm(x, w.T): return (asm_gemm(x, w.T),) return (x @ w.T,) + if MXFP4: + assert x is not None, "MXFP4 matmul requires an unquantized input" + from extra.gemm.cdna_asm_gemm import asm_gemm, can_use_asm_gemm + if can_use_asm_gemm(x, w.T): return (asm_gemm(x, w.T, mxfp4=True),) + return (x @ w.T,) assert w_inv_scale is not None, "fp8 matmul requires w_inv_scale (weights must be stored in fp8 with per-tensor scale)" if MXFP8: from extra.gemm.cdna_asm_gemm import asm_gemm, quantize_mxfp8, mx_pack, can_use_asm_gemm, _mx_block_scale @@ -79,7 +85,7 @@ def matmul(x:Tensor, w:Tensor, fp8:bool=True, amax_x:Tensor|None=None, w_inv_sca def norm_quantize_matmul(x:Tensor, norm:Tensor, w:Tensor, w_inv_scale:Tensor, eps:float, amax_x:Tensor, next_amax_x:Tensor, grad_amax_state:Tensor, next_grad_amax_state:Tensor): - if FUSED_ADD_NORM_MUL_QUANTIZE: + if FUSED_ADD_NORM_MUL_QUANTIZE and not MXFP4: from extra.llama_kernels.fused_rmsnorm_mul_quantize_fp8 import fused_rmsnorm_mul_quantize_fp8 x_fp8, x_normed, rrms = fused_rmsnorm_mul_quantize_fp8(x, norm, amax_x, eps, FP8_DTYPE, next_amax_x) out, *ret = matmul(None, w, w_inv_scale=w_inv_scale, x_fp8=x_fp8, amax_x=amax_x, @@ -92,7 +98,7 @@ def norm_quantize_matmul(x:Tensor, norm:Tensor, w:Tensor, w_inv_scale:Tensor, ep def add_norm_quantize_matmul(x:Tensor, residual:Tensor, norm:Tensor, w:Tensor, w_inv_scale:Tensor, eps:float, amax_x:Tensor, next_amax_x:Tensor, grad_amax_state:Tensor|None=None, next_grad_amax_state:Tensor|None=None): - if FUSED_ADD_NORM_MUL_QUANTIZE: + if FUSED_ADD_NORM_MUL_QUANTIZE and not MXFP4: from extra.llama_kernels.fused_rmsnorm_mul_quantize_fp8 import fused_add_rmsnorm_mul_quantize_fp8 x_fp8, h, x_normed, rrms = fused_add_rmsnorm_mul_quantize_fp8(x, residual, norm, amax_x, eps, FP8_DTYPE, next_amax_x) out, *ret = matmul(None, w, w_inv_scale=w_inv_scale, x_fp8=x_fp8, amax_x=amax_x, @@ -108,7 +114,7 @@ def silu_w13_quantize_matmul(x_w13:Tensor, w2:Tensor, s_2:Tensor, amax_x2:Tensor, next_amax_x2:Tensor, grad_amax_xw13:Tensor, next_grad_amax_xw13:Tensor, grad_amax_xout:Tensor, next_grad_amax_xout:Tensor): - if FUSED_SILU_W13: + if FUSED_SILU_W13 and not MXFP4: from extra.llama_kernels.cast_amax import fused_quantize_fp8_w13 x2_fp8 = fused_quantize_fp8_w13(x_w13, amax_x2, FP8_DTYPE, grad_amax_state=grad_amax_xw13, next_grad_amax_state=next_grad_amax_xw13, amax_out=next_amax_x2) @@ -179,6 +185,9 @@ class FlatTransformer: from extra.gemm.cdna_asm_gemm import quantize_mxfp8 w_q, w_e8, _ = quantize_mxfp8(w.reshape(self.n_layers * out_features, in_features)) return w_q.reshape(self.n_layers, out_features, in_features), w_e8.reshape(self.n_layers, out_features, in_features // 32) + if MXFP4: + # FP4 is produced dynamically so optimizer updates always start from the current BF16 weight. + return w.cast(dtypes.bfloat16), Tensor.ones(self.n_layers) amax = (w.abs().max(axis=2) if COLUMNWISE_WEIGHT_SCALE else w.abs().flatten(1).max(1)).detach() scale = FP8_MAX / (amax + 1e-8) inv_scale = (amax + 1e-8) / FP8_MAX diff --git a/extra/gemm/cdna_asm_gemm.py b/extra/gemm/cdna_asm_gemm.py index feca12a051..2270d7ea4b 100644 --- a/extra/gemm/cdna_asm_gemm.py +++ b/extra/gemm/cdna_asm_gemm.py @@ -1,4 +1,4 @@ -import atexit, functools, pathlib +import atexit, functools, math, pathlib from tinygrad import Tensor, Device, dtypes from tinygrad.dtype import AddrSpace from tinygrad.uop.ops import UOp, Ops, KernelInfo, AxisType @@ -111,17 +111,17 @@ def custom_hk_mxfp8_gemm(C:UOp, A:UOp, B:UOp, scale_A:UOp, scale_B:UOp, *extra:U # ** MXFP4 GEMM custom kernel @functools.cache -def custom_mxfp4_gemm(C:UOp, A:UOp, B:UOp, scale_a:UOp, scale_b:UOp, tile_m:int, tile_n:int) -> UOp: +def custom_mxfp4_gemm(C:UOp, A:UOp, B:UOp, scale_a:UOp, scale_b:UOp, *extra:UOp, tile_m:int, tile_n:int) -> UOp: from extra.gemm.gemm_mxfp4 import build_kernel - M, half_k = A.shape[0]*A.shape[1], A.shape[2] - N, half_k_b = B.shape + M, half_k = math.prod(A.shape[:-1]), A.shape[-1] + N, half_k_b = math.prod(B.shape[:-1]), B.shape[-1] K = half_k * 2 - assert half_k == half_k_b and C.shape == (*A.shape[:-1], N) + assert half_k == half_k_b and math.prod(C.shape[:-1]) == M and C.shape[-1] == N threads = UOp.special(256, "lidx0") groups_x, groups_y = UOp.special(ceildiv(N, tile_n), "gidx0"), UOp.special(ceildiv(M, tile_m), "gidx1") lds = UOp.placeholder((163840,), dtypes.uint8, 0, AddrSpace.LOCAL) - sink = UOp.sink(C.base, A.base, B.base, scale_a.base, scale_b.base, lds, threads, groups_x, groups_y, - arg=KernelInfo(f"custom_mxfp4_{M}_{N}_{K}", estimates=Estimates(ops=2*M*N*K))) + sink = UOp.sink(C.base, A.base, B.base, scale_a.base, scale_b.base, *(x.base for x in extra), lds, threads, groups_x, groups_y, + arg=KernelInfo(f"custom_mxfp4_gemm_{M}_{N}_{K}", estimates=Estimates(ops=2*M*N*K))) insts = build_kernel(M, N, K, tile_m, tile_n) return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=tuple(UOp(Ops.INS, arg=x) for x in insts)))) @@ -139,35 +139,47 @@ def quantize_mxfp8(x:Tensor) -> tuple[Tensor, Tensor, Tensor]: def _mxfp4_shuffle_weight(x:Tensor) -> Tensor: # shuffle_weight(x, layout=(16, 16)) on the packed uint8 buffer. + if x.ndim == 3: + ndev, rows, half_k = x.shape + return x.reshape(ndev, rows//16, 16, half_k//32, 2, 16).permute(0, 1, 3, 4, 2, 5).reshape(ndev, rows, half_k).contiguous() rows, half_k = x.shape return x.reshape(rows//16, 16, half_k//32, 2, 16).permute(0, 2, 3, 1, 4).reshape(rows, half_k).contiguous() def _mxfp4_shuffle_scales(x:Tensor) -> Tensor: # e8m0_shuffle: each 256x8 scale tile is arranged for the raw MFMA scale loads. + if x.ndim == 3: + ndev, rows, scale_k = x.shape + return x.reshape(ndev, rows//32, 2, 16, scale_k//8, 2, 4).permute(0, 1, 4, 6, 3, 5, 2).reshape(ndev, rows, scale_k).contiguous() rows, scale_k = x.shape return x.reshape(rows//32, 2, 16, scale_k//8, 2, 4).permute(0, 3, 5, 2, 4, 1).reshape(rows, scale_k).contiguous() def quantize_mxfp4(x:Tensor) -> tuple[Tensor, Tensor, Tensor]: # OCP MXFP4: 1x32 blocks, E2M1 values packed low-nibble first, and E8M0 scales. - assert x.ndim == 2 and x.shape[1] % 256 == 0 and x.shape[0] % 32 == 0, \ + *batch, K = x.shape + rows = math.prod(batch) + assert x.ndim >= 2 and K % 256 == 0 and rows % 32 == 0, \ f"mxfp4 quantization needs rows%32 and K%256, got {x.shape}" - rows, K = x.shape - xb = x.float().reshape(rows, K//32, 32) + xb = x.float().reshape(*batch, K//32, 32) amax = xb.abs().max(axis=-1) # even scale rounding: round the fp32 significand before choosing 2^(floor(log2)-2). amax_rounded = ((amax.bitcast(dtypes.uint32) + 0x200000) & 0xFF800000).bitcast(dtypes.float32) scale_exp = (amax_rounded.maximum(2**-126).log2().floor() - 2).clamp(-127, 127) e8 = (scale_exp + 127).cast(dtypes.uint8) - scaled = xb * (-scale_exp).exp2().reshape(rows, K//32, 1) + scaled = xb * (-scale_exp).exp2().reshape(*batch, K//32, 1) mag = scaled.abs() code = sum(x.cast(dtypes.uint8) for x in (mag > .25, mag >= .75, mag > 1.25, mag >= 1.75, mag > 2.5, mag >= 3.5, mag > 5.0)) code = code | ((scaled < 0).cast(dtypes.uint8) << 3) - code = code.reshape(rows, K) - packed = code[:, 0::2] | (code[:, 1::2] << 4) - return packed, e8, _mxfp4_shuffle_scales(e8) + code = code.reshape(*batch, K) + packed = code[..., 0::2] | (code[..., 1::2] << 4) + if isinstance(x.device, tuple) and x.uop.axis == x.ndim-2 and x.shape[x.uop.axis] == len(x.device): + axis = x.uop.axis + order = (axis, *range(axis), *range(axis+1, e8.ndim)) + e8_local = e8.permute(order) + return packed, e8, _mxfp4_shuffle_scales(e8_local.reshape(e8_local.shape[0], -1, K//32)) + return packed, e8, _mxfp4_shuffle_scales(e8.reshape(rows, K//32)) def mx_pack(e8:Tensor) -> Tensor: rows, scale_K = e8.shape @@ -390,6 +402,19 @@ def custom_mx_gemm_bw(gradient:UOp, kernel:UOp, has_w_post:bool, w_stored:bool=F if wp is not None: grad_b = grad_b / wp.reshape(-1, 1) return (None, grad_a.uop, grad_b.uop) + tuple(None for _ in inputs[3:]) +# ** mxfp4 gemm backward + +def custom_mxfp4_gemm_bw(gradient:UOp, kernel:UOp): + # The raw kernel consumes quantized buffers, while the final two inputs retain the BF16 operands for STE gradients. + inputs = kernel.src[1:] # (out, a_q, b_q, scale_a, scale_b, a, w) + assert len(inputs) == 7 + a, w = Tensor(inputs[5], device=inputs[5].device), Tensor(inputs[6], device=inputs[6].device) + g = Tensor(gradient, device=a.device)[:a.shape[0]].cast(dtypes.bfloat16) + grad_a = asm_gemm(g, w, mxfp4=True) + a_flat, g_flat = a.reshape(-1, a.shape[-1]), g.reshape(-1, g.shape[-1]) + grad_w = asm_gemm(g_flat.T, a_flat, mxfp4=True) + return (None, None, None, None, None, grad_a.uop, grad_w.uop) + # ** main gemm function def asm_gemm(a:Tensor, b:Tensor, x_scale:Tensor|None=None, w_scale:Tensor|None=None, grad_amax_state:Tensor|None=None, @@ -433,10 +458,17 @@ def asm_gemm(a:Tensor, b:Tensor, x_scale:Tensor|None=None, w_scale:Tensor|None=N if mxfp4: tile_m, tile_n = next((tm, tn) for tm, tn in ((256, 256), (192, 256), (128, 512)) if (batch*M) % tm == N % tn == 0) fxn = functools.partial(custom_mxfp4_gemm, tile_m=tile_m, tile_n=tile_n) - a_q, _, scale_a = quantize_mxfp4(a.reshape(batch*M, K)) - b_q, _, scale_b = quantize_mxfp4(b.T) - a_q, b_q = a_q.reshape(batch, M, K//2).contiguous(), _mxfp4_shuffle_weight(b_q) - out = Tensor.custom_kernel(out, a_q, b_q, scale_a, scale_b, fxn=fxn)[0] + w = b.T + if k_sharded: + ndev = len(a.device) + a_q, _, scale_a = quantize_mxfp4(a.reshape(batch, M, ndev, K)) + b_q, _, scale_b = quantize_mxfp4(w.reshape(w.shape[0], ndev, K)) + b_q = _mxfp4_shuffle_weight(b_q.permute(1, 0, 2)) + else: + a_q, _, scale_a = quantize_mxfp4(a.reshape(batch*M, K)) + b_q, _, scale_b = quantize_mxfp4(w) + a_q, b_q = a_q.reshape(batch, M, K//2).contiguous(), _mxfp4_shuffle_weight(b_q) + out = Tensor.custom_kernel(out, a_q, b_q, scale_a, scale_b, a, w, fxn=fxn, grad_fxn=custom_mxfp4_gemm_bw)[0] elif mx: # mxfp8 1x32 block scaling if mx_scales is not None: From 93b74c75fca0664fa0df8891c05ea5b9ff0865c5 Mon Sep 17 00:00:00 2001 From: wozeparrot Date: Fri, 31 Jul 2026 18:25:32 +0800 Subject: [PATCH 16/44] gptoss: grouped moe (#17322) --- examples/mlperf/model_train.py | 5 +- examples/mlperf/models/gpt_oss.py | 80 ++++++++++++++++++++++++------- 2 files changed, 68 insertions(+), 17 deletions(-) diff --git a/examples/mlperf/model_train.py b/examples/mlperf/model_train.py index 6b70032a13..c1345c3f00 100644 --- a/examples/mlperf/model_train.py +++ b/examples/mlperf/model_train.py @@ -1748,7 +1748,10 @@ def train_gptoss(): from extra.gemm.cdna_asm_gemm import _mx_block_scale model_state = get_state_dict(model) - fp8_scale_names = {n: f"{n}_scale" for n, t in model_state.items() if t.dtype == FP8_DTYPE} + def _scale_key(n): + if "." in n and (c:=f"{(b:=n.rsplit('.',1))[0]}_scale.{b[1]}") in model_state: return c + return f"{n}_scale" + fp8_scale_names = {n: _scale_key(n) for n, t in model_state.items() if t.dtype == FP8_DTYPE} fp8_inv_scales = [model_state[sname] for sname in fp8_scale_names.values()] for wname, sname in fp8_scale_names.items(): w, scale = model_state[wname], model_state[sname] diff --git a/examples/mlperf/models/gpt_oss.py b/examples/mlperf/models/gpt_oss.py index 8e34e7b65f..4bb8613b94 100644 --- a/examples/mlperf/models/gpt_oss.py +++ b/examples/mlperf/models/gpt_oss.py @@ -13,10 +13,14 @@ from tinygrad.uop.ops import Ops, UOp from extra.models.llama import apply_rotary_emb from extra.llama_kernels.rmsnorm import rmsnorm from extra.gemm.cdna_asm_gemm import _mx_block_scale, _mx_block_scale_3d, quantize_mxfp8 +from extra.gemm.moe_gemm import grouped_mx_gemm +from extra.gemm.moe_routing import route, dispatch, combine FP8_DTYPE = dtypes.fp8e4m3 FP8_MAX = 448.0 -INIT_STD = 0.008 +INIT_STD = 0.02 +ASM_GEMM = getenv("ASM_GEMM", 1) + def _quant_dequant_fwd(x:Tensor) -> Tensor: # x (2d bf16) -> bf16 value after an mxfp8 round-trip (1x32 block scaling on the last axis) @@ -59,10 +63,34 @@ def dequant_weight(w_q:Tensor, w_scale:Tensor) -> Tensor: def matmul_mx(x:Tensor, w_q:Tensor, w_scale:Tensor) -> Tensor: l_shape = x.shape[:-1] + if ASM_GEMM: + from extra.gemm.cdna_asm_gemm import asm_gemm, can_use_asm_gemm, mx_pack + x2, K, N = x.reshape(-1, x.shape[-1]), x.shape[-1], w_q.shape[0] + wq, ws = w_q, w_scale + if (pad := (-K) % 256): + x2 = x2.pad(((0, 0), (0, pad))) + wq = wq.pad(((0, 0), (0, pad))) + ws = ws.pad(((0, 0), (0, pad // 32)), value=127).cast(dtypes.uint8) + if (npad := (-N) % 256): + wq = wq.pad(((0, npad), (0, 0))) + ws = ws.pad(((0, npad), (0, 0)), value=127).cast(dtypes.uint8) + x_q, x_e8, x_si = quantize_mxfp8(x2) + if x_si is not None and can_use_asm_gemm(x_q, wq.T): + out = asm_gemm(x_q, wq.T, mx=True, mx_scales=(x_si, x_e8, mx_pack(ws), ws), mx_w_stored=True) + return (out[:, :N] if npad else out).reshape(*l_shape, N).cast(dtypes.bfloat16) x_phys = quant_dequant_mx(x.reshape(-1, x.shape[-1])).reshape(*l_shape, x.shape[-1]) w_phys = dequant_weight(w_q, w_scale) return (x_phys @ w_phys.T).cast(dtypes.bfloat16) +def _pad_to_mult(t:Tensor, axis:int, mult:int=256) -> Tensor: + if (r := (-t.shape[axis]) % mult) == 0: return t + pads = [(0, 0)] * t.ndim + pads[axis] = (0, r) + return t.pad(tuple(pads)) + +def _pad_cols(t:Tensor) -> Tensor: return _pad_to_mult(t, -1) +def _pad_rows(t:Tensor) -> Tensor: return _pad_to_mult(t, -2) + def swiglu(x:Tensor, limit:float=7.0, alpha:float=1.702) -> Tensor: x_glu, x_linear = x[..., ::2], x[..., 1::2] x_glu = x_glu.clamp(max_=limit) @@ -99,9 +127,9 @@ class GPTOSS: self.ffn_norm = Tensor.ones(n_layers, dim).contiguous() self.gate = Tensor.normal(n_layers, n_experts, dim, mean=0.0, std=INIT_STD, dtype=dtypes.bfloat16) self.gate_bias = Tensor.zeros(n_layers, n_experts, dtype=dtypes.bfloat16).contiguous() - self.w_gate_up, self.w_gate_up_scale = self._quant_weight(n_layers, n_experts, intermediate_size * 2, dim) + self.w_gate_up, self.w_gate_up_scale = self._quant_weight(n_layers, n_experts, intermediate_size * 2, dim, moe=True) self.w_gate_up_bias = Tensor.zeros(n_layers, n_experts, intermediate_size * 2, dtype=dtypes.bfloat16).contiguous() - self.w_down, self.w_down_scale = self._quant_weight(n_layers, n_experts, dim, intermediate_size, std=scaled_std) + self.w_down, self.w_down_scale = self._quant_weight(n_layers, n_experts, dim, intermediate_size, std=scaled_std, moe=True) self.w_down_bias = Tensor.zeros(n_layers, n_experts, dim, dtype=dtypes.bfloat16).contiguous() # output @@ -111,10 +139,15 @@ class GPTOSS: self.output = Tensor.normal(vocab_size, dim, mean=0.0, std=INIT_STD, dtype=dtypes.bfloat16) self.freqs_cis = precompute_freqs_cis(head_dim, max_context * 2, rope_theta).contiguous().is_param_(False) - def _quant_weight(self, *shape:int, std:float=INIT_STD): - w = Tensor.zeros(*shape) if getenv("ZEROS") else Tensor.normal(*shape, mean=0.0, std=std) - w_q, w_e8, _ = quantize_mxfp8(w) - return w_q, w_e8.is_param_(False) + def _quant_weight(self, *shape:int, std:float=INIT_STD, moe:bool=False): + def _one(*s:int): + w = Tensor.zeros(*s) if getenv("ZEROS") else Tensor.normal(*s, mean=0.0, std=std) + w_q, w_e8, _ = quantize_mxfp8(_pad_cols(_pad_rows(w)) if moe else w) + return w_q, w_e8.is_param_(False) + if moe: + qs = [_one(*shape[1:]) for _ in range(shape[0])] + return [q[0] for q in qs], [q[1] for q in qs] + return _one(*shape) def _attn_mask(self, seqlen:int, dtype) -> Tensor: i, j = Tensor.arange(seqlen).reshape(seqlen, 1), Tensor.arange(seqlen).reshape(1, seqlen) @@ -173,17 +206,32 @@ class GPTOSS: w_down:Tensor, w_down_scale:Tensor, w_down_bias:Tensor): x_normed, rrms = rmsnorm(x, self.norm_eps) inp = x_normed * ffn_norm - logits = inp.float() @ gate.float().T + gate_bias.float() - thresh = logits.topk(self.experts_per_tok)[0][..., -1:] - weights = (logits >= thresh).where(logits, -float("inf")).softmax(-1) + dim, inter = self.dim, self.intermediate_size - out = None - for e in range(self.n_experts): - gate_up = matmul_mx(inp, w_gate_up[e], w_gate_up_scale[e]) + w_gate_up_bias[e] - y = (matmul_mx(swiglu(gate_up, self.swiglu_limit), w_down[e], w_down_scale[e]) + w_down_bias[e]).contiguous() - contrib = weights[..., e:e+1].cast(y.dtype) * y - out = contrib if out is None else out + contrib + if getenv("GROUPED_MOE", 0): + bsz, seqlen = x.shape[:2] + inp, logits = inp.reshape(-1, dim), logits.reshape(-1, self.n_experts) + r = route(logits, self.experts_per_tok, self.n_experts) + onehot = r.rows_e.one_hot(self.n_experts).float() + xg = dispatch(_pad_cols(inp.cast(dtypes.bfloat16)), r) + h = grouped_mx_gemm(xg, (w_gate_up, w_gate_up_scale), r.off)[:, :2*inter] + (onehot @ w_gate_up_bias.float()).cast(dtypes.bfloat16) + y = swiglu(h, self.swiglu_limit) + z = grouped_mx_gemm(_pad_cols(y.cast(dtypes.bfloat16)), (w_down, w_down_scale), r.off)[:, :dim] \ + + (onehot @ w_down_bias.float()).cast(dtypes.bfloat16) + out = combine(z, r, inp.shape[0], self.experts_per_tok).reshape(bsz, seqlen, dim) + else: + thresh = logits.topk(self.experts_per_tok)[0][..., -1:] + weights = (logits >= thresh).where(logits, -float("inf")).softmax(-1) + + out = None + for e in range(self.n_experts): + gu_q, gu_s = w_gate_up[e][:2*inter, :dim].contiguous(), w_gate_up_scale[e][:2*inter, :dim//32].contiguous() + dn_q, dn_s = w_down[e][:dim, :inter].contiguous(), w_down_scale[e][:dim, :inter//32].contiguous() + gate_up = matmul_mx(inp, gu_q, gu_s) + w_gate_up_bias[e] + y = (matmul_mx(swiglu(gate_up, self.swiglu_limit), dn_q, dn_s) + w_down_bias[e]).contiguous() + contrib = weights[..., e:e+1].cast(y.dtype) * y + out = contrib if out is None else out + contrib return out, [x_normed, rrms] @function(precompile=True, precompile_backward=True) From 6d2700f0b723999770e445e1ff00d7a68f46d4e3 Mon Sep 17 00:00:00 2001 From: qazal <77887910+Qazalin@users.noreply.github.com> Date: Fri, 31 Jul 2026 19:24:44 +0800 Subject: [PATCH 17/44] failing test for unbound _device_num err in BEAM (#17326) * min failing test * switch to cpu * err --- test/backend/test_multitensor.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/test/backend/test_multitensor.py b/test/backend/test_multitensor.py index ce218ac7a9..5ee6a0cc3b 100644 --- a/test/backend/test_multitensor.py +++ b/test/backend/test_multitensor.py @@ -76,6 +76,14 @@ class TestMultiTensor(unittest.TestCase): run_linear(linear) self.assertEqual(len(set(names)), 1, "function was relinearized") + @unittest.expectedFailure + def test_shard_beam(self): + cpu_2 = ("CPU:1", "CPU:2") + src = Tensor.ones(16).shard(cpu_2, 0).realize() + pad = src.to(cpu_2[::-1]).schedule_linear().src[0] + with Context(BEAM=1, IGNORE_BEAM_CACHE=1): prg = compile_linear(UOp(Ops.LINEAR, src=(pad,))).src[0].src[0] + self.assertNotEqual(prg.src[0].arg.applied_opts, ()) + def test_shard_same_device(self): X = Tensor.ones(256).contiguous().realize() X.shard_((d1, X.device), 0) From 0c4bfaeb486012fffe248405c17ddbc277f6ad32 Mon Sep 17 00:00:00 2001 From: nimlgen <138685161+nimlgen@users.noreply.github.com> Date: Fri, 31 Jul 2026 15:49:56 +0300 Subject: [PATCH 18/44] coalesce ints (#17323) * merge ints * fix z3 validation of coalesced loads * fix uint vector names in CUDA and Metal --- test/null/test_validate_oob.py | 8 ++++++++ tinygrad/codegen/late/coalesce.py | 2 +- tinygrad/renderer/cstyle.py | 6 +++--- tinygrad/uop/validate.py | 4 ++-- 4 files changed, 14 insertions(+), 6 deletions(-) diff --git a/test/null/test_validate_oob.py b/test/null/test_validate_oob.py index 4388f2b863..39e509770d 100644 --- a/test/null/test_validate_oob.py +++ b/test/null/test_validate_oob.py @@ -138,6 +138,14 @@ class TestValidateOOB(unittest.TestCase): with self.assertRaises(RuntimeError): to_uops_list([buf1.index((ld0 * 2).valid((ld0 >= 0) & (ld0 < 64))).load(dtype=dtypes.int)]) # oob + def test_load_from_shrink_as_index(self): + with Context(CHECK_OOB=1, SPEC=2): + buf0 = UOp.param(0, dtypes.int, (16,)) + buf1 = UOp.param(1, dtypes.int, (64,)) + shrink = UOp(Ops.SHRINK, src=(buf0, UOp.const(dtypes.int, 0), UOp.const(dtypes.weakint, 4))) + ld0 = shrink.load(dtype=dtypes.int).index(0) + to_uops_list([buf1.index(ld0.valid((ld0 >= 0) & (ld0 < 64))).load(dtype=dtypes.int)]) + def test_load_bool_as_mask(self): with Context(CHECK_OOB=1, SPEC=2): buf_bool = UOp.param(0, dtypes.bool, (16,)) diff --git a/tinygrad/codegen/late/coalesce.py b/tinygrad/codegen/late/coalesce.py index 8510f639bd..077b4a734b 100644 --- a/tinygrad/codegen/late/coalesce.py +++ b/tinygrad/codegen/late/coalesce.py @@ -127,7 +127,7 @@ def memory_coalescing(sink:UOp, ctx:Renderer) -> UOp: if ctx is not None and ctx.target.device == "DSP": lengths = [128,64,32,16,8,4] must_divide = False - elif buf.dtype not in (dtypes.float, dtypes.half, *dtypes.fp8s) and not is_image_shape(buf._shape): + elif buf.dtype not in (dtypes.float, dtypes.half, dtypes.int, dtypes.uint, *dtypes.fp8s) and not is_image_shape(buf._shape): pass elif buf.addrspace == AddrSpace.REG: pass diff --git a/tinygrad/renderer/cstyle.py b/tinygrad/renderer/cstyle.py index ce6d6330de..1b0287b7de 100644 --- a/tinygrad/renderer/cstyle.py +++ b/tinygrad/renderer/cstyle.py @@ -356,7 +356,7 @@ class MetalRenderer(CStyleLanguage): code_for_workitem = {"g": lambda x: f"gid.{chr(120+int(x))}", "l": lambda x: f"lid.{chr(120+int(x))}"} # uint3 used for gid/lid - TODO: this should probably be `ushort3 lid [[thread_position_in_threadgroup]]` extra_args = ['uint3 gid [[threadgroup_position_in_grid]]', 'uint3 lid [[thread_position_in_threadgroup]]'] - type_map = {dtypes.bfloat16: "bfloat"} + type_map = {dtypes.uint32: "uint", dtypes.bfloat16: "bfloat"} # precise::sin code_for_op = {**CStyleLanguage.code_for_op, Ops.SIN: lambda x,dtype: f"precise::sin({x})"} @@ -420,7 +420,7 @@ class CUDARenderer(CStyleLanguage): Ops.EXP2: lambda x,dtype: f"hexp2({x})" if dtype in (dtypes.half, dtypes.bfloat16) else f"exp2({x})", Ops.SQRT: lambda x,dtype: f"hsqrt({x})" if dtype in (dtypes.half, dtypes.bfloat16) else f"sqrt({x})", Ops.RECIPROCAL: lambda x,dtype: f"hrcp({x})" if dtype in (dtypes.half, dtypes.bfloat16) else f"(1/{x})" } - type_map = {dtypes.bfloat16: "nv_bfloat16", dtypes.fp8e4m3: "__nv_fp8_e4m3", dtypes.fp8e5m2: "__nv_fp8_e5m2"} + type_map = {dtypes.uint32: "uint", dtypes.bfloat16: "nv_bfloat16", dtypes.fp8e4m3: "__nv_fp8_e4m3", dtypes.fp8e5m2: "__nv_fp8_e5m2"} extra_matcher = create_non_native_float_pats(dtypes.fp8s, casting=False) + PatternMatcher([ (UPat(Ops.CAST, dtypes.fp8s, UPat.var("x", dtypes.fp8s), name='y'), lambda x,y: x.cast(dtypes.float).cast(y.dtype) if x.dtype!=y.dtype else None), ]) @@ -436,7 +436,7 @@ class CUDARenderer(CStyleLanguage): def render_kernel(self, function_name, kernel, bufs, uops, prefix=None): # TODO: why is dtypes.bfloat16.name == "__bf16"? would be easier not override dtypes.name - prefix = ["#define INFINITY (__int_as_float(0x7f800000))", "#define NAN (__int_as_float(0x7fffffff))", + prefix = ["typedef unsigned int uint;", "#define INFINITY (__int_as_float(0x7f800000))", "#define NAN (__int_as_float(0x7fffffff))", "template __device__ __forceinline__ T tg_bitcast(F v) { union U { F f; T t; }; U u; u.f = v; return u.t; }"] used_dtypes = uops_to_dtypes(uops) if any(dt in dtypes.fp8s for dt, _ in used_dtypes): prefix.append("#include ") diff --git a/tinygrad/uop/validate.py b/tinygrad/uop/validate.py index bcd21c78f1..7c50b76ea4 100644 --- a/tinygrad/uop/validate.py +++ b/tinygrad/uop/validate.py @@ -59,8 +59,8 @@ z3_renderer = PatternMatcher([ ]) def uops_to_z3(solver:z3.Solver, *uops: UOp) -> list[z3.ExprRef]: - # gate on upstream AFTER/BUFFER, but keep INDEX as an unknown LOAD - lst = list(UOp.sink(*uops).toposort(gate=lambda x: x.op not in {Ops.AFTER, Ops.BUFFER} and \ + # gate on upstream memory addressing, but keep INDEX as an unknown LOAD + lst = list(UOp.sink(*uops).toposort(gate=lambda x: x.op not in {Ops.AFTER, Ops.BUFFER, Ops.SHRINK} and \ (x.dtype in dtypes.ints+(dtypes.bool, dtypes.weakint) or x.op is Ops.SINK)))[:-1] z3map: dict[UOp, z3.ExprRef] = {} for u in lst: From 4f5cadd15d7aa2639b20ed3099bb5b838efe8f8c Mon Sep 17 00:00:00 2001 From: wozeparrot Date: Fri, 31 Jul 2026 20:51:58 +0800 Subject: [PATCH 19/44] gptoss ci (#17325) --- .github/workflows/test.yml | 2 ++ examples/mlperf/model_train.py | 3 ++- examples/mlperf/models/gpt_oss.py | 2 +- 3 files changed, 5 insertions(+), 2 deletions(-) diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index a7f37e8250..e4b54e6c21 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -353,6 +353,8 @@ jobs: run: DEV=NULL NULL_ALLOW_COPYOUT=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=24 GPUS=4 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py - name: Test llama 3 training run: DEV=NULL NULL_ALLOW_COPYOUT=1 SAMPLES=300 BS=8 SEQLEN=512 GRADIENT_ACC_STEPS=1 FAKEDATA=1 DEFAULT_FLOAT=bfloat16 OPTIM_DTYPE=bfloat16 LLAMA3_SIZE=1B MODEL=llama3 python3 examples/mlperf/model_train.py + - name: Test gpt-oss training + run: DEV=NULL NULL_ALLOW_COPYOUT=1 SAMPLES=32 BS=2 SEQLEN=128 GRADIENT_ACC_STEPS=1 FAKEDATA=1 DEFAULT_FLOAT=bfloat16 OPTIM_DTYPE=bfloat16 MXFP8=1 VOCAB_SIZE=32000 LAYERS=2 EXPERTS=4 MODEL=gptoss PYTHONPATH=. python3 examples/mlperf/model_train.py - name: Run process replay tests uses: ./.github/actions/process-replay diff --git a/examples/mlperf/model_train.py b/examples/mlperf/model_train.py index c1345c3f00..026e35977c 100644 --- a/examples/mlperf/model_train.py +++ b/examples/mlperf/model_train.py @@ -1711,9 +1711,10 @@ def train_gptoss(): wandb.init(config=config, **wandb_args, project="MLPerf-gpt-oss") model_params = GPT_OSS_20B - model_params['vocab_size'] = 128256 + model_params['vocab_size'] = getenv("VOCAB_SIZE", 128256) real_vocab_size = model_params['vocab_size'] if (layers:=getenv("LAYERS")) != 0: model_params['n_layers'] = layers + if (experts:=getenv("EXPERTS")) != 0: model_params['n_experts'] = experts print(f"model parameters: {model_params}") model = GPTOSS(**model_params, max_context=SEQLEN) diff --git a/examples/mlperf/models/gpt_oss.py b/examples/mlperf/models/gpt_oss.py index 4bb8613b94..1e157a94da 100644 --- a/examples/mlperf/models/gpt_oss.py +++ b/examples/mlperf/models/gpt_oss.py @@ -19,7 +19,7 @@ from extra.gemm.moe_routing import route, dispatch, combine FP8_DTYPE = dtypes.fp8e4m3 FP8_MAX = 448.0 INIT_STD = 0.02 -ASM_GEMM = getenv("ASM_GEMM", 1) +ASM_GEMM = getenv("ASM_GEMM", 0) def _quant_dequant_fwd(x:Tensor) -> Tensor: From 155b84ee80fd8b451f062c085bbba00b268acbdb Mon Sep 17 00:00:00 2001 From: nimlgen <138685161+nimlgen@users.noreply.github.com> Date: Fri, 31 Jul 2026 16:08:18 +0300 Subject: [PATCH 20/44] hcq2 faster schedule (#17324) * avoid quadratic STACK dtype promotion * build HCQ patch stacks directly * pack HCQ command buffers linearly * remove HCQ command buffer simplification --- tinygrad/runtime/support/hcq2.py | 13 +++++++------ tinygrad/uop/ops.py | 3 ++- 2 files changed, 9 insertions(+), 7 deletions(-) diff --git a/tinygrad/runtime/support/hcq2.py b/tinygrad/runtime/support/hcq2.py index cf5cad2718..ce2f68ffc4 100644 --- a/tinygrad/runtime/support/hcq2.py +++ b/tinygrad/runtime/support/hcq2.py @@ -42,8 +42,9 @@ def unwrap_mstack(u): return unwrap_mstack(u.src[0]) if u.op in {Ops.MSELECT, Ops.SLICE} else (u,) def make_patches(buf:UOp, patches:Sequence[tuple[sint, UOp]]) -> UOp: - return buf.index(UOp.stack(*(UOp.const(dtypes.int, off // buf.dtype.itemsize) for off,_ in patches))) \ - .store(UOp.stack(*(val.simplify().cast(buf.dtype) for _,val in patches))) + offsets = UOp(Ops.STACK, dtypes.int, tuple(UOp.const(dtypes.int, off // buf.dtype.itemsize) for off,_ in patches)) + values = UOp(Ops.STACK, buf.dtype, tuple(val.cast(buf.dtype) for _,val in patches)) + return buf.index(offsets).store(values) def make_binary_patch(buf:UOp, blob:bytes) -> UOp: data = UOp(Ops.BINARY, src=(), arg=blob).bitcast(buf.dtype) @@ -51,13 +52,13 @@ def make_binary_patch(buf:UOp, blob:bytes) -> UOp: return buf.index(r).store(data.index(r).load()).end(r) def make_cmdbuf(lin, devs, buf:UOp|None=None, dep:UOp|None=None): - blob, patches = b'', [] + blob, patches = bytearray(), [] for s in (s for ins in lin.src for s in ins.src): - if (ssimp:=s.simplify()).op is not Ops.CONST: patches.append((len(blob), ssimp)) - blob += struct.pack(f'<{ssimp.dtype.fmt}', ssimp.arg if ssimp.op is Ops.CONST else 0x0) + if s.op is not Ops.CONST: patches.append((len(blob), s)) + blob.extend(struct.pack(f'<{s.dtype.fmt}', s.arg if s.op is Ops.CONST else 0x0)) cmdbuf = buf if buf is not None else UOp.placeholder((len(blob) // 4,), dtypes.uint32, next(UOp.unique_num), device=devs).rtag("cmdbuf") writable = cmdbuf.after(dep) if dep is not None else cmdbuf - return cmdbuf.after(make_binary_patch(writable, blob), *((make_patches(writable, patches),) if patches else ())) + return cmdbuf.after(make_binary_patch(writable, bytes(blob)), *((make_patches(writable, patches),) if patches else ())) def make_signal(devs, queue="COMPUTE:0", sentinel=False): return UOp.placeholder((1,), dtypes.uint64, 0, device=devs, volatile=True).rtag("sentinel_signal" if sentinel else f"{queue}_timeline_signal") diff --git a/tinygrad/uop/ops.py b/tinygrad/uop/ops.py index 0670a57b2d..ac996f5767 100644 --- a/tinygrad/uop/ops.py +++ b/tinygrad/uop/ops.py @@ -791,7 +791,8 @@ class UOp(RandMixin, metaclass=UOpMetaClass): case Ops.STACK: # arg is the other srcs; all are cast to the promoted dtype, spec requires STACK srcs to match its dtype srcs = (self,)+tuple(arg) - return UOp(Ops.STACK, src=tuple(u.cast(dtype_from_uop(Ops.STACK, srcs, None)) for u in srcs)) + dtype = cast(DType, dtype_from_uop(Ops.STACK, srcs, None)) + return UOp(Ops.STACK, dtype, tuple(u.cast(dtype) for u in srcs)) case _: raise RuntimeError(f"{op} is not a MovementOp") usrcs = [shape_to_shape_arg(arg) for arg in src_args] if len(usrcs) == 0: return UOp(op, src=(self,), arg=arg) From 8e2f175542aab8d92c29c430eafc042e63bc5f43 Mon Sep 17 00:00:00 2001 From: chenyu Date: Fri, 31 Jul 2026 09:46:37 -0400 Subject: [PATCH 21/44] const(dtype, b) -> const(b, dtype) [PR] (#17328) prep for dtype removal --- extra/gemm/cdna_asm_gemm.py | 6 +- extra/gemm/mi350x_uop_matmul.py | 2 +- extra/gemm/moe_routing.py | 2 +- extra/hcq2/ops_amd2.py | 28 +- .../quantize_fp8_delayed/__init__.py | 2 +- test/amd/test_emu2_pcode.py | 104 ++++---- test/backend/test_call.py | 4 +- test/backend/test_const_folding.py | 4 +- test/backend/test_custom_kernel.py | 2 +- test/backend/test_isel.py | 2 +- test/backend/test_linearizer_dumb.py | 20 +- test/backend/test_pickle.py | 2 +- test/backend/test_renderer_failures.py | 12 +- test/backend/test_tensor.py | 10 +- test/backend/test_uops.py | 18 +- test/external/external_benchmark_op_conv.py | 8 +- test/external/fuzz_symbolic.py | 6 +- test/helpers.py | 2 +- test/mockgpu/amd/emu.py | 244 +++++++++--------- test/mockgpu/amd/pcode.py | 8 +- test/null/test_const_folding.py | 12 +- test/null/test_gpudims.py | 4 +- test/null/test_gradient.py | 4 +- test/null/test_graph_rewrite.py | 86 +++--- test/null/test_helpers.py | 8 +- test/null/test_linearizer_failures.py | 12 +- test/null/test_microbenchmarks.py | 16 +- test/null/test_pattern_matcher.py | 60 ++--- test/null/test_simplify_valid_idx.py | 38 +-- test/null/test_tensor_uop_mixin.py | 4 +- test/null/test_transcendental_helpers.py | 44 ++-- test/null/test_uop_graph.py | 198 +++++++------- test/null/test_uop_repr.py | 8 +- test/null/test_uop_resolve.py | 24 +- test/null/test_uop_symbolic.py | 30 +-- test/null/test_uop_vmin_vmax.py | 34 +-- test/null/test_uops.py | 88 +++---- test/null/test_validate_oob.py | 20 +- test/null/test_viz.py | 24 +- test/unit/test_assign.py | 2 +- test/unit/test_dtype_weak.py | 58 ++--- test/unit/test_function.py | 2 +- test/unit/test_invalid_tensor.py | 4 +- test/unit/test_jit.py | 4 +- test/unit/test_metal_graph.py | 2 +- test/unit/test_multitensor.py | 4 +- test/unit/test_tensor_data.py | 4 +- tinygrad/callify.py | 2 +- tinygrad/codegen/__init__.py | 4 +- tinygrad/codegen/decomp/dtype.py | 10 +- tinygrad/codegen/decomp/op.py | 2 +- tinygrad/codegen/decomp/transcendental.py | 2 +- tinygrad/codegen/late/coalesce.py | 4 +- tinygrad/codegen/late/regalloc.py | 4 +- tinygrad/codegen/opt/postrange.py | 4 +- tinygrad/llm/gguf.py | 6 +- tinygrad/mixin/creation.py | 4 +- tinygrad/mixin/elementwise.py | 2 +- tinygrad/mixin/op.py | 8 +- tinygrad/renderer/cstyle.py | 2 +- tinygrad/renderer/isa/x86.py | 12 +- tinygrad/renderer/llvmir.py | 2 +- tinygrad/renderer/nir.py | 2 +- tinygrad/renderer/wgsl.py | 4 +- tinygrad/runtime/ops_cpu.py | 12 +- tinygrad/runtime/support/hcq2.py | 14 +- tinygrad/schedule/indexing.py | 12 +- tinygrad/schedule/memory.py | 2 +- tinygrad/schedule/rangeify.py | 2 +- tinygrad/tensor.py | 10 +- tinygrad/uop/ops.py | 32 +-- tinygrad/uop/render.py | 2 +- tinygrad/uop/spec.py | 2 +- tinygrad/uop/symbolic.py | 12 +- 74 files changed, 727 insertions(+), 727 deletions(-) diff --git a/extra/gemm/cdna_asm_gemm.py b/extra/gemm/cdna_asm_gemm.py index 2270d7ea4b..ef3d35fd50 100644 --- a/extra/gemm/cdna_asm_gemm.py +++ b/extra/gemm/cdna_asm_gemm.py @@ -236,10 +236,10 @@ def custom_uop_gemm(C:UOp, A:UOp, B:UOp) -> UOp: m = UOp.range(M, 1) n = UOp.range(N, 2) k = UOp.range(K, 0, AxisType.REDUCE) - mul = (A.flatten().index((m*UOp.const(dtypes.weakint, K)+k))* - B.flatten().index((k*UOp.const(dtypes.weakint, N)+n))).cast(dtypes.float32) + mul = (A.flatten().index((m*UOp.const(K)+k))* + B.flatten().index((k*UOp.const(N)+n))).cast(dtypes.float32) red = mul.reduce(k, arg=Ops.ADD, dtype=dtypes.float32).cast(C.dtype) - store = C.flatten().index((m*UOp.const(dtypes.weakint, N)+n)).store(red).end(m, n) + store = C.flatten().index((m*UOp.const(N)+n)).store(red).end(m, n) return store.sink(arg=KernelInfo(name=f'uop_gemm_{M}_{N}_{K}')) # ** bf16 A @ B.T kernel in C diff --git a/extra/gemm/mi350x_uop_matmul.py b/extra/gemm/mi350x_uop_matmul.py index cc3c039e10..bd9e717cde 100644 --- a/extra/gemm/mi350x_uop_matmul.py +++ b/extra/gemm/mi350x_uop_matmul.py @@ -79,7 +79,7 @@ def custom_gemm(C:UOp, A:UOp, B:UOp) -> UOp: # this is the big accumulator acc = UOp.placeholder((BLOCK_N//TC_N, BLOCK_M//TC_M//WARPGROUP_SIZE), dtypes.float, 0, AddrSpace.REG) assert acc.size*WARP_SIZE*WARPGROUP_SIZE*4 == BLOCK_M*BLOCK_N - acc = acc[init_l:=UOp.range(acc.size, 500)].set(UOp.const(dtypes.float, (0.0,)*4), end=init_l) + acc = acc[init_l:=UOp.range(acc.size, 500)].set(UOp.const((0.0,)*4, dtypes.float), end=init_l) # create locals (note A is permuted, and the stride is changed to avoid bank conflicts) def make_locals(slot) -> tuple[UOp, UOp]: diff --git a/extra/gemm/moe_routing.py b/extra/gemm/moe_routing.py index 6420437d54..80b55b42c4 100644 --- a/extra/gemm/moe_routing.py +++ b/extra/gemm/moe_routing.py @@ -34,7 +34,7 @@ def _ggather_fwd_kernel(out:UOp, table:UOp, idx:UOp) -> UOp: def _ggather_zero_kernel(out:UOp) -> UOp: i = UOp.range(out.numel(), 0) - return out.flatten().index(i).store(UOp.const(out.dtype, 0.0)).end(i).sink(arg=KernelInfo(name="ggather_zero")) + return out.flatten().index(i).store(UOp.const(0.0, out.dtype)).end(i).sink(arg=KernelInfo(name="ggather_zero")) def _sharded_zeros(shape:tuple[int, ...], dtype, device) -> Tensor: return Tensor.custom_kernel(_sharded_invalids(shape, dtype, device), fxn=_ggather_zero_kernel)[0] diff --git a/extra/hcq2/ops_amd2.py b/extra/hcq2/ops_amd2.py index b62b14b6c9..640a519688 100644 --- a/extra/hcq2/ops_amd2.py +++ b/extra/hcq2/ops_amd2.py @@ -37,7 +37,7 @@ class PM4Ops(FastEnum): RELEASE_MEM = auto(); DISPATCH_DIRECT = auto(); EVENT_WRITE = auto() # noqa: E702 def pkt3(ctx, op:PM4Ops, *vals): - return UOp(Ops.INS, arg=op, src=tuple(UOp.const(dtypes.uint32, x) + return UOp(Ops.INS, arg=op, src=tuple(UOp.const(x, dtypes.uint32) for x in (ctx.pm4.PACKET3(getattr(ctx.pm4, f"PACKET3_{op.name}"), len(vals) - 1), *vals))) def wreg(ctx, reg:AMDReg, *args:sint, **kwargs:int): @@ -157,18 +157,18 @@ def pm4_submit(ctx, lin): assert size_dw < (1 << 20), f"indirect buffer of {size_dw} dwords doesn't fit one packet" ib = UOp.placeholder((size_dw + 2,), dtypes.uint32, next(UOp.unique_num), device=devs, volatile=True).rtag("cmdbuf") - done_idx, submit_idx = UOp.const(dtypes.int, size_dw + 0), UOp.const(dtypes.int, size_dw + 1) - submitted = (counter:=ib.after(make_patches(ib, [((size_dw + i) * 4, UOp.const(dtypes.uint32, 0)) for i in range(2)])).index(submit_idx)).load() + done_idx, submit_idx = UOp.const(size_dw + 0, dtypes.int), UOp.const(size_dw + 1, dtypes.int) + submitted = (counter:=ib.after(make_patches(ib, [((size_dw + i) * 4, UOp.const(0, dtypes.uint32)) for i in range(2)])).index(submit_idx)).load() completed = ib.after(loop:=UOp.loop(0)).index(done_idx).load() ib_free = completed.end(loop, completed != submitted) - bump_fence = pm4_store(ctx, UOp(Ops.SLICE, dtypes.uint32, (ib, UOp.const(dtypes.weakint, size_dw)), 2), (submitted + 1).cast(dtypes.uint64)) + bump_fence = pm4_store(ctx, UOp(Ops.SLICE, dtypes.uint32, (ib, UOp.const(size_dw)), 2), (submitted + 1).cast(dtypes.uint64)) cmdbuf = make_cmdbuf(lin.replace(src=lin.src + (bump_fence,)), devs, buf=ib, dep=ib_free) # the ring itself only carries a packet pointing at the ib, wrapping the ring - put = put_ptr.index(zero:=UOp.const(dtypes.int, 0)) + put = put_ptr.index(zero:=UOp.const(0, dtypes.int)) pkt = (ctx.pm4.PACKET3(ctx.pm4.PACKET3_INDIRECT_BUFFER, 2), *data64_le(cmdbuf.getaddr(devs)), size_dw | ctx.pm4.INDIRECT_BUFFER_VALID) - write_pkt = UOp.barrier(*[ring.index(((put + off) % q.ring.size).cast(dtypes.int)).store(UOp.const(dtypes.uint32, x)) for off,x in enumerate(pkt)]) + write_pkt = UOp.barrier(*[ring.index(((put + off) % q.ring.size).cast(dtypes.int)).store(UOp.const(x, dtypes.uint32)) for off,x in enumerate(pkt)]) # advance the put/write pointers past the packet bump_put_ptr = put_ptr.index(zero).store(put + len(pkt)) @@ -186,26 +186,26 @@ class SDMAOps(FastEnum): COPY = auto(); POLL_REGMEM = auto(); FENCE = auto(); TR def sdma_copy(ctx, call): sz = call.src[2].max_numel() * call.src[2].dtype.itemsize src_addr, dst_addr = call.src[2].getaddr(ctx.devs), call.src[1].getaddr(ctx.devs) - return call.ins(SDMAOps.COPY, src=tuple(UOp.const(dtypes.uint32, x) for off in range(0, sz, ctx.max_copy_size) for x in ( + return call.ins(SDMAOps.COPY, src=tuple(UOp.const(x, dtypes.uint32) for off in range(0, sz, ctx.max_copy_size) for x in ( ctx.sdma.SDMA_OP_COPY | ctx.sdma.SDMA_PKT_COPY_LINEAR_HEADER_SUB_OP(ctx.sdma.SDMA_SUBOP_COPY_LINEAR), ctx.sdma.SDMA_PKT_COPY_LINEAR_COUNT_COUNT(min(sz-off, ctx.max_copy_size)-1), 0, *data64_le(src_addr+off), *data64_le(dst_addr+off)))) def sdma_wait(ctx, ins, dst, val): op = ctx.sdma.SDMA_OP_POLL_REGMEM | ctx.sdma.SDMA_PKT_POLL_REGMEM_HEADER_FUNC(WAIT_REG_MEM_FUNCTION_GEQ) \ | ctx.sdma.SDMA_PKT_POLL_REGMEM_HEADER_MEM_POLL(1) - return ins.ins(SDMAOps.POLL_REGMEM, src=tuple(UOp.const(dtypes.uint32, x) for x in ( + return ins.ins(SDMAOps.POLL_REGMEM, src=tuple(UOp.const(x, dtypes.uint32) for x in ( op, *data64_le(dst.getaddr(ctx.devs)), val, 0xffffffff, ctx.sdma.SDMA_PKT_POLL_REGMEM_DW5_INTERVAL(0x04) | ctx.sdma.SDMA_PKT_POLL_REGMEM_DW5_RETRY_COUNT(0xfff)))) def sdma_store(ctx, ins, dst, val): op = ctx.sdma.SDMA_OP_FENCE | (ctx.sdma.SDMA_PKT_FENCE_HEADER_MTYPE(3) if ctx.target[0] != 9 else 0) return UOp(Ops.LINEAR, src=( - ins.ins(SDMAOps.FENCE, src=tuple(UOp.const(dtypes.uint32, x) for x in (op, *data64_le(dst.getaddr(ctx.devs)), val))), - ins.ins(SDMAOps.TRAP, src=tuple(UOp.const(dtypes.uint32, x) for x in (ctx.sdma.SDMA_OP_TRAP, 0))))) + ins.ins(SDMAOps.FENCE, src=tuple(UOp.const(x, dtypes.uint32) for x in (op, *data64_le(dst.getaddr(ctx.devs)), val))), + ins.ins(SDMAOps.TRAP, src=tuple(UOp.const(x, dtypes.uint32) for x in (ctx.sdma.SDMA_OP_TRAP, 0))))) def sdma_timestamp(ctx, ins, dst): op = ctx.sdma.SDMA_OP_TIMESTAMP | ctx.sdma.SDMA_PKT_TIMESTAMP_GET_HEADER_SUB_OP(ctx.sdma.SDMA_SUBOP_TIMESTAMP_GET_GLOBAL) - return ins.ins(SDMAOps.TIMESTAMP, src=tuple(UOp.const(dtypes.uint32, x) for x in (op, *data64_le(dst.getaddr(ctx.devs))))) + return ins.ins(SDMAOps.TIMESTAMP, src=tuple(UOp.const(x, dtypes.uint32) for x in (op, *data64_le(dst.getaddr(ctx.devs))))) pm_sdma_opsel = PatternMatcher([ (UPat(Ops.CALL, src=(UPat(Ops.COPY),), name="call", allow_any_len=True), sdma_copy), @@ -218,7 +218,7 @@ pm_sdma_opsel = PatternMatcher([ def sdma_submit(cmdbuf, devs): # the cmdbuf to submit + the patch writes that fill it - size_dw, zero = cmdbuf.nbytes() // dtypes.uint32.itemsize, UOp.const(dtypes.int, 0) + size_dw, zero = cmdbuf.nbytes() // dtypes.uint32.itemsize, UOp.const(0, dtypes.int) # the sdma queue's ring and its host-side ring/write/put pointers for d in devs: q = Device[d].sdma_queue(0) @@ -234,8 +234,8 @@ def sdma_submit(cmdbuf, devs): # zero the wrapped tail, then copy the cmdbuf into the ring zi = UOp.range(zero_amt_dw, 0, dtype=dtypes.int, src=(cmdbuf,)) - zero_tail = ring.index(tail_off_dw + zi).store(UOp.const(dtypes.uint32, 0)).end(zi) - i = UOp.range(UOp.const(dtypes.int, size_dw), 0, dtype=dtypes.int, src=(cmdbuf,)) + zero_tail = ring.index(tail_off_dw + zi).store(UOp.const(0, dtypes.uint32)).end(zi) + i = UOp.range(UOp.const(size_dw, dtypes.int), 0, dtype=dtypes.int, src=(cmdbuf,)) copy_to_ring = ring.index(start_dw + i).store(cmdbuf.index(i).load()).end(i) # advance the put/write pointers past the zeroed tail and the cmdbuf diff --git a/extra/llama_kernels/quantize_fp8_delayed/__init__.py b/extra/llama_kernels/quantize_fp8_delayed/__init__.py index 7f2f189c79..6e1c25a1fc 100644 --- a/extra/llama_kernels/quantize_fp8_delayed/__init__.py +++ b/extra/llama_kernels/quantize_fp8_delayed/__init__.py @@ -48,7 +48,7 @@ def _custom_quantize_fp8_with_amax(fp8_out:UOp, amax_out:UOp, x:UOp, amax_state: device = device[0].split(":")[0] if isinstance(device, tuple) else device.split(":")[0] if device in {"AMD", "NULL"}: atomic_arg = "if ({2} > {3}) __hip_atomic_fetch_max((int*){0}, {1}, __ATOMIC_RELAXED, __HIP_MEMORY_SCOPE_AGENT);" else: raise NotImplementedError(f"no atomic max for device {device}") - amax_idx = amax_out.reshape((1,)).index(UOp.const(dtypes.weakint, 0)) + amax_idx = amax_out.reshape((1,)).index(UOp.const(0)) max_val = lds[0].load() atomic = UOp(Ops.CUSTOM, dtypes.void, (amax_idx, max_val.bitcast(dtypes.int32), max_val, amax_idx.load()), arg=atomic_arg) return atomic.end(tid, wg).sink(arg=KernelInfo(f"quantize_fp8_with_amax_{n_elems}", opts_to_apply=())) diff --git a/test/amd/test_emu2_pcode.py b/test/amd/test_emu2_pcode.py index be095c1ac9..5cdf9206e5 100644 --- a/test/amd/test_emu2_pcode.py +++ b/test/amd/test_emu2_pcode.py @@ -11,8 +11,8 @@ from tinygrad.runtime.autogen.amd.rdna3.enum import VOP1Op, VOP2Op, SOP2Op, DSOp def _srcs(): """Create minimal source variables for pcode parsing.""" - def u32(v=0): return UOp.const(dtypes.uint32, v) - return {'S0': u32(), 'S1': u32(), 'S2': u32(), 'SCC': u32(), 'VCC': UOp.const(dtypes.uint64, 0), 'laneId': u32()} + def u32(v=0): return UOp.const(v, dtypes.uint32) + return {'S0': u32(), 'S1': u32(), 'S2': u32(), 'SCC': u32(), 'VCC': UOp.const(0, dtypes.uint64), 'laneId': u32()} class TestBasicParsing(unittest.TestCase): """Test basic pcode parsing for common instruction patterns.""" @@ -44,8 +44,8 @@ class TestWithSources(unittest.TestCase): def test_v_add_f32_with_sources(self): """Test V_ADD_F32 with actual float constants.""" - s0 = UOp.const(dtypes.uint32, 0x3f800000) # 1.0f - s1 = UOp.const(dtypes.uint32, 0x40000000) # 2.0f + s0 = UOp.const(0x3f800000, dtypes.uint32) # 1.0f + s1 = UOp.const(0x40000000, dtypes.uint32) # 2.0f _, assigns = parse_pcode(PCODE[VOP2Op.V_ADD_F32_E32], {'S0': s0, 'S1': s1}) self.assertEqual(len(assigns), 1) dest, val = assigns[0] @@ -55,8 +55,8 @@ class TestWithSources(unittest.TestCase): def test_v_mul_f32_with_sources(self): """Test V_MUL_F32 with actual float constants.""" - s0 = UOp.const(dtypes.uint32, 0x40000000) # 2.0f - s1 = UOp.const(dtypes.uint32, 0x40400000) # 3.0f + s0 = UOp.const(0x40000000, dtypes.uint32) # 2.0f + s1 = UOp.const(0x40400000, dtypes.uint32) # 3.0f _, assigns = parse_pcode(PCODE[VOP2Op.V_MUL_F32_E32], {'S0': s0, 'S1': s1}) self.assertEqual(len(assigns), 1) dest, val = assigns[0] @@ -90,13 +90,13 @@ class TestParseExpr(unittest.TestCase): def test_variable_lookup(self): """Test variable lookup in parse_expr.""" - vrs = {'x': UOp.const(dtypes.uint32, 42)} + vrs = {'x': UOp.const(42, dtypes.uint32)} result = parse_expr('x', vrs) self.assertEqual(result.arg, 42) def test_binary_ops(self): """Test parsing binary operations.""" - vrs = {'a': UOp.const(dtypes.uint32, 10), 'b': UOp.const(dtypes.uint32, 5)} + vrs = {'a': UOp.const(10, dtypes.uint32), 'b': UOp.const(5, dtypes.uint32)} # Addition result = parse_expr('a + b', vrs) @@ -109,7 +109,7 @@ class TestParseExpr(unittest.TestCase): def test_ternary(self): """Test parsing ternary expressions.""" - vrs = {'cond': UOp.const(dtypes.bool, True), 'a': UOp.const(dtypes.uint32, 1), 'b': UOp.const(dtypes.uint32, 0)} + vrs = {'cond': UOp.const(True), 'a': UOp.const(1, dtypes.uint32), 'b': UOp.const(0, dtypes.uint32)} result = parse_expr('cond ? a : b', vrs) self.assertEqual(result.op, Ops.WHERE) @@ -127,7 +127,7 @@ class TestForLoopParsing(unittest.TestCase): def test_clz_parsing(self): """Test CLZ pcode parsing produces correct structure.""" pcode = PCODE[VOP1Op.V_CLZ_I32_U32_E32] - S0 = UOp.const(dtypes.uint32, 0xFFFFFFFF) # All ones - CLZ should be 0 + S0 = UOp.const(0xFFFFFFFF, dtypes.uint32) # All ones - CLZ should be 0 _vrs, assigns = parse_pcode(pcode, {'S0': S0}) self.assertEqual(len(assigns), 1) @@ -139,7 +139,7 @@ class TestForLoopParsing(unittest.TestCase): def test_clz_with_zero(self): """Test CLZ with input 0 - should return -1.""" pcode = PCODE[VOP1Op.V_CLZ_I32_U32_E32] - S0 = UOp.const(dtypes.uint32, 0) + S0 = UOp.const(0, dtypes.uint32) _vrs, assigns = parse_pcode(pcode, {'S0': S0}) # Check that the innermost value (default) is -1 (may be wrapped in CAST) @@ -158,7 +158,7 @@ class TestForLoopParsing(unittest.TestCase): if pcode is None: self.skipTest("V_CTZ_I32_B32_E32 pcode not available") - S0 = UOp.const(dtypes.uint32, 1) # LSB set - CTZ should be 0 + S0 = UOp.const(1, dtypes.uint32) # LSB set - CTZ should be 0 _vrs, assigns = parse_pcode(pcode, {'S0': S0}) self.assertEqual(len(assigns), 1) @@ -169,8 +169,8 @@ class TestDSPcodePatterns(unittest.TestCase): """Test GLOBAL_ATOMIC_ADD_F32 keeps memory values in float dtype.""" vmem = UOp.param(2, dtypes.uint32, (1024,)) srcs = { - 'ADDR': UOp.const(dtypes.uint64, 0), - 'DATA': UOp.const(dtypes.uint32, 0x3f800000), + 'ADDR': UOp.const(0, dtypes.uint64), + 'DATA': UOp.const(0x3f800000, dtypes.uint32), '_vmem': vmem, } @@ -199,8 +199,8 @@ class TestDSPcodePatterns(unittest.TestCase): """Test MEM[addr].type read expression parsing.""" # Create a mock LDS buffer lds = UOp.param(3, dtypes.uint32, (16384,)) - addr = UOp.const(dtypes.uint32, 0) - vrs = {'_lds': lds, 'ADDR': addr, 'OFFSET': UOp.const(dtypes.uint32, 0)} + addr = UOp.const(0, dtypes.uint32) + vrs = {'_lds': lds, 'ADDR': addr, 'OFFSET': UOp.const(0, dtypes.uint32)} result = parse_expr('MEM[ADDR + OFFSET].b32', vrs) # Should be an INDEX operation into LDS @@ -212,13 +212,13 @@ class TestDSPcodePatterns(unittest.TestCase): self.assertIsNotNone(pcode) assert pcode is not None srcs = { - 'ADDR': UOp.const(dtypes.uint32, 0), - 'OFFSET0': UOp.const(dtypes.uint32, 0), - 'OFFSET1': UOp.const(dtypes.uint32, 1), - 'DATA': UOp.const(dtypes.uint32, 0xAAAAAAAA), - 'DATA2': UOp.const(dtypes.uint32, 0xBBBBBBBB), + 'ADDR': UOp.const(0, dtypes.uint32), + 'OFFSET0': UOp.const(0, dtypes.uint32), + 'OFFSET1': UOp.const(1, dtypes.uint32), + 'DATA': UOp.const(0xAAAAAAAA, dtypes.uint32), + 'DATA2': UOp.const(0xBBBBBBBB, dtypes.uint32), } - srcs['laneId'] = UOp.const(dtypes.uint32, 0) + srcs['laneId'] = UOp.const(0, dtypes.uint32) _, assigns = parse_pcode(pcode, srcs) # Should have 2 MEM write assignments self.assertEqual(len(assigns), 2) @@ -235,12 +235,12 @@ class TestDSPcodePatterns(unittest.TestCase): assert pcode is not None lds = UOp.param(3, dtypes.uint32, (16384,)) srcs = { - 'ADDR': UOp.const(dtypes.uint32, 0), - 'OFFSET0': UOp.const(dtypes.uint32, 0), - 'OFFSET1': UOp.const(dtypes.uint32, 1), + 'ADDR': UOp.const(0, dtypes.uint32), + 'OFFSET0': UOp.const(0, dtypes.uint32), + 'OFFSET1': UOp.const(1, dtypes.uint32), '_lds': lds, } - srcs['laneId'] = UOp.const(dtypes.uint32, 0) + srcs['laneId'] = UOp.const(0, dtypes.uint32) _, assigns = parse_pcode(pcode, srcs) # Should have 2 RETURN_DATA assignments self.assertEqual(len(assigns), 2) @@ -252,13 +252,13 @@ class TestDSPcodePatterns(unittest.TestCase): pcode = PCODE.get(DSOp.DS_STORE_2ADDR_B32) assert pcode is not None srcs = { - 'ADDR': UOp.const(dtypes.uint32, 100), - 'OFFSET0': UOp.const(dtypes.uint32, 2), - 'OFFSET1': UOp.const(dtypes.uint32, 5), - 'DATA': UOp.const(dtypes.uint32, 0xAAAAAAAA), - 'DATA2': UOp.const(dtypes.uint32, 0xBBBBBBBB), + 'ADDR': UOp.const(100, dtypes.uint32), + 'OFFSET0': UOp.const(2, dtypes.uint32), + 'OFFSET1': UOp.const(5, dtypes.uint32), + 'DATA': UOp.const(0xAAAAAAAA, dtypes.uint32), + 'DATA2': UOp.const(0xBBBBBBBB, dtypes.uint32), } - srcs['laneId'] = UOp.const(dtypes.uint32, 0) + srcs['laneId'] = UOp.const(0, dtypes.uint32) _, assigns = parse_pcode(pcode, srcs) # Check addresses: 100 + 2*4 = 108, 100 + 5*4 = 120 # assigns[i][1] is (addr, val) tuple for MEM writes; mypy sees UOp @@ -270,13 +270,13 @@ class TestDSPcodePatterns(unittest.TestCase): pcode = PCODE.get(DSOp.DS_STORE_2ADDR_B32) assert pcode is not None srcs = { - 'ADDR': UOp.const(dtypes.uint32, 0), - 'OFFSET0': UOp.const(dtypes.uint32, 0), - 'OFFSET1': UOp.const(dtypes.uint32, 1), - 'DATA': UOp.const(dtypes.uint32, 0xAAAAAAAA), - 'DATA2': UOp.const(dtypes.uint32, 0xBBBBBBBB), + 'ADDR': UOp.const(0, dtypes.uint32), + 'OFFSET0': UOp.const(0, dtypes.uint32), + 'OFFSET1': UOp.const(1, dtypes.uint32), + 'DATA': UOp.const(0xAAAAAAAA, dtypes.uint32), + 'DATA2': UOp.const(0xBBBBBBBB, dtypes.uint32), } - srcs['laneId'] = UOp.const(dtypes.uint32, 0) + srcs['laneId'] = UOp.const(0, dtypes.uint32) _, assigns = parse_pcode(pcode, srcs) # assigns[i][1] is (addr, val) tuple for MEM writes; mypy sees UOp # DATA[31:0] should preserve the value @@ -290,9 +290,9 @@ class TestConditionalParsing(unittest.TestCase): """Test parsing ternary expression (which becomes WHERE).""" # S_CSELECT_B32: D0.u32 = SCC ? S0.u32 : S1.u32 pcode = PCODE[SOP2Op.S_CSELECT_B32] - s0 = UOp.const(dtypes.uint32, 10) - s1 = UOp.const(dtypes.uint32, 20) - scc = UOp.const(dtypes.uint32, 1) + s0 = UOp.const(10, dtypes.uint32) + s1 = UOp.const(20, dtypes.uint32) + scc = UOp.const(1, dtypes.uint32) _vrs, assigns = parse_pcode(pcode, {'S0': s0, 'S1': s1, 'SCC': scc}) self.assertEqual(len(assigns), 1) dest, val = assigns[0] @@ -305,26 +305,26 @@ class TestConcatWidthParsing(unittest.TestCase): def test_permlanex16_altrow_concat(self): for row, expected in [(0, 1), (1, 0), (2, 3), (3, 2)]: - parsed = parse_expr('{ row[1], ~row[0] }', {'row': UOp.const(dtypes.uint32, row)}) + parsed = parse_expr('{ row[1], ~row[0] }', {'row': UOp.const(row, dtypes.uint32)}) self.assertEqual(parsed.simplify().arg, expected) def test_permlane64_altlane_concat(self): for lane, expected in [(0, 32), (1, 33), (31, 63), (32, 0), (63, 31)]: - parsed = parse_expr('{ ~lane[5], lane[4:0] }', {'lane': UOp.const(dtypes.uint32, lane)}) + parsed = parse_expr('{ ~lane[5], lane[4:0] }', {'lane': UOp.const(lane, dtypes.uint32)}) self.assertEqual(parsed.simplify().arg, expected) def test_permlane64_wave64_pcode_indices(self): vgpr = UOp.param(0, dtypes.uint32, (256,)) srcs = { - 'SRC0': UOp.const(dtypes.uint32, 0), - 'VDST': UOp.const(dtypes.uint32, 1), - 'EXEC_LO': UOp.const(dtypes.uint32, 0xFFFFFFFF), - 'EXEC': UOp.const(dtypes.uint64, 0xFFFFFFFFFFFFFFFF), + 'SRC0': UOp.const(0, dtypes.uint32), + 'VDST': UOp.const(1, dtypes.uint32), + 'EXEC_LO': UOp.const(0xFFFFFFFF, dtypes.uint32), + 'EXEC': UOp.const(0xFFFFFFFFFFFFFFFF, dtypes.uint64), '_vgpr': vgpr, '_wave_size': 64, - 'S0': UOp.const(dtypes.uint32, 0), - 'S1': UOp.const(dtypes.uint32, 0), - 'S2': UOp.const(dtypes.uint32, 0), + 'S0': UOp.const(0, dtypes.uint32), + 'S1': UOp.const(0, dtypes.uint32), + 'S2': UOp.const(0, dtypes.uint32), } def load_idx(v: UOp) -> int: @@ -346,7 +346,7 @@ class TestAllPcode(unittest.TestCase): def _make_srcs(self): """Create dummy source variables for pcode parsing.""" - u32, u64 = lambda v=0: UOp.const(dtypes.uint32, v), lambda v=0: UOp.const(dtypes.uint64, v) + u32, u64 = lambda v=0: UOp.const(v, dtypes.uint32), lambda v=0: UOp.const(v, dtypes.uint64) lds = UOp.param(3, dtypes.uint32, (16384,)) return {'laneId': u32(), 'laneID': u32(), 'S0': u32(), 'S1': u32(), 'S2': u32(), 'S3': u32(), 'SRC0': u32(), 'D0': u32(), 'D1': u32(), 'DST': u32(), 'VDST': u32(), 'SDST': u32(), @@ -358,7 +358,7 @@ class TestAllPcode(unittest.TestCase): 'M0': u32(), 'PC': u64(), 'DENORM': u32(1), 'ROUND_MODE': u32(), 'ROUND_TOWARD_ZERO': u32(), 'ROUND_NEAREST_EVEN': u32(), 'WAVE_STATUS': u32(), 'MAX_FLOAT_F32': u32(0x7f7fffff), 'Unsigned': u32(1), 'clampedLOD': u32(), - '_lds': lds, '_vmem': lds, '_active': UOp.const(dtypes.bool, True)} + '_lds': lds, '_vmem': lds, '_active': UOp.const(True)} def _parse_all_pcode(self, pcode_dict, arch: str, min_pct: float): """Parse all pcode. RuntimeError = parser limitation (ok), other exceptions = real bugs.""" diff --git a/test/backend/test_call.py b/test/backend/test_call.py index 0a9e039b9f..53d8bb0ff2 100644 --- a/test/backend/test_call.py +++ b/test/backend/test_call.py @@ -6,11 +6,11 @@ from tinygrad.renderer.cstyle import CStyleLanguage from tinygrad.uop.ops import KernelInfo def call_out_kernel(F:UOp, C:UOp) -> UOp: - call = F[0].load().call(UOp.const(dtypes.int, 3), C[0], ret_dtype=dtypes.void) + call = F[0].load().call(UOp.const(3, dtypes.int), C[0], ret_dtype=dtypes.void) return C.after(call)[1].store(C.after(call)[0].load() + 1).sink(arg=KernelInfo(name="call_out")) def call_ret_kernel(F:UOp, C:UOp) -> UOp: - val = F[0].load().call(UOp.const(dtypes.int, 21), ret_dtype=dtypes.int) + val = F[0].load().call(UOp.const(21, dtypes.int), ret_dtype=dtypes.int) return C[0].store(val * 2).sink(arg=KernelInfo(name="call_ret")) @unittest.skipUnless(isinstance(Device["CPU"].renderer, CStyleLanguage), "TODO: CALL is rendered in C style only") diff --git a/test/backend/test_const_folding.py b/test/backend/test_const_folding.py index 4a8d4a350f..0e4239ecfe 100644 --- a/test/backend/test_const_folding.py +++ b/test/backend/test_const_folding.py @@ -16,7 +16,7 @@ def _check_ast_count(desired_count:int, t:Tensor): class TestMovedConstFolding(unittest.TestCase): def test_contiguous_deviceless_const(self): - t = Tensor(UOp.const(dtypes.float, 2.0)).contiguous() + t = Tensor(UOp.const(2.0, dtypes.float)).contiguous() self.assertIs(t.uop.op, Ops.CONST) self.assertIsNone(t.uop.device) @@ -169,7 +169,7 @@ class TestMultiConstFolding(unittest.TestCase): class TestThreefryConstFolding(unittest.TestCase): def test_threefry(self): # THREEFRY(const,const) folds to a const once decomposed - x = threefry2x32(UOp.const(dtypes.uint64, 5), UOp.const(dtypes.uint64, 10)) + x = threefry2x32(UOp.const(5, dtypes.uint64), UOp.const(10, dtypes.uint64)) self.assertIs(x.simplify().op, Ops.CONST) class TestTautologicalCompare(unittest.TestCase): diff --git a/test/backend/test_custom_kernel.py b/test/backend/test_custom_kernel.py index 7e087da373..0d9345ec73 100644 --- a/test/backend/test_custom_kernel.py +++ b/test/backend/test_custom_kernel.py @@ -341,7 +341,7 @@ class TestCustomKernel(unittest.TestCase): def test_partial_invalid_store_keeps_uncovered_reads(self): x = Tensor([10., 20., 30., 40.]) - after = x.uop.after(x.uop.shrink(((0, 2),)).store(UOp.const(dtypes.float, Invalid, shape=(2,)))) + after = x.uop.after(x.uop.shrink(((0, 2),)).store(UOp.const(Invalid, dtypes.float, shape=(2,)))) self.assertEqual(Tensor(after).contiguous().tolist(), [10., 20., 30., 40.]) def test_multi_after_invalid_store_dep_removed(self): diff --git a/test/backend/test_isel.py b/test/backend/test_isel.py index 9eb629b303..cd357af06d 100644 --- a/test/backend/test_isel.py +++ b/test/backend/test_isel.py @@ -7,7 +7,7 @@ from tinygrad.renderer.isa.x86 import X86Renderer, X86Ops from tinygrad.renderer.isa import IselContext # INDEX on a register value with a constant index extracts a single element (the old GEP) -def lane(y:UOp, i:int) -> UOp: return y.index(UOp.const(dtypes.int, i), dtype=y.dtype.scalar()) +def lane(y:UOp, i:int) -> UOp: return y.index(UOp.const(i, dtypes.int), dtype=y.dtype.scalar()) @unittest.skipUnless(isinstance(Device[Device.DEFAULT].renderer, X86Renderer), "only x86") class TestIselX86(unittest.TestCase): diff --git a/test/backend/test_linearizer_dumb.py b/test/backend/test_linearizer_dumb.py index b6128cd66a..6604171484 100644 --- a/test/backend/test_linearizer_dumb.py +++ b/test/backend/test_linearizer_dumb.py @@ -12,18 +12,18 @@ class TestLinearizerFailure(unittest.TestCase): @unittest.skipUnless(Device.DEFAULT == "METAL", "only tested on METAL") def test_failure_beam_mnist(self): c0 = UOp.param(0, dtypes.uchar, (4014080,)) - c1 = UOp.range(UOp.const(dtypes.weakint, 512), 0, AxisType.GLOBAL) - c2 = UOp.range(UOp.const(dtypes.weakint, 784), 1, AxisType.GLOBAL) - c3 = UOp.range(UOp.const(dtypes.weakint, 10), 3, AxisType.GLOBAL) + c1 = UOp.range(UOp.const(512), 0, AxisType.GLOBAL) + c2 = UOp.range(UOp.const(784), 1, AxisType.GLOBAL) + c3 = UOp.range(UOp.const(10), 3, AxisType.GLOBAL) c4 = UOp.param(1, dtypes.int, (512,)) - c5 = c4.index(c1.valid(UOp.const(dtypes.bool, True))) - c6 = UOp.range(UOp.const(dtypes.weakint, 6000), 1004, AxisType.REDUCE) - c7 = UOp.range(UOp.const(dtypes.weakint, 3750), 2006, AxisType.REDUCE) - c8 = UOp.range(UOp.const(dtypes.weakint, 16), 2007, AxisType.GROUP_REDUCE) + c5 = c4.index(c1.valid(UOp.const(True))) + c6 = UOp.range(UOp.const(6000), 1004, AxisType.REDUCE) + c7 = UOp.range(UOp.const(3750), 2006, AxisType.REDUCE) + c8 = UOp.range(UOp.const(16), 2007, AxisType.GROUP_REDUCE) c9 = UOp.param(2, dtypes.uchar, (47040000,)) - c10 = c9.index((((c3*UOp.const(dtypes.weakint, 4704000))+c2)+(c6*UOp.const(dtypes.weakint, 784))).valid(UOp.const(dtypes.bool, True))) - c11 = c5.alu(Ops.CMPNE, ((((c3*UOp.const(dtypes.weakint, 6000))+c6)+((c7*UOp.const(dtypes.weakint, 16))+c8)).alu(Ops.CMPLT, UOp.const(dtypes.weakint, 59999)).where(UOp.const(dtypes.int, 0), UOp.const(dtypes.int, 1)).reduce(c7, c8, arg=Ops.ADD)+UOp.const(dtypes.int, -1))).where(UOp.const(dtypes.uchar, 0), c10).reduce(c6, arg=Ops.ADD) - c12 = c0.index((((c1*UOp.const(dtypes.weakint, 7840))+(c2*UOp.const(dtypes.weakint, 10)))+c3).valid(UOp.const(dtypes.bool, True))).store(c11).end(c1, c2, c3) + c10 = c9.index((((c3*UOp.const(4704000))+c2)+(c6*UOp.const(784))).valid(UOp.const(True))) + c11 = c5.alu(Ops.CMPNE, ((((c3*UOp.const(6000))+c6)+((c7*UOp.const(16))+c8)).alu(Ops.CMPLT, UOp.const(59999)).where(UOp.const(0, dtypes.int), UOp.const(1, dtypes.int)).reduce(c7, c8, arg=Ops.ADD)+UOp.const(-1, dtypes.int))).where(UOp.const(0, dtypes.uchar), c10).reduce(c6, arg=Ops.ADD) + c12 = c0.index((((c1*UOp.const(7840))+(c2*UOp.const(10)))+c3).valid(UOp.const(True))).store(c11).end(c1, c2, c3) ast = c12.sink(arg=KernelInfo(name='test', axis_types=(), dont_use_locals=False, applied_opts=(Opt(op=OptOps.GROUP, axis=1, arg=16),), opts_to_apply=None)) _ = to_program(ast, Device["METAL"].renderer) diff --git a/test/backend/test_pickle.py b/test/backend/test_pickle.py index 9f31ce53c7..435270bfea 100644 --- a/test/backend/test_pickle.py +++ b/test/backend/test_pickle.py @@ -13,7 +13,7 @@ class TestPickle(unittest.TestCase): def test_pickle_pattern_matcher(self): pm = PatternMatcher([(UPat.cvar('x'), lambda x: x*2)]) - sink = UOp.const(dtypes.int, 2) + sink = UOp.const(2, dtypes.int) tt = pm.rewrite(sink) pm_str = pickle.dumps(pm) pm2 = pickle.loads(pm_str) diff --git a/test/backend/test_renderer_failures.py b/test/backend/test_renderer_failures.py index b347dcc8d1..8d3c09d4b0 100644 --- a/test/backend/test_renderer_failures.py +++ b/test/backend/test_renderer_failures.py @@ -24,7 +24,7 @@ def _setup_and_test_alu(alu_op:Ops, input_val:ConstType, *alu_src_uops:UOp): dtype = alu_src_uops[0].dtype a = UOp.param(0, dtype, (1,)) b = UOp.param(1, dtype, (1,)) - idx = UOp.const(dtypes.int, 0) + idx = UOp.const(0, dtypes.int) ld = b.index(idx).load() alu = ld.alu(alu_op, *alu_src_uops) store = UOp.store(a.index(idx), alu) @@ -35,7 +35,7 @@ class TestRendererFailures(unittest.TestCase): def test_gated_store_with_alu(self): a = UOp.param(0, dtypes.int, (4,)) gate_alu = (lidx0:=UOp.special(4, 'lidx0')).ne(0) - gated_alu_store = UOp(Ops.STORE, src=(a.index(lidx0.valid(gate_alu)), UOp.const(dtypes.int, 1))) + gated_alu_store = UOp(Ops.STORE, src=(a.index(lidx0.valid(gate_alu)), UOp.const(1, dtypes.int))) sink = UOp(Ops.SINK, src=(gated_alu_store,), arg=KernelInfo()) ret = _test_uop_result([], sink, local_size=[4, 1, 1])[0] np.testing.assert_equal(ret, [0, 1, 1, 1]) @@ -45,7 +45,7 @@ class TestRendererFailures(unittest.TestCase): a = UOp.param(0, dtypes.int, (8,)) gate_alu_0 = (lidx0:=UOp.special(4, 'lidx0')).ne(0) gate_alu_1 = (lidx1:=UOp.special(2, 'lidx1')).ne(0) - gated_alu_store = UOp(Ops.STORE, src=(a.index((lidx0+lidx1*4).valid(gate_alu_0&gate_alu_1)), UOp.const(dtypes.int, 1))) + gated_alu_store = UOp(Ops.STORE, src=(a.index((lidx0+lidx1*4).valid(gate_alu_0&gate_alu_1)), UOp.const(1, dtypes.int))) sink = UOp(Ops.SINK, src=(gated_alu_store,), arg=KernelInfo()) ret = _test_uop_result([], sink, local_size=[4, 2, 1])[0] np.testing.assert_equal(ret, [0, 0, 0, 0, 0, 1, 1, 1]) @@ -54,7 +54,7 @@ class TestRendererFailures(unittest.TestCase): class TestCStyleFailures(unittest.TestCase): def test_inline_const_alu(self): # CPU doesn't use the max function - ret = _setup_and_test_alu(Ops.MAX, 1, UOp.const(dtypes.int, dtypes.int.min+1)) + ret = _setup_and_test_alu(Ops.MAX, 1, UOp.const(dtypes.int.min+1, dtypes.int)) self.assertEqual(ret[0], 1) def _test_src_strip_paren(self, op: Ops, should_strip_paren:bool=True): @@ -80,7 +80,7 @@ class TestWGSLFailures(unittest.TestCase): def test_multiply_infinity(self): # multiplying a positive constant by infinity should return infinity # WGSL pipelines do not handle this reliably, some of which return zero, unless infinity always comes from a read on a dynamic buffer - ret = _setup_and_test_alu(Ops.MUL, 5.0, UOp.const(dtypes.float32, float("inf"))) + ret = _setup_and_test_alu(Ops.MUL, 5.0, UOp.const(float("inf"), dtypes.float32)) self.assertEqual(ret[0], float("inf")) # WGSL has a specific select(alt, val, gate) ternary operator instead of gate?val:alt @@ -104,7 +104,7 @@ class TestPTXFailures(unittest.TestCase): def test_gated_store_with_if(self): a = UOp.param(0, dtypes.int, (4,)) gate_alu = (lidx0:=UOp.special(4, 'lidx0')).ne(0) - val = UOp.const(dtypes.int, 1) + val = UOp.const(1, dtypes.int) if_uop = UOp(Ops.IF, src=(gate_alu,)) gated_alu_store = UOp(Ops.STORE, src=(a.index(lidx0, if_uop), val)) sink = UOp(Ops.SINK, src=(gated_alu_store,), arg=KernelInfo()) diff --git a/test/backend/test_tensor.py b/test/backend/test_tensor.py index c1f65c5dc6..421343706c 100644 --- a/test/backend/test_tensor.py +++ b/test/backend/test_tensor.py @@ -24,13 +24,13 @@ class TestTinygrad(unittest.TestCase): self.assertEqual(Tensor(3.14).shape, ()) def test_deviceless_const_construct_device_repr(self): - t = Tensor(UOp.const(dtypes.float, 2.0)) + t = Tensor(UOp.const(2.0, dtypes.float)) self.assertIsNone(t.uop.device) self.assertIsNone(t.device) self.assertIn(" UOp: - if op is Ops.CONST: uops.append(UOp.const(dtype, arg)) + if op is Ops.CONST: uops.append(UOp.const(arg, dtype)) elif op is Ops.PARAM: uops.append(UOp.param(arg, dtype, shape=(1,))) else: uops.append(UOp(op, dtype, tuple(src), arg)) return uops[-1] @@ -43,7 +43,7 @@ def _test_single_value_const(vals, op, dts): buf_store = uop(uops, Ops.PARAM, output_dtype, (), 0) loads = (uop(uops, Ops.CONST, dtype, [], a) for a,dtype in zip(vals, dts)) alu = uop(uops, op, output_dtype, loads) - out = buf_store[UOp.const(dtypes.int32, 0)].store(alu) + out = buf_store[UOp.const(0, dtypes.int32)].store(alu) buf = Buffer(Device.DEFAULT, 1, output_dtype).allocate() run_uops([out], [buf]) return np.frombuffer(buf.as_memoryview(), _to_np_dtype(output_dtype))[0] @@ -221,12 +221,12 @@ class TestAssembly(unittest.TestCase): def test_bitshift_left(self): g1 = UOp.param(0, dtypes.int32, shape=(3,)) out = UOp.param(1, dtypes.int32, shape=(2,)) - c1 = UOp.const(dtypes.int, 2) - c2 = UOp.const(dtypes.int, 3) + c1 = UOp.const(2, dtypes.int) + c2 = UOp.const(3, dtypes.int) l1 = g1.index(c1) a1 = UOp(Ops.MUL, src=(l1, c1)) a2 = UOp(Ops.MUL, src=(l1, c2)) - uops = to_uops_list([out.index(UOp.const(dtypes.int, 0)).store(a1), out.index(UOp.const(dtypes.int, 1)).store(a2)], + uops = to_uops_list([out.index(UOp.const(0, dtypes.int)).store(a1), out.index(UOp.const(1, dtypes.int)).store(a2)], ren=Device[Device.DEFAULT].renderer) Device[Device.DEFAULT].renderer.render(uops) ops = [x.op for x in uops] @@ -249,16 +249,16 @@ class TestAssembly(unittest.TestCase): def test_mulacc_shl(self): g1 = UOp.param(0, dtypes.int32, shape=(2,)) - c1 = UOp.const(dtypes.int, 0) - c2 = UOp.const(dtypes.int, 1) - expr = g1.index(c1) * UOp.const(dtypes.int, 4096) + g1.index(c2) + c1 = UOp.const(0, dtypes.int) + c2 = UOp.const(1, dtypes.int) + expr = g1.index(c1) * UOp.const(4096, dtypes.int) + g1.index(c2) uops = to_uops_list([expr], ren=Device[Device.DEFAULT].renderer) Device[Device.DEFAULT].renderer.render(uops) self.assertIn(Ops.MULACC, [x.op for x in uops]) def test_use_cmpeq(self): g = UOp.param(0, dtypes.uint32, shape=(8,)) - c = UOp.const(dtypes.uint, 7) + c = UOp.const(7, dtypes.uint) comp = g.index(c).ne(c).ne(True) uops = to_uops_list([comp], ren=Device[Device.DEFAULT].renderer) Device[Device.DEFAULT].renderer.render(uops) diff --git a/test/external/external_benchmark_op_conv.py b/test/external/external_benchmark_op_conv.py index bc7e774862..4076facd5d 100644 --- a/test/external/external_benchmark_op_conv.py +++ b/test/external/external_benchmark_op_conv.py @@ -24,8 +24,8 @@ def vision_conv_143(): c32 = ((c27<3)!=True)&(c27<67) c34 = UOp.param(1, dtypes.half, shape=(32, 1024, 4)) c38 = c5//2 - c45 = (c32&c24).where((c27*64+c38+c17*4096+-12480), UOp.const(dtypes.weakint, Invalid)) - c48 = (c24&c32).where(c34.index(c45), UOp.const(dtypes.float, 0.0)) + c45 = (c32&c24).where((c27*64+c38+c17*4096+-12480), UOp.const(Invalid, dtypes.weakint)) + c48 = (c24&c32).where(c34.index(c45), UOp.const(0.0, dtypes.float)) c49 = UOp.param(2, dtypes.half, shape=(64, 49, 4)) c61 = c48*c49.index((c26*4+c5%2+c16*28+c38*196)) c63 = UOp.param(3, dtypes.float, (128,)) @@ -50,8 +50,8 @@ def vision_conv_153(): c32 = ((c27<3)!=True)&(c27<35) c34 = UOp.param(1, dtypes.half, shape=(16, 1024, 4)) c38 = c5//2 - c45 = (c32&c24).where((c27*128+c38+c17*4096+-12672), UOp.const(dtypes.weakint, Invalid)) - c48 = (c24&c32).where(c34.index(c45), UOp.const(dtypes.float, 0.0)) + c45 = (c32&c24).where((c27*128+c38+c17*4096+-12672), UOp.const(Invalid, dtypes.weakint)) + c48 = (c24&c32).where(c34.index(c45), UOp.const(0.0, dtypes.float)) c49 = UOp.param(2, dtypes.half, shape=(128, 49, 4)) c61 = c48*c49.index((c26*4+c5%2+c16*28+c38*196)) c63 = UOp.param(3, dtypes.float, (256,)) diff --git a/test/external/fuzz_symbolic.py b/test/external/fuzz_symbolic.py index d9a3144667..08c78a4c5d 100644 --- a/test/external/fuzz_symbolic.py +++ b/test/external/fuzz_symbolic.py @@ -40,7 +40,7 @@ def random_int_expr(depth=10): def random_bool_expr(depth=10, expr1=None): if depth == 0: return True if expr1 is None: expr1 = random_int_expr(depth-1) - expr2 = random.choice([random_or_sub_expression_int(depth-1, expr1), UOp.const(dtypes.weakint, random.randint(-10, 10))]) + expr2 = random.choice([random_or_sub_expression_int(depth-1, expr1), UOp.const(random.randint(-10, 10))]) return random.choice(comp_ops)(expr1, expr2) @@ -82,8 +82,8 @@ if __name__ == "__main__": f"v2=Variable(\"{u2.arg[0]}\", {u2.arg[1]}, {u2.arg[2]})\n" +\ f"v3=Variable(\"{u3.arg[0]}\", {u3.arg[1]}, {u3.arg[2]})\n" +\ f"expr = {expr}\n" +\ - f"v1_val, v2_val, v3_val = UOp.const(dtypes.weakint, {n1.as_long()}), UOp.const(dtypes.weakint, {n2.as_long()})," +\ - f"UOp.const(dtypes.weakint, {n3.as_long()})\n" +\ + f"v1_val, v2_val, v3_val = UOp.const({n1.as_long()}), UOp.const({n2.as_long()})," +\ + f"UOp.const({n3.as_long()})\n" +\ "num = expr.simplify().substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify()\n" +\ "rn = expr.substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify()\n" +\ "assert num==rn, f\"{num} != {rn}\"\n" diff --git a/test/helpers.py b/test/helpers.py index 626ef789ab..88c92787dd 100644 --- a/test/helpers.py +++ b/test/helpers.py @@ -90,7 +90,7 @@ def eval_uop(uop:UOp, inputs:list[tuple[DType, list[Any]]]|None=None, vals:tuple bufs.append(buf:=allocator.alloc(len(data) * buf_dt.itemsize)) allocator._copyin(buf, memoryview(struct.pack(str(len(data)) + (buf_dt.fmt or ""), *data))) g = UOp.param(0, uop.dtype, (1,)) - prg = to_program(UOp.store(g.index(UOp.const(dtypes.int, 0)), uop).sink(arg=KernelInfo()), PythonRenderer(Target("PYTHON"))) + prg = to_program(UOp.store(g.index(UOp.const(0, dtypes.int)), uop).sink(arg=KernelInfo()), PythonRenderer(Target("PYTHON"))) prog = dev.runtime(prg.to_elf()) prog(out_buf:=allocator.alloc(uop.dtype.itemsize), *bufs, vals=vals) return out_buf.cast(uop.dtype.fmt or "").tolist()[0] diff --git a/test/mockgpu/amd/emu.py b/test/mockgpu/amd/emu.py index 61dda4f20e..c6add529c1 100644 --- a/test/mockgpu/amd/emu.py +++ b/test/mockgpu/amd/emu.py @@ -192,16 +192,16 @@ def _init_sqtt_encoder(): return emit, finish, finalize -def _c(val, dtype=dtypes.uint32): return UOp.const(dtype, val) +def _c(val, dtype=dtypes.uint32): return UOp.const(val, dtype) def _u64(lo: UOp, hi: UOp) -> UOp: """Combine two 32-bit UOps into a 64-bit UOp.""" - return lo.cast(dtypes.uint64) | (hi.cast(dtypes.uint64) << UOp.const(dtypes.uint64, 32)) + return lo.cast(dtypes.uint64) | (hi.cast(dtypes.uint64) << UOp.const(32, dtypes.uint64)) def _split64(val: UOp) -> tuple[UOp, UOp]: """Split a 64-bit value into (lo, hi) 32-bit values.""" v64 = val.bitcast(dtypes.uint64) if val.dtype == dtypes.float64 else val.cast(dtypes.uint64) if val.dtype != dtypes.uint64 else val - return v64.cast(dtypes.uint32), (v64 >> UOp.const(dtypes.uint64, 32)).cast(dtypes.uint32) + return v64.cast(dtypes.uint32), (v64 >> UOp.const(32, dtypes.uint64)).cast(dtypes.uint32) _SRC_MOD_TYPES = {16: (dtypes.uint16, dtypes.half, 0x7FFF), 32: (dtypes.uint32, dtypes.float32, 0x7FFFFFFF), 64: (dtypes.uint64, dtypes.float64, 0x7FFFFFFFFFFFFFFF)} @@ -210,7 +210,7 @@ def _apply_src_mods(val: UOp, mod_bit: int, abs_bits: int, neg_bits: int, bits: if not (abs_bits & (1 << mod_bit)) and not (neg_bits & (1 << mod_bit)): return val ut, ft, mask = _SRC_MOD_TYPES[bits] fv = val.cast(ut).bitcast(ft) if bits == 16 else val.bitcast(ft) if val.dtype == ut else val - if abs_bits & (1 << mod_bit): fv = (fv.bitcast(ut) & UOp.const(ut, mask)).bitcast(ft) + if abs_bits & (1 << mod_bit): fv = (fv.bitcast(ut) & UOp.const(mask, ut)).bitcast(ft) if neg_bits & (1 << mod_bit): fv = fv.neg() return fv.bitcast(ut).cast(dtypes.uint32) if bits == 16 else fv.bitcast(ut) @@ -251,7 +251,7 @@ def _to_u32(val: UOp) -> UOp: if val.dtype.itemsize == 4: return val.bitcast(dtypes.uint32) # same size: bitcast (float32->uint32) return val.cast(dtypes.uint32) # different size: cast (bool, int16, etc) def _lane_active(exec_mask: UOp, lane: UOp) -> UOp: - if exec_mask.dtype == dtypes.uint64: return ((exec_mask >> lane.cast(dtypes.uint64)) & UOp.const(dtypes.uint64, 1)).ne(UOp.const(dtypes.uint64, 0)) + if exec_mask.dtype == dtypes.uint64: return ((exec_mask >> lane.cast(dtypes.uint64)) & UOp.const(1, dtypes.uint64)).ne(UOp.const(0, dtypes.uint64)) return ((exec_mask >> lane.cast(dtypes.uint32)) & _c(1)).ne(_c(0)) def _hi16(v: UOp) -> UOp: return (v >> _c(16)) & _c(0xFFFF) def _cond(cond, if_true, if_false): @@ -264,9 +264,9 @@ def _set_lane_bit(old: UOp, lane: UOp, val: UOp, exec_mask: UOp) -> UOp: """Set/clear a single bit in a mask based on lane index, respecting exec mask.""" if old.dtype in (dtypes.uint64, dtypes.int64): dt = dtypes.uint64 - mask = UOp.const(dt, 1) << lane.cast(dt) + mask = UOp.const(1, dt) << lane.cast(dt) new_bit = _to_u32(val).cast(dt) << lane.cast(dt) - cleared = old.cast(dt) & (mask ^ UOp.const(dt, 0xFFFFFFFFFFFFFFFF)) + cleared = old.cast(dt) & (mask ^ UOp.const(0xFFFFFFFFFFFFFFFF, dt)) return _lane_active(exec_mask, lane).where(cleared | new_bit, old.cast(dt)) mask = _c(1) << lane.cast(dtypes.uint32) new_bit = _to_u32(val) << lane.cast(dtypes.uint32) @@ -365,7 +365,7 @@ def _write_64bit(val: UOp, wfn, reg_or_addr, is_mem: bool, *args) -> list[UOp]: """Write a 64-bit value as two 32-bit writes. args passed to wfn after reg/addr and lo/hi value.""" lo, hi = _split64(val) incr = 4 if is_mem else 1 # 4 bytes for memory addresses, 1 for register indices - return [wfn(reg_or_addr, lo, *args), wfn(reg_or_addr + (UOp.const(reg_or_addr.dtype, incr) if isinstance(reg_or_addr, UOp) else incr), hi, *args)] + return [wfn(reg_or_addr, lo, *args), wfn(reg_or_addr + (UOp.const(incr, reg_or_addr.dtype) if isinstance(reg_or_addr, UOp) else incr), hi, *args)] def _write_val(bits: int, val: UOp, wfn, reg_or_addr, *args, is_mem: bool = False) -> list[UOp]: """Write value, splitting 64-bit if needed. bits=64 for 64-bit writes, otherwise 32-bit.""" @@ -374,7 +374,7 @@ def _write_val(bits: int, val: UOp, wfn, reg_or_addr, *args, is_mem: bool = Fals def _mem_store(mem: UOp, addr: UOp, val: UOp, active: UOp, addr_bits: int = 32, data_bits: int = 32) -> list[UOp]: """Conditional memory store with sub-word support. Returns list of store UOps.""" adt = dtypes.uint64 if addr_bits == 64 else dtypes.uint32 - word_addr = addr >> UOp.const(adt, 2) + word_addr = addr >> UOp.const(2, adt) idx = mem.index(word_addr.valid(active)) if data_bits == 32: return [idx.store(active.where(_to_u32(val), idx))] # Sub-word store: read-modify-write with mask @@ -388,7 +388,7 @@ def _mem_store(mem: UOp, addr: UOp, val: UOp, active: UOp, addr_bits: int = 32, is_cross = byte_pos.eq(_c(3)) cross_word0 = (idx & _c(0x00FFFFFF)) | ((val_u32 & _c(0xFF)) << _c(24)) store0 = idx.store(active.where(is_cross.where(cross_word0, new_word), idx)) - next_idx = mem.index((word_addr + UOp.const(adt, 1)).valid(active & is_cross)) + next_idx = mem.index((word_addr + UOp.const(1, adt)).valid(active & is_cross)) cross_word1 = (next_idx & _c(0xFFFFFF00)) | ((val_u32 >> _c(8)) & _c(0xFF)) return [store0, next_idx.store((active & is_cross).where(cross_word1, next_idx))] @@ -397,8 +397,8 @@ def _mem_store_bytes(mem: UOp, addr: UOp, val: UOp, active: UOp, data_bits: int stores = [] val_u32 = val.cast(dtypes.uint32) if val.dtype != dtypes.uint32 else val for i in range(data_bits // 8): - byte_val = (val_u32 >> UOp.const(dtypes.uint32, i * 8)) & UOp.const(dtypes.uint32, 0xFF) - stores.append(mem.index((addr + UOp.const(dtypes.uint64, i)).valid(active)).store(byte_val.cast(dtypes.uint8))) + byte_val = (val_u32 >> UOp.const(i * 8, dtypes.uint32)) & UOp.const(0xFF, dtypes.uint32) + stores.append(mem.index((addr + UOp.const(i, dtypes.uint64)).valid(active)).store(byte_val.cast(dtypes.uint8))) return stores def _collect_data_slices(assigns: list[tuple[str, UOp]], data_prefix: str, pcode_vars: dict | None = None, op_name: str = "") -> dict[int, UOp]: @@ -462,8 +462,8 @@ class _Ctx: def inst_word(self, dword_idx: int) -> UOp: """Read instruction dword from vmem at PC + dword_idx*4.""" pc = self.rpc() - addr = pc if dword_idx == 0 else pc + UOp.const(dtypes.uint64, dword_idx * 4) - return self.vmem.index(addr >> UOp.const(dtypes.uint64, 2)).load() + addr = pc if dword_idx == 0 else pc + UOp.const(dword_idx * 4, dtypes.uint64) + return self.vmem.index(addr >> UOp.const(2, dtypes.uint64)).load() def inst_field(self, field) -> UOp: """Extract field bits from instruction encoding. Tracks field for canonical key computation.""" @@ -475,15 +475,15 @@ class _Ctx: word = self.inst_word(dword_idx) if lo // 32 == hi // 32: # Same dword mask = (1 << (hi - lo + 1)) - 1 - shifted = word if lo_in_dword == 0 else word >> UOp.const(dtypes.uint32, lo_in_dword) - return shifted & UOp.const(dtypes.uint32, mask) + shifted = word if lo_in_dword == 0 else word >> UOp.const(lo_in_dword, dtypes.uint32) + return shifted & UOp.const(mask, dtypes.uint32) else: # Spans two dwords lo_bits = 32 - lo_in_dword lo_mask = (1 << lo_bits) - 1 hi_mask = (1 << (hi_in_dword + 1)) - 1 - lo_part = (word >> UOp.const(dtypes.uint32, lo_in_dword)) & UOp.const(dtypes.uint32, lo_mask) - hi_part = self.inst_word(dword_idx + 1) & UOp.const(dtypes.uint32, hi_mask) - return lo_part | (hi_part << UOp.const(dtypes.uint32, lo_bits)) + lo_part = (word >> UOp.const(lo_in_dword, dtypes.uint32)) & UOp.const(lo_mask, dtypes.uint32) + hi_part = self.inst_word(dword_idx + 1) & UOp.const(hi_mask, dtypes.uint32) + return lo_part | (hi_part << UOp.const(lo_bits, dtypes.uint32)) def inst_field_signed(self, field) -> UOp: """Extract field and sign-extend based on field width.""" @@ -584,7 +584,7 @@ class _Ctx: inline = is_float_const.where(float_inline.bitcast(dtypes.uint64), int_inline.bitcast(dtypes.uint64)) # Literal handling: F64 VOP puts literal in high 32 bits; B64/I64/U64 VOP and SOP zero-extend if literal is not None: - lit_val = literal.cast(dtypes.uint64) << UOp.const(dtypes.uint64, 32) if is_f64 else literal.cast(dtypes.uint64) + lit_val = literal.cast(dtypes.uint64) << UOp.const(32, dtypes.uint64) if is_f64 else literal.cast(dtypes.uint64) inline = off.eq(_c(255)).where(lit_val, inline) scalar_val = (off < _c(128)).where(sgpr_val, inline) else: @@ -602,7 +602,7 @@ class _Ctx: def inc_pc(self) -> list[UOp]: """Increment PC by instruction size in bytes. Returns [store].""" - new_pc = self.rpc() + UOp.const(dtypes.uint64, self.inst_size) + new_pc = self.rpc() + UOp.const(self.inst_size, dtypes.uint64) lo, hi = _split64(new_pc) return [self.wsgpr_dyn(_c(PC_LO_IDX), lo), self.wsgpr_dyn(_c(PC_HI_IDX), hi)] @@ -668,7 +668,7 @@ class _Ctx: if 'VCC' not in srcs: srcs['VCC'] = self.rmask(_c(vcc_reg)) srcs.update({'EXEC': exec_mask, 'SCC': self.rsgpr_dyn(_c(SCC.offset)), 'laneId': lane, 'VDST': vdst_reg, 'ROUND_MODE': _c(0), 'ROUND_TOWARD_ZERO': _c(0), 'ROUND_NEAREST_EVEN': _c(0), '_vgpr': self.vgpr, '_wave_size': self.wave_size, - 'MAX_FLOAT_F32': UOp.const(dtypes.float32, 3.4028234663852886e38), + 'MAX_FLOAT_F32': UOp.const(3.4028234663852886e38, dtypes.float32), # CDNA SDWA byte/word select constants (E32 always uses BYTE0/WORD0 defaults) 'SDWA_SRC0_SEL': _c(0), 'BYTE0': _c(0), 'BYTE1': _c(1), 'BYTE2': _c(2), 'BYTE3': _c(3), 'WORD0': _c(0), 'WORD1': _c(1)}) # rounding mode and SDWA constants @@ -732,24 +732,24 @@ class _Ctx: d0_width, slice_mask = d0_hi_bit - d0_lo_bit + 1, (1 << (d0_hi_bit - d0_lo_bit + 1)) - 1 val_bits = val.bitcast(dtypes.uint16).cast(dtypes.uint32) if val.dtype == dtypes.half else \ val.cast(dtypes.uint32) if val.dtype in (dtypes.uint16, dtypes.int16) else \ - val.cast(dtypes.uint32) & UOp.const(dtypes.uint32, slice_mask) + val.cast(dtypes.uint32) & UOp.const(slice_mask, dtypes.uint32) raw_stores.append(('vgpr_slice', (d0_lo_bit, d0_width, val_bits))) continue # For integer ops with clamp, use pre-computed saturated value; for floats, clamp to [0,1] if int_saturate is not None: val = int_saturate elif clmp and val.dtype in (dtypes.float32, dtypes.half, dtypes.float64): - clamped = val.maximum(UOp.const(val.dtype, 0.0)).minimum(UOp.const(val.dtype, 1.0)) - val = _FUNCS['isNAN'](val).where(UOp.const(val.dtype, 0.0), clamped) + clamped = val.maximum(UOp.const(0.0, val.dtype)).minimum(UOp.const(1.0, val.dtype)) + val = _FUNCS['isNAN'](val).where(UOp.const(0.0, val.dtype), clamped) if val.dtype in (dtypes.uint64, dtypes.int64, dtypes.float64): lo, hi = _split64(val) raw_stores.extend([('vgpr', self.wvgpr_dyn(vdst_reg, lane, lo, exec_mask)), ('vgpr', self.wvgpr_dyn(vdst_reg + _c(1), lane, hi, exec_mask))]) elif val.dtype in (dtypes.half, dtypes.uint16, dtypes.int16): result, old_val = _val_to_u32(val), self.rvgpr_dyn(vdst_reg, lane) - hi_result = (old_val & UOp.const(dtypes.uint32, 0xFFFF)) | (result << UOp.const(dtypes.uint32, 16)) + hi_result = (old_val & UOp.const(0xFFFF, dtypes.uint32)) | (result << UOp.const(16, dtypes.uint32)) # GFX9/CDNA zeroes upper 16 bits on lo-half write; RDNA preserves them - lo_result = (result & UOp.const(dtypes.uint32, 0xFFFF)) if self.wave_size == 64 else \ - (old_val & UOp.const(dtypes.uint32, 0xFFFF0000)) | (result & UOp.const(dtypes.uint32, 0xFFFF)) + lo_result = (result & UOp.const(0xFFFF, dtypes.uint32)) if self.wave_size == 64 else \ + (old_val & UOp.const(0xFFFF0000, dtypes.uint32)) | (result & UOp.const(0xFFFF, dtypes.uint32)) result = opsel_dst_hi.where(hi_result, lo_result) if isinstance(opsel_dst_hi, UOp) else hi_result if opsel_dst_hi else lo_result raw_stores.append(('vgpr', self.wvgpr_dyn(vdst_reg, lane, result, exec_mask))) else: raw_stores.append(('vgpr', self.wvgpr_dyn(vdst_reg, lane, _val_to_u32(val), exec_mask))) @@ -767,8 +767,8 @@ class _Ctx: if slice_stores: result = self.rvgpr_dyn(vdst_reg, lane) for lo_bit, width, val_bits in slice_stores: - mask = UOp.const(dtypes.uint32, ((1 << width) - 1) << lo_bit) - result = (result & (mask ^ UOp.const(dtypes.uint32, 0xFFFFFFFF))) | (val_bits << UOp.const(dtypes.uint32, lo_bit)) + mask = UOp.const(((1 << width) - 1) << lo_bit, dtypes.uint32) + result = (result & (mask ^ UOp.const(0xFFFFFFFF, dtypes.uint32))) | (val_bits << UOp.const(lo_bit, dtypes.uint32)) lane_stores.append(self.wvgpr_dyn(vdst_reg, lane, result, exec_mask)) # VCC/EXEC mask writes must be computed BEFORE VGPR stores to avoid reading modified VGPRs. # When vdst overlaps with src operands (e.g. v_add_co_u32 v[0], vcc, s[8], v[0]), the carry @@ -790,8 +790,8 @@ class _Ctx: def _compile_sopp(inst: ir3.SOPP | ir4.SOPP, ctx: _Ctx) -> UOp: simm16 = ctx.inst_field_signed(type(inst).simm16).cast(dtypes.int16) if inst.op in (ir3.SOPPOp.S_ENDPGM, ir4.SOPPOp.S_ENDPGM, irc.SOPPOp.S_ENDPGM): - return UOp.sink(ctx.wsgpr_dyn(_c(PC_LO_IDX), UOp.const(dtypes.uint32, 0xFFFFFFFF)), - ctx.wsgpr_dyn(_c(PC_HI_IDX), UOp.const(dtypes.uint32, 0xFFFFFFFF))) + return UOp.sink(ctx.wsgpr_dyn(_c(PC_LO_IDX), UOp.const(0xFFFFFFFF, dtypes.uint32)), + ctx.wsgpr_dyn(_c(PC_HI_IDX), UOp.const(0xFFFFFFFF, dtypes.uint32))) # S_BARRIER: advance PC past the barrier instruction. The execution loop detects barriers before executing and handles synchronization. barrier_ops = {ir3.SOPPOp.S_BARRIER, irc.SOPPOp.S_BARRIER} if hasattr(ir4.SOPPOp, 'S_BARRIER_WAIT'): barrier_ops.add(ir4.SOPPOp.S_BARRIER_WAIT) @@ -804,8 +804,8 @@ def _compile_sopp(inst: ir3.SOPP | ir4.SOPP, ctx: _Ctx) -> UOp: pc_bytes = ctx.rpc() # PC is already 64-bit byte address vcc, exec_val = ctx.rmask(_c(VCC_LO.offset)), ctx.rexec() srcs: dict[str, UOp|int] = {'PC': pc_bytes.cast(dtypes.int64), 'SIMM16': simm16, 'SCC': ctx.rsgpr_dyn(_c(SCC.offset)), 'VCC': vcc, - 'VCCZ': vcc.eq(UOp.const(vcc.dtype, 0)).cast(dtypes.uint32), - 'EXECZ': exec_val.eq(UOp.const(exec_val.dtype, 0)).cast(dtypes.uint32)} + 'VCCZ': vcc.eq(UOp.const(0, vcc.dtype)).cast(dtypes.uint32), + 'EXECZ': exec_val.eq(UOp.const(0, exec_val.dtype)).cast(dtypes.uint32)} for dest, val in parse_pcode(pcode, srcs)[1]: if dest == 'PC' or dest.startswith('PC.'): lo, hi = _split64(val.cast(dtypes.uint64)) @@ -833,12 +833,12 @@ def _compile_smem(inst: ir3.SMEM | ir4.SMEM, ctx: _Ctx) -> UOp: part = op_name.rsplit('_', 1)[1] # B32, DWORD, DWORDX2, U8, I8, etc. nval = int(part.removeprefix('DWORD').removeprefix('X') or '1') if 'DWORD' in part else int(part[1:]) / 32 * (-1 if part[0] == 'I' else 1) ndwords = max(1, int(abs(nval))) - dword_base = addr >> UOp.const(dtypes.uint64, 2) - vals = [ctx.vmem.index(dword_base + UOp.const(dtypes.uint64, i)) for i in range(ndwords)] + dword_base = addr >> UOp.const(2, dtypes.uint64) + vals = [ctx.vmem.index(dword_base + UOp.const(i, dtypes.uint64)) for i in range(ndwords)] if abs(nval) < 1: nbits = int(abs(nval) * 32) - byte_off = (addr & UOp.const(dtypes.uint64, 3)).cast(dtypes.uint32) * UOp.const(dtypes.uint32, 8) - extracted = (vals[0] >> byte_off) & UOp.const(dtypes.uint32, (1 << nbits) - 1) + byte_off = (addr & UOp.const(3, dtypes.uint64)).cast(dtypes.uint32) * UOp.const(8, dtypes.uint32) + extracted = (vals[0] >> byte_off) & UOp.const((1 << nbits) - 1, dtypes.uint32) vals[0] = extracted.cast({8: dtypes.int8, 16: dtypes.int16}[nbits]).cast(dtypes.int32).bitcast(dtypes.uint32) if nval < 0 else extracted stores = [ctx.wsgpr_dyn(sdata_reg + _c(i), vals[i]) for i in range(ndwords)] return UOp.sink(*stores, *ctx.inc_pc()) @@ -941,7 +941,7 @@ def _dpp16_ctrl(lane: UOp, dpp: int, row_mask: int, bank_mask: int, wave_size: i enabled = (((_c(row_mask) >> row.cast(dtypes.uint32)) & _c(1)).ne(_c(0)) & (((_c(bank_mask) >> bank.cast(dtypes.uint32)) & _c(1)).ne(_c(0)))) op, arg = decode_dpp16(dpp) - src_lane, valid = lane_i, UOp.const(dtypes.bool, True) + src_lane, valid = lane_i, UOp.const(True) if op == 'quad_perm': assert isinstance(arg, tuple) @@ -967,7 +967,7 @@ def _load_dpp16_src0(ctx: _Ctx, inst, lane: UOp, fallback: UOp) -> UOp: getattr(inst, 'bank_mask', 0xf) or 0xf, ctx.wave_size) safe_src_lane = (enabled & valid).where(src_lane, _c(0, dtypes.int)) swizzled = ctx.rvgpr_dyn(ctx.inst_field(type(inst).vsrc0), safe_src_lane) - invalid = UOp.const(fallback.dtype, 0) if getattr(inst, 'bc', 0) else fallback + invalid = UOp.const(0, fallback.dtype) if getattr(inst, 'bc', 0) else fallback return enabled.where(valid.where(swizzled, invalid), fallback) def _compile_sdwa(inst: irc.VOP1_SDWA | irc.VOP2_SDWA | irc.VOP2_SDWA_SDST | irc.VOPC_SDWA_SDST, ctx: _Ctx) -> UOp: @@ -1159,7 +1159,7 @@ def _compile_vopc(inst: ir3.VOPC|ir3.VOPC_DPP16|ir3.VOP3|ir4.VOPC|ir4.VOPC_DPP16 s1 = _apply_src_mods(s1, 0, 1 if getattr(inst, 'src1_abs', 0) else 0, 1 if getattr(inst, 'src1_neg', 0) else 0, bits['s1']) s0 = _apply_src_mods(s0, 0, abs_bits, neg_bits, bits['s0']) s1 = _apply_src_mods(s1, 1, abs_bits, neg_bits, bits['s1']) - for dest, val in parse_pcode(pcode, {'S0': s0, 'S1': s1, 'laneId': lc, 'D0': UOp.const(dtypes.uint64, 0)})[1]: + for dest, val in parse_pcode(pcode, {'S0': s0, 'S1': s1, 'laneId': lc, 'D0': UOp.const(0, dtypes.uint64)})[1]: if '[laneId]' in dest and ('D0' in dest or 'EXEC' in dest): return val.cast(dtypes.uint32) return _c(0) @@ -1189,8 +1189,8 @@ def _compile_bitop3(inst, ctx: _Ctx, exec_mask: UOp, bits: dict, op_name: str) - is_16 = 'B16' in op_name dt, mask = (dtypes.uint16, 0xFFFF) if is_16 else (dtypes.uint32, 0xFFFFFFFF) s0, s1, s2 = src0.cast(dt), src1.cast(dt), src2.cast(dt) - def bnot(v): return v ^ UOp.const(dt, mask) - result = UOp.const(dt, 0) + def bnot(v): return v ^ UOp.const(mask, dt) + result = UOp.const(0, dt) for i in range(8): if not (ttbl & (1 << i)): continue result = result | ((s0 if i & 4 else bnot(s0)) & (s1 if i & 2 else bnot(s1)) & (s2 if i & 1 else bnot(s2))) @@ -1244,7 +1244,7 @@ def _compile_vop3(inst: ir3.VOP3 | ir4.VOP3 | irc.VOP3, ctx: _Ctx) -> UOp: src0 = _apply_src_mods(src0, 0, abs_bits, neg_bits, bits['s0']) src1 = _apply_src_mods(src1, 1, abs_bits, neg_bits, bits['s1']) src2 = _apply_src_mods(src2, 2, abs_bits, neg_bits, bits['s2']) - srcs = {'S0': src0, 'S1': src1, 'S2': src2, 'OPSEL': UOp.const(dtypes.uint32, opsel)} + srcs = {'S0': src0, 'S1': src1, 'S2': src2, 'OPSEL': UOp.const(opsel, dtypes.uint32)} if 'CNDMASK' in op_name and src2 is not None: srcs['VCC'] = src2 # FMAC instructions need D0 (accumulator) from destination register if 'FMAC' in op_name: srcs['D0'] = ctx.rvgpr_dyn(vdst_reg, lane) @@ -1398,7 +1398,7 @@ def _compile_mfma(inst: irc.VOP3P, ctx: _Ctx) -> UOp: # Layout: tmp[0..n_a_elems-1] = A[m][k], tmp[n_a_elems..n_a_elems+n_b_elems-1] = B[n][k] # Within each group of lanes, lane%grp_sub gives M/N index, lane//grp_sub gives sub-block grp_sub = min(M, 16) # lanes within group mapped to M/N dimension - b_off = UOp.const(dtypes.int, n_a_elems) + b_off = UOp.const(n_a_elems, dtypes.int) acc_dt = dtypes.int32 if is_int_out else dtypes.float32 # Use uint32 temp array to prevent optimizer from eliminating f16→f32 bitcast chains. # The optimizer folds bitcast(uint32→float32) stores to float32 arrays, losing the conversion. @@ -1407,44 +1407,44 @@ def _compile_mfma(inst: irc.VOP3P, ctx: _Ctx) -> UOp: def cvt_elem(raw: UOp, sub_idx: int) -> UOp: if is_i8: # Extract i8, sign-extend to i32 - byte_val = (raw >> UOp.const(dtypes.uint32, sub_idx * 8)) & UOp.const(dtypes.uint32, 0xFF) - return (byte_val.cast(dtypes.int32) ^ UOp.const(dtypes.int32, 0x80)) - UOp.const(dtypes.int32, 0x80) + byte_val = (raw >> UOp.const(sub_idx * 8, dtypes.uint32)) & UOp.const(0xFF, dtypes.uint32) + return (byte_val.cast(dtypes.int32) ^ UOp.const(0x80, dtypes.int32)) - UOp.const(0x80, dtypes.int32) elif is_f32_src: return raw # already uint32 (f32 bit pattern) elif is_fp8: - return ((raw >> UOp.const(dtypes.uint32, sub_idx * 8)) & UOp.const(dtypes.uint32, 0xFF)).cast(dtypes.uint32) + return ((raw >> UOp.const(sub_idx * 8, dtypes.uint32)) & UOp.const(0xFF, dtypes.uint32)).cast(dtypes.uint32) elif is_bf16: # bf16→f32 bits: just shift left by 16 (bf16 is upper 16 bits of f32) - return ((raw >> UOp.const(dtypes.uint32, sub_idx * 16)) & UOp.const(dtypes.uint32, 0xFFFF)) << UOp.const(dtypes.uint32, 16) + return ((raw >> UOp.const(sub_idx * 16, dtypes.uint32)) & UOp.const(0xFFFF, dtypes.uint32)) << UOp.const(16, dtypes.uint32) else: # f16→f32 conversion using float arithmetic to avoid UOp optimizer eliminating the conversion. # The optimizer folds bitcast(uint32→float32) chains, so we compute the float value directly. - h = (raw >> UOp.const(dtypes.uint32, sub_idx * 16)) & UOp.const(dtypes.uint32, 0xFFFF) - sign = (h >> UOp.const(dtypes.uint32, 15)) & UOp.const(dtypes.uint32, 1) - exp = (h >> UOp.const(dtypes.uint32, 10)) & UOp.const(dtypes.uint32, 0x1F) - mant = h & UOp.const(dtypes.uint32, 0x3FF) + h = (raw >> UOp.const(sub_idx * 16, dtypes.uint32)) & UOp.const(0xFFFF, dtypes.uint32) + sign = (h >> UOp.const(15, dtypes.uint32)) & UOp.const(1, dtypes.uint32) + exp = (h >> UOp.const(10, dtypes.uint32)) & UOp.const(0x1F, dtypes.uint32) + mant = h & UOp.const(0x3FF, dtypes.uint32) # Use bf16 path: shift left by 16 to create bf16 bits, then shift mantissa and adjust exponent in float domain # bf16 bits = (sign << 15) | (exp_bf16 << 7) | mant_bf16 -- but f16 and bf16 have different formats # Instead: construct f32 bits properly, use a local uint32 array to force materialization - f32_bits = (sign << UOp.const(dtypes.uint32, 31)) | \ - ((exp + UOp.const(dtypes.uint32, 112)) << UOp.const(dtypes.uint32, 23)) | \ - (mant << UOp.const(dtypes.uint32, 13)) - is_zero = exp.eq(UOp.const(dtypes.uint32, 0)) + f32_bits = (sign << UOp.const(31, dtypes.uint32)) | \ + ((exp + UOp.const(112, dtypes.uint32)) << UOp.const(23, dtypes.uint32)) | \ + (mant << UOp.const(13, dtypes.uint32)) + is_zero = exp.eq(UOp.const(0, dtypes.uint32)) # Return uint32 (f32 bit pattern) — stored directly to uint32 temp array, bitcast to float on read - return is_zero.where(UOp.const(dtypes.uint32, 0), f32_bits) + return is_zero.where(UOp.const(0, dtypes.uint32), f32_bits) read_lane = ctx.range() # For 32x32: lane%16 = M/N index within 16-wide block, lane//16 = which of 4 quarter-waves # Groups: lanes 0-31 = group 0, lanes 32-63 = group 1 # Within group: (lane%32)%16 = M/N[0-15], (lane%32)//16 selects M/N[0-15] or [16-31] - lane_in_grp = read_lane % UOp.const(dtypes.int, grp_size) - grp_idx = read_lane // UOp.const(dtypes.int, grp_size) + lane_in_grp = read_lane % UOp.const(grp_size, dtypes.int) + grp_idx = read_lane // UOp.const(grp_size, dtypes.int) if M == 32: # 32x32: lane_in_grp%16 = sub-row/col (0-15), lane_in_grp//16 = block (0=rows 0-15, 1=rows 16-31) - sub_mn = lane_in_grp % UOp.const(dtypes.int, 16) - block_mn = lane_in_grp // UOp.const(dtypes.int, 16) - mn_idx = block_mn * UOp.const(dtypes.int, 16) + sub_mn # actual M/N index (0-31) + sub_mn = lane_in_grp % UOp.const(16, dtypes.int) + block_mn = lane_in_grp // UOp.const(16, dtypes.int) + mn_idx = block_mn * UOp.const(16, dtypes.int) + sub_mn # actual M/N index (0-31) else: mn_idx = lane_in_grp # for 16x16 and 4x4 @@ -1456,18 +1456,18 @@ def _compile_mfma(inst: irc.VOP3P, ctx: _Ctx) -> UOp: ctx.rsrc_dyn(src0_off, _c(0, dtypes.int), 32)) a_val = cvt_elem(a_raw, sub_idx) if M == 4: - a_idx = grp_idx * UOp.const(dtypes.int, M * K) + mn_idx * UOp.const(dtypes.int, K) + UOp.const(dtypes.int, kl) + a_idx = grp_idx * UOp.const(M * K, dtypes.int) + mn_idx * UOp.const(K, dtypes.int) + UOp.const(kl, dtypes.int) else: - a_idx = mn_idx * UOp.const(dtypes.int, K) + grp_idx * UOp.const(dtypes.int, k_per_grp) + UOp.const(dtypes.int, kl) + a_idx = mn_idx * UOp.const(K, dtypes.int) + grp_idx * UOp.const(k_per_grp, dtypes.int) + UOp.const(kl, dtypes.int) read_stores.append(tmp.index(a_idx).store(a_val)) b_raw = src1_is_vgpr.where(ctx.rvgpr_dyn(src1_r + _c(reg_idx), read_lane), ctx.rsrc_dyn(src1_off, _c(0, dtypes.int), 32)) b_val = cvt_elem(b_raw, sub_idx) if M == 4: - b_idx = b_off + grp_idx * UOp.const(dtypes.int, N * K) + mn_idx * UOp.const(dtypes.int, K) + UOp.const(dtypes.int, kl) + b_idx = b_off + grp_idx * UOp.const(N * K, dtypes.int) + mn_idx * UOp.const(K, dtypes.int) + UOp.const(kl, dtypes.int) else: - b_idx = b_off + mn_idx * UOp.const(dtypes.int, K) + grp_idx * UOp.const(dtypes.int, k_per_grp) + UOp.const(dtypes.int, kl) + b_idx = b_off + mn_idx * UOp.const(K, dtypes.int) + grp_idx * UOp.const(k_per_grp, dtypes.int) + UOp.const(kl, dtypes.int) read_stores.append(tmp.index(b_idx).store(b_val)) read_phase = UOp.group(*read_stores).end(read_lane) @@ -1488,11 +1488,11 @@ def _compile_mfma(inst: irc.VOP3P, ctx: _Ctx) -> UOp: # Lane mapping: n = (lane%32)%16 + ((lane%32)//16)*16, gives column 0-31 # Row groups: 4 groups of 4, covering rows 0-31. Group g (0-3): rows g*4 .. g*4+3 # group assignment: lane//16 gives quarter (0-3), each quarter maps to 4 rows - c_lane_in_32 = compute_lane % UOp.const(dtypes.int, 32) - c_sub = c_lane_in_32 % UOp.const(dtypes.int, 16) - c_block = c_lane_in_32 // UOp.const(dtypes.int, 16) - n_idx = c_block * UOp.const(dtypes.int, 16) + c_sub - c_half = compute_lane // UOp.const(dtypes.int, 32) # 0 or 1 + c_lane_in_32 = compute_lane % UOp.const(32, dtypes.int) + c_sub = c_lane_in_32 % UOp.const(16, dtypes.int) + c_block = c_lane_in_32 // UOp.const(16, dtypes.int) + n_idx = c_block * UOp.const(16, dtypes.int) + c_sub + c_half = compute_lane // UOp.const(32, dtypes.int) # 0 or 1 for out_reg in range(16): # Each half covers 8 rows. out_reg 0-3: rows 0-3 (half0) or 16-19 (half1) @@ -1503,7 +1503,7 @@ def _compile_mfma(inst: irc.VOP3P, ctx: _Ctx) -> UOp: # acc[r] at lane l -> D[m][n] where n = (l%32)%16 + ((l%32)//16)*16 # m = (l//32)*16 + (r//4)*4 + (r%4) ... giving rows in blocks of 4 # So: m_base = half * 16 + (out_reg // 4) * 4 + (out_reg % 4) - m_base = c_half * UOp.const(dtypes.int, 16) + UOp.const(dtypes.int, (out_reg // 4) * 4 + (out_reg % 4)) + m_base = c_half * UOp.const(16, dtypes.int) + UOp.const((out_reg // 4) * 4 + (out_reg % 4), dtypes.int) acc_v = (ctx.raccvgpr_dyn if use_acc else ctx.rvgpr_dyn)(src2_r + _c(out_reg), compute_lane, src2_is_vgpr) if is_int_out: acc_v = acc_v.cast(dtypes.int32) @@ -1511,8 +1511,8 @@ def _compile_mfma(inst: irc.VOP3P, ctx: _Ctx) -> UOp: acc = src2_is_vgpr.where(acc_v, acc_scalar) for k in range(K): - a_val = tmp2.index(m_base * UOp.const(dtypes.int, K) + UOp.const(dtypes.int, k)).bitcast(acc_dt) - b_val = tmp2.index(b_off + n_idx * UOp.const(dtypes.int, K) + UOp.const(dtypes.int, k)).bitcast(acc_dt) + a_val = tmp2.index(m_base * UOp.const(K, dtypes.int) + UOp.const(k, dtypes.int)).bitcast(acc_dt) + b_val = tmp2.index(b_off + n_idx * UOp.const(K, dtypes.int) + UOp.const(k, dtypes.int)).bitcast(acc_dt) acc = acc + a_val * b_val if is_int_out: @@ -1523,8 +1523,8 @@ def _compile_mfma(inst: irc.VOP3P, ctx: _Ctx) -> UOp: vdst_reg + _c(out_reg), compute_lane, acc.bitcast(dtypes.uint32), exec_mask)) else: # 16x16 and 4x4: each lane computes out_per_lane outputs - n_idx = compute_lane % UOp.const(dtypes.int, grp_sub) - c_grp = compute_lane // UOp.const(dtypes.int, grp_sub) + n_idx = compute_lane % UOp.const(grp_sub, dtypes.int) + c_grp = compute_lane // UOp.const(grp_sub, dtypes.int) for out_reg in range(out_per_lane): acc_v = (ctx.raccvgpr_dyn if use_acc else ctx.rvgpr_dyn)(src2_r + _c(out_reg), compute_lane, src2_is_vgpr) @@ -1534,17 +1534,17 @@ def _compile_mfma(inst: irc.VOP3P, ctx: _Ctx) -> UOp: if M == 4: # 4x4: each group is independent. A/B indexed per-group. - m_base = c_grp * UOp.const(dtypes.int, M * K) + UOp.const(dtypes.int, out_reg * K) + m_base = c_grp * UOp.const(M * K, dtypes.int) + UOp.const(out_reg * K, dtypes.int) for k in range(K): - a_val = tmp2.index(m_base + UOp.const(dtypes.int, k)).bitcast(acc_dt) - b_val = tmp2.index(b_off + c_grp * UOp.const(dtypes.int, N*K) + n_idx * UOp.const(dtypes.int, K)+UOp.const(dtypes.int, k)).bitcast(acc_dt) + a_val = tmp2.index(m_base + UOp.const(k, dtypes.int)).bitcast(acc_dt) + b_val = tmp2.index(b_off + c_grp * UOp.const(N*K, dtypes.int) + n_idx * UOp.const(K, dtypes.int)+UOp.const(k, dtypes.int)).bitcast(acc_dt) acc = acc + a_val * b_val else: # 16x16: K is split across groups. Shared MxK/NxK arrays. - m_base = c_grp * UOp.const(dtypes.int, out_per_lane) + UOp.const(dtypes.int, out_reg) + m_base = c_grp * UOp.const(out_per_lane, dtypes.int) + UOp.const(out_reg, dtypes.int) for k in range(K): - a_val = tmp2.index(m_base * UOp.const(dtypes.int, K) + UOp.const(dtypes.int, k)).bitcast(acc_dt) - b_val = tmp2.index(b_off + n_idx * UOp.const(dtypes.int, K) + UOp.const(dtypes.int, k)).bitcast(acc_dt) + a_val = tmp2.index(m_base * UOp.const(K, dtypes.int) + UOp.const(k, dtypes.int)).bitcast(acc_dt) + b_val = tmp2.index(b_off + n_idx * UOp.const(K, dtypes.int) + UOp.const(k, dtypes.int)).bitcast(acc_dt) acc = acc + a_val * b_val if is_int_out: @@ -1570,8 +1570,8 @@ def _compile_wmma(inst: ir3.VOP3P | ir4.VOP3P | irc.VOP3P, ctx: _Ctx) -> UOp: is_rdna4 = isinstance(inst, ir4.VOP3P) # read 16x16 F16/BF16 matrix from VGPRs → flat f32 array[row*16+k] def read_f16_val(src, lane, vgpr, half): - v = ctx.rvgpr_dyn(src + _c(vgpr), UOp.const(dtypes.int, lane)) - return cvt((v >> UOp.const(dtypes.uint32, 16)) if half else (v & UOp.const(dtypes.uint32, 0xFFFF))) + v = ctx.rvgpr_dyn(src + _c(vgpr), UOp.const(lane, dtypes.int)) + return cvt((v >> UOp.const(16, dtypes.uint32)) if half else (v & UOp.const(0xFFFF, dtypes.uint32))) # RDNA3: 16 lanes × 8 VGPRs × 2 halves, k maps linearly # RDNA4: 32 lanes × 4 VGPRs × 2 halves, k bits are scrambled (k[2] goes to lane bit 4) @@ -1593,20 +1593,20 @@ def _compile_wmma(inst: ir3.VOP3P | ir4.VOP3P | irc.VOP3P, ctx: _Ctx) -> UOp: for m in range(16) for n in range(16) for lane, vgpr in [d_map(m, n)]] mat_d = [sum(mat_a[r*16+k] * mat_b[c*16+k] for k in range(16)) + mat_c[r*16+c] for r in range(16) for c in range(16)] def f32_to_f16_bits(v: UOp) -> UOp: return v.cast(dtypes.half).bitcast(dtypes.uint16).cast(dtypes.uint32) - def f32_to_bf16_bits(v: UOp) -> UOp: return (v.bitcast(dtypes.uint32) >> UOp.const(dtypes.uint32, 16)) & UOp.const(dtypes.uint32, 0xFFFF) + def f32_to_bf16_bits(v: UOp) -> UOp: return (v.bitcast(dtypes.uint32) >> UOp.const(16, dtypes.uint32)) & UOp.const(0xFFFF, dtypes.uint32) out_cvt = f32_to_bf16_bits if is_bf16 else f32_to_f16_bits if is_rdna4: # pack 2 f16 per VGPR: adjacent m values share (lane, vgpr) since vgpr=m&7, half=m&1 - stores = [ctx.wvgpr_dyn(vdst_reg + _c(d_map(m, n)[1] // 2), UOp.const(dtypes.int, d_map(m, n)[0]), - out_cvt(mat_d[m*16+n]) | (out_cvt(mat_d[(m+1)*16+n]) << UOp.const(dtypes.uint32, 16)), exec_mask) + stores = [ctx.wvgpr_dyn(vdst_reg + _c(d_map(m, n)[1] // 2), UOp.const(d_map(m, n)[0], dtypes.int), + out_cvt(mat_d[m*16+n]) | (out_cvt(mat_d[(m+1)*16+n]) << UOp.const(16, dtypes.uint32)), exec_mask) for n in range(16) for m in range(0, 16, 2)] else: # (rdna3) 1 f16 per VGPR (lo half only) - stores = [ctx.wvgpr_dyn(vdst_reg + _c(d_map(m, n)[1]), UOp.const(dtypes.int, d_map(m, n)[0]), out_cvt(mat_d[m*16+n]), exec_mask) + stores = [ctx.wvgpr_dyn(vdst_reg + _c(d_map(m, n)[1]), UOp.const(d_map(m, n)[0], dtypes.int), out_cvt(mat_d[m*16+n]), exec_mask) for m in range(16) for n in range(16)] else: # f32 - mat_c = [ctx.rvgpr_dyn(src2_r + _c(d_map(m, n)[1]), UOp.const(dtypes.int, d_map(m, n)[0])).bitcast(dtypes.float32) + mat_c = [ctx.rvgpr_dyn(src2_r + _c(d_map(m, n)[1]), UOp.const(d_map(m, n)[0], dtypes.int)).bitcast(dtypes.float32) for m in range(16) for n in range(16)] mat_d = [sum(mat_a[r*16+k] * mat_b[c*16+k] for k in range(16)) + mat_c[r*16+c] for r in range(16) for c in range(16)] - stores = [ctx.wvgpr_dyn(vdst_reg + _c(d_map(m, n)[1]), UOp.const(dtypes.int, d_map(m, n)[0]), mat_d[m*16+n].bitcast(dtypes.uint32), exec_mask) + stores = [ctx.wvgpr_dyn(vdst_reg + _c(d_map(m, n)[1]), UOp.const(d_map(m, n)[0], dtypes.int), mat_d[m*16+n].bitcast(dtypes.uint32), exec_mask) for m in range(16) for n in range(16)] return UOp.sink(*stores, *ctx.inc_pc()) @@ -1688,8 +1688,8 @@ def _compile_vop3p(inst: ir3.VOP3P | ir4.VOP3P | irc.VOP3P, ctx: _Ctx) -> UOp: scalar_hi_sel = src_lo if not opsel_hi_bit else is_sgpr_pair.where(sgpr_hi, src_lo) lo = is_vgpr.where(vgpr_hi if opsel_lo else vgpr_lo, scalar_lo_sel) hi = is_vgpr.where(vgpr_hi if opsel_hi_bit else vgpr_lo, scalar_hi_sel) - if neg_lo: lo = lo ^ UOp.const(dtypes.uint32, 0x80000000) - if neg_hi_bit: hi = hi ^ UOp.const(dtypes.uint32, 0x80000000) + if neg_lo: lo = lo ^ UOp.const(0x80000000, dtypes.uint32) + if neg_hi_bit: hi = hi ^ UOp.const(0x80000000, dtypes.uint32) return _u64(lo, hi) srcs = {'S0': build_pk_f32(src0, src_offs[0], opsel & 1, opsel_hi & 1, neg & 1, neg_hi & 1), 'S1': build_pk_f32(src1, src_offs[1], opsel & 2, opsel_hi & 2, neg & 2, neg_hi & 2), @@ -1700,35 +1700,35 @@ def _compile_vop3p(inst: ir3.VOP3P | ir4.VOP3P | irc.VOP3P, ctx: _Ctx) -> UOp: def apply_abs(v, bit, opsel_hi_bit, opsel_bit): if not (neg_hi & bit): return v # Apply abs based on whether source is f32 or f16 - if not (combined_opsel_hi & opsel_hi_bit): return v & UOp.const(dtypes.uint32, 0x7FFFFFFF) # f32 abs - if opsel & opsel_bit: return v & UOp.const(dtypes.uint32, 0x7FFF0000) # f16 hi abs (preserve lo) - return v & UOp.const(dtypes.uint32, 0xFFFF7FFF) # f16 lo abs (preserve hi) + if not (combined_opsel_hi & opsel_hi_bit): return v & UOp.const(0x7FFFFFFF, dtypes.uint32) # f32 abs + if opsel & opsel_bit: return v & UOp.const(0x7FFF0000, dtypes.uint32) # f16 hi abs (preserve lo) + return v & UOp.const(0xFFFF7FFF, dtypes.uint32) # f16 lo abs (preserve hi) def apply_neg_mix(v, bit, opsel_hi_bit, opsel_bit): if not (neg & bit): return v - if not (combined_opsel_hi & opsel_hi_bit): return v ^ UOp.const(dtypes.uint32, 0x80000000) # f32 neg - if opsel & opsel_bit: return v ^ UOp.const(dtypes.uint32, 0x80000000) # f16 hi neg - return v ^ UOp.const(dtypes.uint32, 0x00008000) # f16 lo neg + if not (combined_opsel_hi & opsel_hi_bit): return v ^ UOp.const(0x80000000, dtypes.uint32) # f32 neg + if opsel & opsel_bit: return v ^ UOp.const(0x80000000, dtypes.uint32) # f16 hi neg + return v ^ UOp.const(0x00008000, dtypes.uint32) # f16 lo neg s0_mod = apply_neg_mix(apply_abs(src0, 1, 1, 1), 1, 1, 1) s1_mod = apply_neg_mix(apply_abs(src1, 2, 2, 2), 2, 2, 2) s2_mod = apply_neg_mix(apply_abs(src2, 4, 4, 4), 4, 4, 4) srcs = {'S@0': s0_mod, 'S@1': s1_mod, 'S@2': s2_mod, - 'OPSEL_HI': UOp.const(dtypes.uint32, combined_opsel_hi), 'OPSEL': UOp.const(dtypes.uint32, opsel)} + 'OPSEL_HI': UOp.const(combined_opsel_hi, dtypes.uint32), 'OPSEL': UOp.const(opsel, dtypes.uint32)} else: def get_half_bits(val: UOp, use_hi: bool, apply_neg: bool = False) -> UOp: - bits = ((val >> UOp.const(dtypes.uint32, 16)) if use_hi else val) & UOp.const(dtypes.uint32, 0xFFFF) + bits = ((val >> UOp.const(16, dtypes.uint32)) if use_hi else val) & UOp.const(0xFFFF, dtypes.uint32) if apply_neg: bits = bits.cast(dtypes.uint16).bitcast(dtypes.half).neg().bitcast(dtypes.uint16).cast(dtypes.uint32) return bits def build_remapped_src(src: UOp, opsel_lo_bit: int, opsel_hi_bit: int, neg_lo_bit: int, neg_hi_bit: int) -> UOp: lo = get_half_bits(src, bool(opsel_lo_bit), bool(neg_lo_bit)) hi = get_half_bits(src, bool(opsel_hi_bit), bool(neg_hi_bit)) - return lo | (hi << UOp.const(dtypes.uint32, 16)) + return lo | (hi << UOp.const(16, dtypes.uint32)) # DOT IU instructions use NEG bits for signed/unsigned selection, not fp16 negation is_dot_iu = 'DOT' in op_name and 'IU' in op_name n0, n1, n2, nh0, nh1, nh2 = (0, 0, 0, 0, 0, 0) if is_dot_iu else (neg & 1, neg & 2, neg & 4, neg_hi & 1, neg_hi & 2, neg_hi & 4) srcs = {'S0': build_remapped_src(src0, opsel & 1, opsel_hi & 1, n0, nh0), 'S1': build_remapped_src(src1, opsel & 2, opsel_hi & 2, n1, nh1), 'S2': build_remapped_src(src2, opsel & 4, 1 if opsel_hi2 else 0, n2, nh2)} - if is_dot_iu: srcs['NEG'] = UOp.const(dtypes.uint32, neg) + if is_dot_iu: srcs['NEG'] = UOp.const(neg, dtypes.uint32) return ctx.compile_vop_pcode(inst.op, srcs, lane, vdst_reg, exec_mask) def _compile_vopd(inst: ir3.VOPD | ir4.VOPD, ctx: _Ctx) -> UOp: @@ -1781,7 +1781,7 @@ def _compile_mem_op(inst: ir3.DS|ir3.FLAT|ir3.GLOBAL|ir3.SCRATCH|ir4.DS|ir4.VFLA # CDNA acc bit: when set, VGPR operands (vdst/vdata) target ACCVGPR file instead of VGPR use_acc = bool(getattr(inst, 'acc', 0)) mem = ctx.lds if is_lds else ctx.scratch if is_scratch else ctx.vmem - addr_shift = UOp.const(dtypes.uint32 if is_lds else dtypes.uint64, 2) + addr_shift = UOp.const(2, dtypes.uint32 if is_lds else dtypes.uint64) # Extract register info - all dynamic for deduplication if is_lds: @@ -1831,20 +1831,20 @@ def _compile_mem_op(inst: ir3.DS|ir3.FLAT|ir3.GLOBAL|ir3.SCRATCH|ir4.DS|ir4.VFLA return addr offset64 = offset.cast(dtypes.uint64) # Dynamic saddr check: saddr < 124 means valid SGPR, otherwise use VGPR pair for address - use_saddr = (saddr_reg < _c(124)) if saddr_reg is not None else UOp.const(dtypes.bool, False) + use_saddr = (saddr_reg < _c(124)) if saddr_reg is not None else UOp.const(False) if is_scratch: scratch_stride = ctx.rsgpr_dyn(_c(SCRATCH_STRIDE_IDX)).cast(dtypes.uint64) base = lane.cast(dtypes.uint64) * scratch_stride # SVE (Scratch VGPR Enable): when SVE=1, VADDR is used as offset; when SVE=0, VADDR is ignored sve = getattr(inst, 'sve', 0) vaddr = ctx.rvgpr_dyn(addr_reg, lane).cast(dtypes.uint64) - addr_offset = vaddr if sve == 1 else UOp.const(dtypes.uint64, 0) + addr_offset = vaddr if sve == 1 else UOp.const(0, dtypes.uint64) # Add saddr value only if use_saddr is true (saddr < 124) - saddr_contrib = use_saddr.where(ctx.rsgpr_dyn(saddr_reg).cast(dtypes.uint64), UOp.const(dtypes.uint64, 0)) \ - if saddr_reg is not None else UOp.const(dtypes.uint64, 0) + saddr_contrib = use_saddr.where(ctx.rsgpr_dyn(saddr_reg).cast(dtypes.uint64), UOp.const(0, dtypes.uint64)) \ + if saddr_reg is not None else UOp.const(0, dtypes.uint64) return base + addr_offset + saddr_contrib + offset64 # FLAT/GLOBAL: choose between SGPR base (saddr) or VGPR pair (addr) based on saddr validity - saddr_base = _u64(ctx.rsgpr_dyn(saddr_reg), ctx.rsgpr_dyn(saddr_reg + _c(1))) if saddr_reg is not None else UOp.const(dtypes.uint64, 0) + saddr_base = _u64(ctx.rsgpr_dyn(saddr_reg), ctx.rsgpr_dyn(saddr_reg + _c(1))) if saddr_reg is not None else UOp.const(0, dtypes.uint64) vaddr_base = _u64(ctx.rvgpr_dyn(addr_reg, lane), ctx.rvgpr_dyn(addr_reg + _c(1), lane)) # When saddr is valid: base = saddr pair, vaddr is 32-bit offset; otherwise: base = 0, vaddr is 64-bit address base_addr = use_saddr.where(saddr_base + ctx.rvgpr_dyn(addr_reg, lane).cast(dtypes.uint64), vaddr_base) @@ -1874,18 +1874,18 @@ def _compile_mem_op(inst: ir3.DS|ir3.FLAT|ir3.GLOBAL|ir3.SCRATCH|ir4.DS|ir4.VFLA data = {'DATA': ctx.rvgpr_dyn(vdata_reg, lane), 'DATA1': ctx.rvgpr_dyn(vdata_reg + _c(1), lane), 'DATA2': ctx.rvgpr_dyn(vdata_reg + _c(2), lane)} elif data_bits_mem <= 32: - data = {'DATA': ctx.rvgpr_dyn(vdata_reg, lane), 'DATA2': ctx.rvgpr_dyn(data1_reg, lane) if has_data1 else UOp.const(dtypes.uint32, 0)} + data = {'DATA': ctx.rvgpr_dyn(vdata_reg, lane), 'DATA2': ctx.rvgpr_dyn(data1_reg, lane) if has_data1 else UOp.const(0, dtypes.uint32)} else: data = {'DATA': _u64(ctx.rvgpr_dyn(vdata_reg, lane), ctx.rvgpr_dyn(vdata_reg + _c(1), lane)), - 'DATA2': _u64(ctx.rvgpr_dyn(data1_reg, lane), ctx.rvgpr_dyn(data1_reg + _c(1), lane)) if has_data1 else UOp.const(dtypes.uint64, 0)} + 'DATA2': _u64(ctx.rvgpr_dyn(data1_reg, lane), ctx.rvgpr_dyn(data1_reg + _c(1), lane)) if has_data1 else UOp.const(0, dtypes.uint64)} # RDNA3 uses ADDR/OFFSET, RDNA4 uses vgpr_a/offset (lowercase) + CalcDsAddr function return {'ADDR': addr, 'ADDR_BASE': addr, 'OFFSET': offset, 'OFFSET0': offset0, 'OFFSET1': offset1, '_lds': mem, 'laneId': lane, 'vgpr_a': ctx.rvgpr_dyn(addr_reg, lane), 'offset': offset, 'offset0': offset0, 'offset1': offset1, **data} active = _lane_active(exec_mask, lane) # saddr < 124 means valid SGPR pair, otherwise use 0 (NULL means no saddr contribution) - use_saddr = (saddr_reg < _c(124)) if saddr_reg is not None else UOp.const(dtypes.bool, False) - saddr_raw = _u64(ctx.rsgpr_dyn(saddr_reg), ctx.rsgpr_dyn(saddr_reg + _c(1))) if saddr_reg is not None else UOp.const(dtypes.uint64, 0) - saddr_base = use_saddr.where(saddr_raw, UOp.const(dtypes.uint64, 0)) + use_saddr = (saddr_reg < _c(124)) if saddr_reg is not None else UOp.const(False) + saddr_raw = _u64(ctx.rsgpr_dyn(saddr_reg), ctx.rsgpr_dyn(saddr_reg + _c(1))) if saddr_reg is not None else UOp.const(0, dtypes.uint64) + saddr_base = use_saddr.where(saddr_raw, UOp.const(0, dtypes.uint64)) # Sign-extend offset to 64-bit for the final address calculation ioffset64 = offset.cast(dtypes.int64).cast(dtypes.uint64) # v_addr for CalcGlobalAddr: when saddr valid, use low 32 bits as offset; otherwise full 64-bit address. Include ioffset. @@ -1900,13 +1900,13 @@ def _compile_mem_op(inst: ir3.DS|ir3.FLAT|ir3.GLOBAL|ir3.SCRATCH|ir4.DS|ir4.VFLA # acc bit: read/write ACCVGPR instead of VGPR for data operands _rvdata = (lambda r, l, *a: ctx.raccvgpr_dyn(r, l)) if use_acc else ctx.rvgpr_dyn vdata = _rvdata(vdata_reg, lane).cast(dtypes.uint64) if 'STORE' in op_name \ - else _rvdata(vdst_reg, lane) if 'D16' in op_name else UOp.const(dtypes.uint32, 0) + else _rvdata(vdst_reg, lane) if 'D16' in op_name else UOp.const(0, dtypes.uint32) if 'STORE' in op_name and data_bits_mem >= 64: - vdata = vdata | (_rvdata(vdata_reg + _c(1), lane).cast(dtypes.uint64) << UOp.const(dtypes.uint64, 32)) + vdata = vdata | (_rvdata(vdata_reg + _c(1), lane).cast(dtypes.uint64) << UOp.const(32, dtypes.uint64)) srcs = {'ADDR': addr, 'VDATA': vdata, '_vmem': mem, '_active': active, 'laneId': lane, 'v_addr': vaddr_base, 's_saddr': saddr_base, 'SADDR': saddr_base, 'OFFSET': offset} for i in range(data_bits_mem // 32): - srcs[f'VDATA{i}'] = _rvdata(vdata_reg + _c(i), lane) if 'STORE' in op_name else UOp.const(dtypes.uint32, 0) + srcs[f'VDATA{i}'] = _rvdata(vdata_reg + _c(i), lane) if 'STORE' in op_name else UOp.const(0, dtypes.uint32) return srcs def make_stores(dest: str, val: UOp, lane: UOp, active: UOp, writes_return_data: bool) -> list[UOp]: @@ -1983,7 +1983,7 @@ def _compile_mubuf(inst: irc.MUBUF, ctx: _Ctx) -> UOp: offset, offen, idxen = ctx.inst_field(type(inst).offset), ctx.inst_field(type(inst).offen), ctx.inst_field(type(inst).idxen) # V# descriptor: base[0:1], num_records[2], stride=word3[13:0] - base = _u64(ctx.rsgpr_dyn(srsrc), ctx.rsgpr_dyn(srsrc + _c(1))) & UOp.const(dtypes.uint64, 0xFFFFFFFFFFFF) + base = _u64(ctx.rsgpr_dyn(srsrc), ctx.rsgpr_dyn(srsrc + _c(1))) & UOp.const(0xFFFFFFFFFFFF, dtypes.uint64) num_records = ctx.rsgpr_dyn(srsrc + _c(2)) stride = (ctx.rsgpr_dyn(srsrc + _c(3)) & _c(0x3FFF)).cast(dtypes.uint64) @@ -2000,7 +2000,7 @@ def _compile_mubuf(inst: irc.MUBUF, ctx: _Ctx) -> UOp: buffer_offset = (stride * index + voff + offset.cast(dtypes.uint64)).cast(dtypes.uint32) in_bounds = active & buffer_offset.__lt__(num_records) addr = base + soff + buffer_offset.cast(dtypes.uint64) - addr = in_bounds.where(addr, UOp.const(dtypes.uint64, 0)) # safe address when OOB + addr = in_bounds.where(addr, UOp.const(0, dtypes.uint64)) # safe address when OOB mem = ctx.vmem stores: list[UOp] = [] @@ -2009,20 +2009,20 @@ def _compile_mubuf(inst: irc.MUBUF, ctx: _Ctx) -> UOp: lds_base = ctx.rsgpr_dyn(_c(124)) & _c(0x3FFFF) lds_addr = lds_base + lane.cast(dtypes.uint32) * _c(n_dwords * 4) for i in range(n_dwords): - word_addr = (addr + UOp.const(dtypes.uint64, i * 4)) >> UOp.const(dtypes.uint64, 2) + word_addr = (addr + UOp.const(i * 4, dtypes.uint64)) >> UOp.const(2, dtypes.uint64) val = in_bounds.where(mem.index(word_addr.cast(dtypes.int64)).load(), _c(0)) lds_idx = (lds_addr + _c(i * 4)) >> _c(2) lds_slot = ctx.lds.index(lds_idx.valid(active)) stores.append(lds_slot.store(active.where(val, lds_slot))) elif is_store: for i in range(n_dwords): - word_addr = (addr + UOp.const(dtypes.uint64, i * 4)) >> UOp.const(dtypes.uint64, 2) + word_addr = (addr + UOp.const(i * 4, dtypes.uint64)) >> UOp.const(2, dtypes.uint64) idx = mem.index(word_addr.cast(dtypes.int64).valid(in_bounds)) val = (ctx.raccvgpr_dyn if use_acc else ctx.rvgpr_dyn)(vdata + _c(i), lane) stores.append(idx.store(in_bounds.where(_to_u32(val), idx))) else: for i in range(n_dwords): - word_addr = (addr + UOp.const(dtypes.uint64, i * 4)) >> UOp.const(dtypes.uint64, 2) + word_addr = (addr + UOp.const(i * 4, dtypes.uint64)) >> UOp.const(2, dtypes.uint64) val = in_bounds.where(mem.index(word_addr.cast(dtypes.int64).valid(in_bounds)).load(), _c(0)) stores.append((ctx.waccvgpr_dyn if use_acc else ctx.wvgpr_dyn)(vdata + _c(i), lane, val, exec_mask)) return UOp.sink(UOp.group(*stores).end(lane), *ctx.inc_pc()) diff --git a/test/mockgpu/amd/pcode.py b/test/mockgpu/amd/pcode.py index b53dd1e97c..f7fbd8587b 100644 --- a/test/mockgpu/amd/pcode.py +++ b/test/mockgpu/amd/pcode.py @@ -7,7 +7,7 @@ from tinygrad.codegen.decomp.dtype import f2f # Type alias for vars dict: stores UOps and tuples for lambda definitions VarVal = UOp | tuple[str, list[str], str] -def _const(dt, v): return UOp.const(dt, v) +def _const(dt, v): return UOp.const(v, dt) def _u32(v): return _const(dtypes.uint32, v) def _u64(v): return _const(dtypes.uint64, v) def _to_u32(v): return v if v.dtype == dtypes.uint32 else v.bitcast(dtypes.uint32) if v.dtype.itemsize == 4 else v.cast(dtypes.uint32) @@ -866,8 +866,8 @@ class Parser: idx_hi_native = ((addr + _const(adt, 4)) >> _const(adt, 2)).cast(dtypes.int64) safe_idx_hi = is_unaligned.where(idx_hi_native, idx_native) hi = mindex(safe_idx_hi) - combined = val.cast(dtypes.uint64) | (hi.cast(dtypes.uint64) << UOp.const(dtypes.uint64, 32)) - val = is_unaligned.where((combined >> (byte_off.cast(dtypes.uint64) * UOp.const(dtypes.uint64, 8))).cast(dtypes.uint32), val) + combined = val.cast(dtypes.uint64) | (hi.cast(dtypes.uint64) << UOp.const(32, dtypes.uint64)) + val = is_unaligned.where((combined >> (byte_off.cast(dtypes.uint64) * UOp.const(8, dtypes.uint64))).cast(dtypes.uint32), val) return _cast_to(val, dt) def _coerce_cmp(self, l: UOp, r: UOp) -> tuple[UOp, UOp]: @@ -1339,7 +1339,7 @@ def parse_block(lines: list[str], start: int, env: dict[str, VarVal], funcs: dic return (dest, (val[0], cnd.where(val[1], val[1]))) return (dest, val) # Build combined condition: each branch fires when its cond is true AND no earlier cond was true - remaining = UOp.const(dtypes.bool, True) + remaining = UOp.const(True) for bc, bse in branch_assigns: effective = remaining & bc if remaining.op != Ops.CONST else bc for dest, val in bse: assigns.append(_cond_side_effect(effective, dest, val)) diff --git a/test/null/test_const_folding.py b/test/null/test_const_folding.py index c15b38fdae..e2a7bbfcd2 100644 --- a/test/null/test_const_folding.py +++ b/test/null/test_const_folding.py @@ -36,21 +36,21 @@ class TestUnaryOpsConstFolding(unittest.TestCase): class TestWeakConstFolding(unittest.TestCase): def test_weakint_math(self): - out = (UOp.const(dtypes.weakint, 2**40) + UOp.const(dtypes.weakint, 2**40)).simplify() + out = (UOp.const(2**40) + UOp.const(2**40)).simplify() self.assertEqual((out.op, out.dtype, out.arg), (Ops.CONST, dtypes.weakint, 2**41)) def test_float_unaries(self): for dtype in (dtypes.weakfloat,): for op in (Ops.SIN, Ops.LOG2, Ops.EXP2, Ops.SQRT, Ops.RECIPROCAL): - out = UOp.const(dtype, 4).alu(op).simplify() + out = UOp.const(4, dtype).alu(op).simplify() self.assertEqual((out.op, out.dtype), (Ops.CONST, dtypes.weakfloat)) def test_weakfloat_math(self): - out = (UOp.const(dtypes.weakfloat, 1.25) + UOp.const(dtypes.weakfloat, 2.5)).simplify() + out = (UOp.const(1.25) + UOp.const(2.5)).simplify() self.assertEqual((out.op, out.dtype, out.arg), (Ops.CONST, dtypes.weakfloat, 3.75)) def test_invalid_poison(self): - self.assertIs(UOp.invalid().alu(Ops.CDIV, UOp.const(dtypes.weakint, 0)).simplify().arg, Invalid) + self.assertIs(UOp.invalid().alu(Ops.CDIV, UOp.const(0)).simplify().arg, Invalid) class TestBinaryOpsConstFolding(unittest.TestCase): def test_add_literal_zero(self): @@ -121,7 +121,7 @@ class TestBitcastConstFolding(unittest.TestCase): def t(cases: dict[DType, ConstType]): for (from_dt, from_v), (to_dt, to_v) in itertools.product(cases.items(), cases.items()): if not math.isnan(from_v): - r = full_rewrite(UOp.const(from_dt, from_v).bitcast(to_dt).sink()).src[0] + r = full_rewrite(UOp.const(from_v, from_dt).bitcast(to_dt).sink()).src[0] self.assertEqual(r.op, Ops.CONST, msg:=f"{from_dt} -> {to_dt} ({from_v} -> {to_v})") self.assertEqual(r.dtype, to_dt, msg) np.testing.assert_equal(r.arg, to_v, msg) @@ -145,7 +145,7 @@ class TestBitcastConstFolding(unittest.TestCase): def test_vec_bitcast(self): with Context(SPEC=0): - srcs = full_rewrite(UOp.const(dtypes.int32, (-1, -2**31, 75)).bitcast(dtypes.uint32).sink()).src + srcs = full_rewrite(UOp.const((-1, -2**31, 75), dtypes.int32).bitcast(dtypes.uint32).sink()).src self.assertTrue(all(r.op is Ops.CONST and r.dtype == dtypes.uint32 for r in srcs)) self.assertEqual(tuple(x.arg for x in srcs), (2**32-1, 2**31, 75)) diff --git a/test/null/test_gpudims.py b/test/null/test_gpudims.py index 7a93d272a6..dbca42b954 100644 --- a/test/null/test_gpudims.py +++ b/test/null/test_gpudims.py @@ -24,7 +24,7 @@ class TestGroupedDims(unittest.TestCase): total = math.prod(dims) specials = sorted(dedup(flatten([[y for y in x.toposort() if y.op is Ops.SPECIAL] for x in idxs])), key=lambda u: u.arg) # build flat index and primed flat (same expression with renamed SPECIALs) - flat = UOp.const(dtypes.weakint, 0) + flat = UOp.const(0) for i, idx in enumerate(idxs): flat = flat + idx * int(math.prod(dims[i+1:])) flat_p = flat.substitute({s: UOp(Ops.SPECIAL, src=s.src, arg=s.arg+"_p") for s in specials}) @@ -107,7 +107,7 @@ class TestGroupedDims(unittest.TestCase): def test_global_prod_max(self): g, l = UOp.range(256, 0, AxisType.GLOBAL), UOp.range(256, 1, AxisType.LOCAL) - sink = UOp.param(0, dtypes.float, (512,)).index(g + l).store(UOp.const(dtypes.float, 1.0)).end(g, l).sink(arg=KernelInfo()) + sink = UOp.param(0, dtypes.float, (512,)).index(g + l).store(UOp.const(1.0, dtypes.float)).end(g, l).sink(arg=KernelInfo()) class R(Renderer): global_max, local_max, global_prod_max = (256, 256, 256), (128, 128, 128), (128, 128, 128) specials = [u for u in add_gpudims(R(Target()), sink).toposort() if u.op is Ops.SPECIAL] self.assertGreater(len([s for s in specials if "lidx" in s.arg]), 1) diff --git a/test/null/test_gradient.py b/test/null/test_gradient.py index 00282f0cdf..ad39dfa7c9 100644 --- a/test/null/test_gradient.py +++ b/test/null/test_gradient.py @@ -14,7 +14,7 @@ class TestGradient(unittest.TestCase): def _test_one_input_function(self, f:Callable, jf:Callable|None=None): if jf is None: jf = f x = UOp.variable('x', -math.inf, math.inf, dtype=dtypes.float) - gx = compute_gradient(f(x), UOp.const(dtypes.float, 1.0), set([x]))[x] + gx = compute_gradient(f(x), UOp.const(1.0, dtypes.float), set([x]))[x] for val in [-5., -2.0, 0.0, 2.0, 5.]: tg_out = gx.substitute({x: x.const_like(val)}).ssimplify() @@ -26,7 +26,7 @@ class TestGradient(unittest.TestCase): if jf is None: jf = f x = UOp.variable('x', -math.inf, math.inf, dtype=dtypes.float) y = UOp.variable('y', -math.inf, math.inf, dtype=dtypes.float) - grads = compute_gradient(f(x, y), UOp.const(dtypes.float, 1.0), set([x, y])) + grads = compute_gradient(f(x, y), UOp.const(1.0, dtypes.float), set([x, y])) gx, gy = grads[x], grads[y] for valx in [-5., -2.0, 0.0, 2.0, 5.]: diff --git a/test/null/test_graph_rewrite.py b/test/null/test_graph_rewrite.py index cbb0895092..56d19f5965 100644 --- a/test/null/test_graph_rewrite.py +++ b/test/null/test_graph_rewrite.py @@ -32,30 +32,30 @@ def evaluate_uop(uop, variables): class TestArithmeticSimplifications(unittest.TestCase): def test_full_graph_rewrite_division_by_zero(self): - optimized_div_uop = apply_rewrite(UOp.const(dtypes.float32, 10.0) / UOp.const(dtypes.float32, 0.0)) + optimized_div_uop = apply_rewrite(UOp.const(10.0, dtypes.float32) / UOp.const(0.0, dtypes.float32)) self.assertEqual(optimized_div_uop.op, Ops.CONST) self.assertTrue(math.isinf(optimized_div_uop.arg) or math.isnan(optimized_div_uop.arg)) def test_full_graph_rewrite_redundant_operations(self): - optimized_uop = apply_rewrite((UOp.const(dtypes.float32, 10.0) + UOp.const(dtypes.float32, 0.0)) * UOp.const(dtypes.float32, 1.0)) + optimized_uop = apply_rewrite((UOp.const(10.0, dtypes.float32) + UOp.const(0.0, dtypes.float32)) * UOp.const(1.0, dtypes.float32)) self.assertEqual(optimized_uop.op, Ops.CONST) self.assertEqual(optimized_uop.arg, 10.0) def test_full_graph_rewrite_large_graph(self): - prev_uop = UOp.const(dtypes.int32, 0) + prev_uop = UOp.const(0, dtypes.int32) for i in range(1, 101): - prev_uop += UOp.const(dtypes.int32, i) + prev_uop += UOp.const(i, dtypes.int32) optimized_uop = apply_rewrite(prev_uop) self.assertEqual(optimized_uop.op, Ops.CONST) self.assertEqual(optimized_uop.arg, sum(range(1, 101))) def test_full_graph_rewrite_division_by_one(self): - optimized_uop = apply_rewrite(UOp.const(dtypes.float32, 42.0) / UOp.const(dtypes.float32, 1.0)) + optimized_uop = apply_rewrite(UOp.const(42.0, dtypes.float32) / UOp.const(1.0, dtypes.float32)) self.assertEqual(optimized_uop.op, Ops.CONST) self.assertEqual(optimized_uop.arg, 42.0) def test_full_graph_rewrite_modulo_by_one(self): - optimized_uop = apply_rewrite(UOp.const(dtypes.int32, 42) % UOp.const(dtypes.int32, 1)) + optimized_uop = apply_rewrite(UOp.const(42, dtypes.int32) % UOp.const(1, dtypes.int32)) self.assertEqual(optimized_uop.op, Ops.CONST) self.assertEqual(optimized_uop.arg, 0) @@ -63,17 +63,17 @@ class TestArithmeticSimplifications(unittest.TestCase): class TestFoldingAndReduction(unittest.TestCase): @unittest.skip("reduce is removed now") def test_full_graph_rewrite_constant_reduction_folding(self): - const1 = UOp.const(dtypes.int32, 5) - const2 = UOp.const(dtypes.int32, 10) - const3 = UOp.const(dtypes.int32, 20) + const1 = UOp.const(5, dtypes.int32) + const2 = UOp.const(10, dtypes.int32) + const3 = UOp.const(20, dtypes.int32) optimized_sink = apply_rewrite((const1 + const2 + const3).reduce(Ops.ADD)) expected_sum = 5 + 10 + 20 self.assertEqual(optimized_sink.arg, expected_sum) @unittest.skip("reduce is removed now") def test_full_graph_rewrite_reduction_with_unused_range(self): - const1 = UOp.const(dtypes.int32, 15) - const2 = UOp.const(dtypes.int32, 25) + const1 = UOp.const(15, dtypes.int32) + const2 = UOp.const(25, dtypes.int32) rng = UOp.range(10, idx=0) optimized_sink = apply_rewrite((const1 + const2).reduce(Ops.ADD, rng)) expected_sum = 10 * (15 + 25) @@ -89,7 +89,7 @@ class TestFoldingAndReduction(unittest.TestCase): @unittest.skip("currently failing") def test_full_graph_rewrite_simple_reduction_folding(self): simple_range = UOp.range(4, idx=0) - add_uop = simple_range + UOp.const(dtypes.int32, 1) + add_uop = simple_range + UOp.const(1, dtypes.int32) optimized_sink = apply_rewrite(add_uop.reduce(Ops.ADD, simple_range)) expected_sum = sum(i + 1 for i in range(4)) self.assertEqual(optimized_sink.arg, expected_sum) @@ -128,9 +128,9 @@ class TestModuloAndDivisionFolding(unittest.TestCase): def test_graph_rewrite_div_folding_bug(self): lhs = UOp(Ops.ADD, src=( - UOp(Ops.STACK, arg=None, src=(UOp(Ops.SPECIAL, src=(UOp.const(dtypes.int, 32),), arg='lidx0'),)*4), - UOp.const(dtypes.int, (0, 256, 512, 768)))) - rhs = UOp.const(dtypes.int, (2,)*4) + UOp(Ops.STACK, arg=None, src=(UOp(Ops.SPECIAL, src=(UOp.const(32, dtypes.int),), arg='lidx0'),)*4), + UOp.const((0, 256, 512, 768), dtypes.int))) + rhs = UOp.const((2,)*4, dtypes.int) unopt = lhs 0, @@ -182,27 +182,27 @@ class TestEdgeCasesAndSpecialOperations(unittest.TestCase): class TestGEPAndVectorizeRewrite(unittest.TestCase): def test_gep_single_element_extraction(self): # GEP on a vector dtype to extract a single element - base_vector = UOp.const(dtypes.float32, (1.0, 2.0, 3.0, 4.0)) + base_vector = UOp.const((1.0, 2.0, 3.0, 4.0), dtypes.float32) self.assertEqual(apply_rewrite(base_vector.index(2)).arg, 3.0) def test_gep_tuple_extraction(self): # GEP on a vector dtype to extract multiple elements as a vector - base_vector = UOp.const(dtypes.float32, (1.0, 2.0, 3.0, 4.0)) + base_vector = UOp.const((1.0, 2.0, 3.0, 4.0), dtypes.float32) self.assertEqual(list(apply_rewrite_values(UOp.stack(*[base_vector.index(i) for i in (2, 3)]))), [3.0, 4.0]) def test_gep_on_const_stack(self): # GEP on a const STACK to extract a single element - const_stack = UOp.const(dtypes.float32, (1.0, 2.0, 3.0, 4.0)) + const_stack = UOp.const((1.0, 2.0, 3.0, 4.0), dtypes.float32) self.assertEqual(apply_rewrite(const_stack.index(2)).arg, 3.0) def test_gep_tuple_on_const_stack(self): # GEP on a const STACK using a tuple to extract multiple elements - const_stack = UOp.const(dtypes.float32, (7.0, 8.0, 9.0, 10.0)) + const_stack = UOp.const((7.0, 8.0, 9.0, 10.0), dtypes.float32) self.assertEqual(list(apply_rewrite_values(UOp.stack(*[const_stack.index(i) for i in (1, 3)]))), [8.0, 10.0]) def test_vectorize_multiple_elements(self): # Vectorizing multiple elements using GEP - base_vector = UOp.const(dtypes.float32, (5.0, 10.0, 15.0, 20.0)) + base_vector = UOp.const((5.0, 10.0, 15.0, 20.0), dtypes.float32) vectorized_uop = UOp(Ops.STACK, src=tuple(base_vector.index(i) for i in range(4))) self.assertEqual(list(apply_rewrite_values(vectorized_uop)), [5.0, 10.0, 15.0, 20.0]) @@ -213,7 +213,7 @@ from tinygrad.uop.symbolic import symbolic_simple class TestBottomUpRewrite(unittest.TestCase): def test_const_folding(self): - a = UOp.const(dtypes.int, 5) + a = UOp.const(5, dtypes.int) ret = (a*3) + (a*7) gt = graph_rewrite(ret, symbolic_simple) ret = graph_rewrite(ret, symbolic_simple, bottom_up=True) @@ -305,7 +305,7 @@ class TestRecurse(unittest.TestCase): graph_rewrite(a, pm, bottom_up=True) def test_inf_loop(self): - a = UOp.const(dtypes.int, 3) + a = UOp.const(3, dtypes.int) pm = PatternMatcher([ (UPat(Ops.CONST, arg=3, name="x"), lambda x: x.replace(arg=4)), (UPat(Ops.CONST, arg=4, name="x"), lambda x: x.replace(arg=3)), @@ -314,7 +314,7 @@ class TestRecurse(unittest.TestCase): graph_rewrite(a, pm) def test_inf_loop_bottom_up(self): - a = UOp.const(dtypes.int, 3) + a = UOp.const(3, dtypes.int) pm = PatternMatcher([ (UPat(Ops.CONST, arg=3, name="x"), lambda x: x.replace(arg=4)), (UPat(Ops.CONST, arg=4, name="x"), lambda x: x.replace(arg=3)), @@ -325,8 +325,8 @@ class TestRecurse(unittest.TestCase): def bidir_append(ctx, x, b): ctx.append((x.arg if x.op is Ops.CONST else "+", b)) class TestBidirectional(unittest.TestCase): def test_simple(self): - a = UOp.const(dtypes.int, 1) - b = UOp.const(dtypes.int, 2) + a = UOp.const(1, dtypes.int) + b = UOp.const(2, dtypes.int) c = a + b pm = PatternMatcher([ (UPat(GroupOp.All, name="x"), lambda ctx,x: bidir_append(ctx, x, False)) ]) bpm = PatternMatcher([ (UPat(GroupOp.All, name="x"), lambda ctx,x: bidir_append(ctx, x, True)) ]) @@ -336,11 +336,11 @@ class TestBidirectional(unittest.TestCase): class TestStopEarly(unittest.TestCase): def test_stop_early(self): - a = UOp.const(dtypes.int, 3) - b = UOp.const(dtypes.int, 4) + a = UOp.const(3, dtypes.int) + b = UOp.const(4, dtypes.int) c = a+b - cn = UOp.const(dtypes.int, 7) - d = UOp.const(dtypes.int, 2) + cn = UOp.const(7, dtypes.int) + d = UOp.const(2, dtypes.int) def visit_const(c:UOp): print(f"visit {c.arg}") assert c.arg not in (3,4) @@ -376,7 +376,7 @@ class TestWalkRewrite(unittest.TestCase): def test_walk_topdown_no_fixed_point(self): """A bouncing pattern applies once and stops instead of looping.""" - a = UOp.const(dtypes.int, 3) + a = UOp.const(3, dtypes.int) pm = PatternMatcher([ (UPat(Ops.CONST, arg=3, name="x"), lambda x: x.replace(arg=4)), (UPat(Ops.CONST, arg=4, name="x"), lambda x: x.replace(arg=3)), @@ -384,7 +384,7 @@ class TestWalkRewrite(unittest.TestCase): with self.assertRaises(RuntimeError): graph_rewrite(a, pm, bottom_up=True) ret = graph_rewrite(a, pm, walk=True) - self.assertIs(ret, UOp.const(dtypes.int, 4)) + self.assertIs(ret, UOp.const(4, dtypes.int)) def test_walk_topdown_rewrites_children(self): a = UOp.variable('a', 0, 10) @@ -421,8 +421,8 @@ class TestWalkRewrite(unittest.TestCase): ctx.append(x.arg if x.op is Ops.CONST else x.op) return None pm = PatternMatcher([(UPat(GroupOp.All, name="x"), track_visit)]) - a = UOp.const(dtypes.int, 1) - b = UOp.const(dtypes.int, 2) + a = UOp.const(1, dtypes.int) + b = UOp.const(2, dtypes.int) graph_rewrite(a + b, pm, ctx=visited, walk=True) self.assertEqual(visited, [1, 2, Ops.ADD]) @@ -454,13 +454,13 @@ class TestWalkRewrite(unittest.TestCase): def test_walk_bottomup_no_fixed_point(self): """Bottom-up walk also applies once per node, no fixed-point iteration.""" - a = UOp.const(dtypes.int, 3) + a = UOp.const(3, dtypes.int) pm = PatternMatcher([ (UPat(Ops.CONST, arg=3, name="x"), lambda x: x.replace(arg=4)), (UPat(Ops.CONST, arg=4, name="x"), lambda x: x.replace(arg=3)), ]) ret = graph_rewrite(a, pm, bottom_up=True, walk=True) - self.assertIs(ret, UOp.const(dtypes.int, 4)) + self.assertIs(ret, UOp.const(4, dtypes.int)) def test_walk_bottomup_visit_order(self): """Bottom-up walk fires bpm before descending (pre-order).""" @@ -469,8 +469,8 @@ class TestWalkRewrite(unittest.TestCase): ctx.append(x.arg if x.op is Ops.CONST else x.op) return None pm = PatternMatcher([(UPat(GroupOp.All, name="x"), track_visit)]) - a = UOp.const(dtypes.int, 1) - b = UOp.const(dtypes.int, 2) + a = UOp.const(1, dtypes.int) + b = UOp.const(2, dtypes.int) graph_rewrite(a + b, pm, ctx=visited, bottom_up=True, walk=True) # bpm fires on each node before children: +, 1, 2 self.assertEqual(visited, [Ops.ADD, 1, 2]) @@ -497,8 +497,8 @@ class TestWalkRewrite(unittest.TestCase): return None bpm = PatternMatcher([(UPat(GroupOp.All, name="x"), bpm_visit)]) pm = PatternMatcher([(UPat(GroupOp.All, name="x"), pm_visit)]) - a = UOp.const(dtypes.int, 1) - b = UOp.const(dtypes.int, 2) + a = UOp.const(1, dtypes.int) + b = UOp.const(2, dtypes.int) graph_rewrite(a + b, pm, ctx=visited, bpm=bpm, walk=True) # bpm fires pre-order, pm fires post-order self.assertEqual(visited, [ @@ -518,14 +518,14 @@ class TestWalkRewrite(unittest.TestCase): return None bpm = PatternMatcher([(UPat(GroupOp.All, name="x"), bpm_match)]) pm = PatternMatcher([(UPat(GroupOp.All, name="x"), pm_match)]) - a = UOp.const(dtypes.int, 1) - b = UOp.const(dtypes.int, 2) + a = UOp.const(1, dtypes.int) + b = UOp.const(2, dtypes.int) ret = graph_rewrite(a + b, pm, ctx=visited, bpm=bpm, walk=True) # bpm matches const(1) and short-circuits it, so pm never fires on const(1) self.assertNotIn((1, "pm"), visited) # but pm still fires on const(2) and the rebuilt ADD self.assertIn((2, "pm"), visited) - self.assertIs(ret, UOp.const(dtypes.int, 10) + b) + self.assertIs(ret, UOp.const(10, dtypes.int) + b) if __name__ == '__main__': unittest.main() diff --git a/test/null/test_helpers.py b/test/null/test_helpers.py index ca65513c04..de4c04c6f6 100644 --- a/test/null/test_helpers.py +++ b/test/null/test_helpers.py @@ -297,10 +297,10 @@ class TestPolyN(unittest.TestCase): from tinygrad.dtype import dtypes from tinygrad.uop.ops import UOp from test.helpers import eval_uop - np.testing.assert_allclose(eval_uop(polyN(UOp.const(dtypes.float, 1.0), [1.0, -2.0, 1.0])), 0.0) - np.testing.assert_allclose(eval_uop(polyN(UOp.const(dtypes.float, 2.0), [1.0, -2.0, 1.0])), 1.0) - np.testing.assert_allclose(eval_uop(polyN(UOp.const(dtypes.float, 3.0), [1.0, -2.0, 1.0])), 4.0) - np.testing.assert_allclose(eval_uop(polyN(UOp.const(dtypes.float, 4.0), [1.0, -2.0, 1.0])), 9.0) + np.testing.assert_allclose(eval_uop(polyN(UOp.const(1.0, dtypes.float), [1.0, -2.0, 1.0])), 0.0) + np.testing.assert_allclose(eval_uop(polyN(UOp.const(2.0, dtypes.float), [1.0, -2.0, 1.0])), 1.0) + np.testing.assert_allclose(eval_uop(polyN(UOp.const(3.0, dtypes.float), [1.0, -2.0, 1.0])), 4.0) + np.testing.assert_allclose(eval_uop(polyN(UOp.const(4.0, dtypes.float), [1.0, -2.0, 1.0])), 9.0) class TestTimeToStr(unittest.TestCase): def test_seconds(self): self.assertEqual(" 10.01s ", time_to_str(10.01)) diff --git a/test/null/test_linearizer_failures.py b/test/null/test_linearizer_failures.py index fe23315741..a5df41f53f 100644 --- a/test/null/test_linearizer_failures.py +++ b/test/null/test_linearizer_failures.py @@ -8,15 +8,15 @@ from tinygrad.codegen import to_program class TestLinearizerFailures(unittest.TestCase): def test_fail_1(self): c0 = UOp.param(0, dtypes.float, (64,)) - c1 = UOp.range(UOp.const(dtypes.weakint, 2), 1, AxisType.WEAK) - c2 = UOp.range(UOp.const(dtypes.weakint, 32), 2, AxisType.WEAK) - c3 = ((c1*UOp.const(dtypes.weakint, 32))+c2) + c1 = UOp.range(UOp.const(2), 1, AxisType.WEAK) + c2 = UOp.range(UOp.const(32), 2, AxisType.WEAK) + c3 = ((c1*UOp.const(32))+c2) c4 = UOp.param(1, dtypes.float, (163840,)) - c5 = UOp.range(UOp.const(dtypes.weakint, 2560), 0, AxisType.REDUCE) - c6 = c4.index(((((((c5//UOp.const(dtypes.weakint, 8))%UOp.const(dtypes.weakint, 8))*UOp.const(dtypes.weakint, 8))+(c5%UOp.const(dtypes.weakint, 8)))+(((c2*UOp.const(dtypes.weakint, 40))+(c5//UOp.const(dtypes.weakint, 64)))*UOp.const(dtypes.weakint, 64)))+(c1*UOp.const(dtypes.weakint, 81920)))) + c5 = UOp.range(UOp.const(2560), 0, AxisType.REDUCE) + c6 = c4.index(((((((c5//UOp.const(8))%UOp.const(8))*UOp.const(8))+(c5%UOp.const(8)))+(((c2*UOp.const(40))+(c5//UOp.const(64)))*UOp.const(64)))+(c1*UOp.const(81920)))) c7 = UOp.param(2, dtypes.float, (64,)) c8 = c7.index(c3) - c9 = ((((c6+(c8*UOp.const(dtypes.float, -1.0)))*(c6+(c8*UOp.const(dtypes.float, -1.0)))).reduce(c5, arg=Ops.ADD)*UOp.const(dtypes.float, 0.000390625))+UOp.const(dtypes.float, 1e-05)).sqrt().reciprocal() + c9 = ((((c6+(c8*UOp.const(-1.0, dtypes.float)))*(c6+(c8*UOp.const(-1.0, dtypes.float)))).reduce(c5, arg=Ops.ADD)*UOp.const(0.000390625, dtypes.float))+UOp.const(1e-05, dtypes.float)).sqrt().reciprocal() c10 = c0.index(c3).store(c9).end(c1, c2) ast = c10.sink(arg=KernelInfo()) to_program(ast, renderer=Device[Device.DEFAULT].renderer) diff --git a/test/null/test_microbenchmarks.py b/test/null/test_microbenchmarks.py index 2efb64d211..c09fa997a4 100644 --- a/test/null/test_microbenchmarks.py +++ b/test/null/test_microbenchmarks.py @@ -27,29 +27,29 @@ class TestBench(unittest.TestCase): print(f"{self._testMethodName:30s} {et*1e6/self.N:.2f} us") def test_uop_instant_creation(self): - for i in range(self.N): UOp.const(dtypes.int, 100+i) + for i in range(self.N): UOp.const(100+i, dtypes.int) def test_uop_list_creation(self): - [UOp.const(dtypes.int, 100+i) for i in range(self.N)] + [UOp.const(100+i, dtypes.int) for i in range(self.N)] def test_uop_add_2n(self): - a = UOp.const(dtypes.int, 2) + a = UOp.const(2, dtypes.int) for _ in range(self.N): a = a + a def test_uop_toposort(self): - a = UOp.const(dtypes.int, 0) - for i in range(self.N): a = a + UOp.const(dtypes.int, 100+i) + a = UOp.const(0, dtypes.int) + for i in range(self.N): a = a + UOp.const(100+i, dtypes.int) self.start_time() self.assertEqual(len(a.toposort()), 2*self.N+1) def test_uop_toposort_2n(self): - a = UOp.const(dtypes.int, 0) + a = UOp.const(0, dtypes.int) for _ in range(self.N): a = a + a self.start_time() self.assertEqual(len(a.toposort()), self.N+1) def test_uop_simplify(self): - a = UOp.const(dtypes.int, 2) + a = UOp.const(2, dtypes.int) for _ in range(self.N): (a+a).simplify() def test_uop_simplify_complex(self): @@ -68,7 +68,7 @@ class TestBench(unittest.TestCase): for _ in range(self.N): expr.simplify() def test_uop_chain_free(self): - a = UOp.const(dtypes.int, 2) + a = UOp.const(2, dtypes.int) for _ in range(self.N): a = a + a self.start_time() del a diff --git a/test/null/test_pattern_matcher.py b/test/null/test_pattern_matcher.py index ff839cdd10..3ed8deb8e5 100644 --- a/test/null/test_pattern_matcher.py +++ b/test/null/test_pattern_matcher.py @@ -6,8 +6,8 @@ from tinygrad.uop.ops import PatternMatcher, UPat class TestPatternMatcher(unittest.TestCase): def test_simple_match(self): matcher = PatternMatcher([(UPat(Ops.CONST, name="x", dtype=dtypes.float), lambda x: x.rtag())]) - c1 = UOp.const(dtypes.float, 1.0) - c2 = UOp.const(dtypes.int, 1) + c1 = UOp.const(1.0, dtypes.float) + c2 = UOp.const(1, dtypes.int) self.assertEqual(matcher.rewrite(c1), c1.rtag()) self.assertEqual(matcher.rewrite(c2), None) @@ -61,16 +61,16 @@ class TestPatternMatcher(unittest.TestCase): def test_uop(self): matcher = PatternMatcher([(UPat(Ops.CONST, name="x"), lambda x: x.rtag())]) - c1 = UOp.const(dtypes.float, 1.0) + c1 = UOp.const(1.0, dtypes.float) c2 = UOp(Ops.ADD, src=(c1, c1)) self.assertEqual(matcher.rewrite(c1), c1.rtag()) self.assertEqual(matcher.rewrite(c2), None) def test_uop_set(self): matcher = PatternMatcher([(UPat((Ops.CONST, Ops.CAST), name="x"), lambda x: x.rtag())]) - c1 = UOp.const(dtypes.bool, False) + c1 = UOp.const(False) c2 = UOp(Ops.CAST, arg=dtypes.int, src=(c1,)) - c3 = UOp.const(dtypes.float, 1.0) + c3 = UOp.const(1.0, dtypes.float) c4 = UOp(Ops.ADD, src=(c3, c3)) self.assertEqual(matcher.rewrite(c1), c1.rtag()) self.assertEqual(matcher.rewrite(c2), c2.rtag()) @@ -82,11 +82,11 @@ class TestPatternMatcher(unittest.TestCase): (UPat(Ops.CONST, arg=False, name="x"), lambda x: x.rtag()), (UPat(Ops.MAX, name="x"), lambda x: x.rtag()), ]) - c1 = UOp.const(dtypes.float, 0.0) - c2 = UOp.const(dtypes.bool, False) + c1 = UOp.const(0.0, dtypes.float) + c2 = UOp.const(False) c3 = UOp(Ops.MAX, src=(c1, c1)) c4 = UOp(Ops.MUL, src=(c1, c1)) - c5 = UOp.const(dtypes.int, -1) + c5 = UOp.const(-1, dtypes.int) self.assertEqual(matcher.rewrite(c1), c1.rtag()) self.assertEqual(matcher.rewrite(c2), c2.rtag()) self.assertEqual(matcher.rewrite(c3), c3.rtag()) @@ -98,9 +98,9 @@ class TestPatternMatcher(unittest.TestCase): (UPat(Ops.MUL, src=[UPat(Ops.CONST, name="c"), UPat(Ops.CONST, arg=2)], name="x"), lambda x,c: x.rtag() if c.arg in {1, -1} else None) ]) - y1 = UOp.const(dtypes.int, 1) - y2 = UOp.const(dtypes.int, 2) - y3 = UOp.const(dtypes.int, -1) + y1 = UOp.const(1, dtypes.int) + y2 = UOp.const(2, dtypes.int) + y3 = UOp.const(-1, dtypes.int) c1 = UOp(Ops.MUL, src=(y1, y2)) c2 = UOp(Ops.MUL, src=(y2, y2)) c3 = UOp(Ops.MUL, src=(y3, y2)) @@ -114,8 +114,8 @@ class TestPatternMatcher(unittest.TestCase): def test_dup_name(self): matcher = PatternMatcher([(UPat(GroupOp.ALU, name="x", src=(UPat(Ops.CONST, name="y"), UPat(Ops.CONST, name="y"))), lambda x, y: x.rtag())]) - y1 = UOp.const(dtypes.float, 1.0) - y2 = UOp.const(dtypes.float, 1.0) + y1 = UOp.const(1.0, dtypes.float) + y2 = UOp.const(1.0, dtypes.float) c1 = UOp(Ops.ADD, src=(y1, y1)) c2 = UOp(Ops.ADD, src=(y1, y2)) self.assertEqual(matcher.rewrite(c1), c1.rtag()) @@ -123,17 +123,17 @@ class TestPatternMatcher(unittest.TestCase): def test_dtype(self): matcher = PatternMatcher([(UPat(Ops.CONST, name="x", dtype=dtypes.float32), lambda x: x.rtag())]) - c1 = UOp.const(dtypes.float, 1.0) - c2 = UOp.const(dtypes.float64, 1.0) + c1 = UOp.const(1.0, dtypes.float) + c2 = UOp.const(1.0, dtypes.float64) self.assertEqual(matcher.rewrite(c1), c1.rtag()) self.assertEqual(matcher.rewrite(c2), None) def test_dtype_set(self): matcher = PatternMatcher([(UPat(Ops.CONST, name="x", dtype={dtypes.float32, dtypes.float64}), lambda x: x.rtag())]) - c1 = UOp.const(dtypes.float, 1.0) - c2 = UOp.const(dtypes.float64, 1.0) - c3 = UOp.const(dtypes.float16, 1.0) - c4 = UOp.const(dtypes.int, 1) + c1 = UOp.const(1.0, dtypes.float) + c2 = UOp.const(1.0, dtypes.float64) + c3 = UOp.const(1.0, dtypes.float16) + c4 = UOp.const(1, dtypes.int) self.assertEqual(matcher.rewrite(c1), c1.rtag()) self.assertEqual(matcher.rewrite(c2), c2.rtag()) self.assertEqual(matcher.rewrite(c3), None) @@ -141,8 +141,8 @@ class TestPatternMatcher(unittest.TestCase): def test_src_one(self): matcher = PatternMatcher([(UPat(GroupOp.ALU, name="x", src=(UPat(Ops.CONST), UPat(Ops.CONST))), lambda x: x.rtag())]) - c1 = UOp.const(dtypes.float, 1.0) - c2 = UOp.const(dtypes.float, 2.0) + c1 = UOp.const(1.0, dtypes.float) + c2 = UOp.const(2.0, dtypes.float) c3 = UOp(Ops.ADD, src=(c1,c2)) self.assertEqual(matcher.rewrite(c3), c3.rtag()) self.assertEqual(matcher.rewrite(c2), None) @@ -158,8 +158,8 @@ class TestPatternMatcher(unittest.TestCase): def test_src_permutations(self): matcher = PatternMatcher([(UPat(GroupOp.ALU, name="x", src=[UPat(Ops.CONST), UPat(GroupOp.ALU)]), lambda x: x.rtag())]) - c1 = UOp.const(dtypes.float, 1.0) - c2 = UOp.const(dtypes.float, 2.0) + c1 = UOp.const(1.0, dtypes.float) + c2 = UOp.const(2.0, dtypes.float) c3 = UOp(Ops.ADD, src=(c1,c2)) c4 = UOp(Ops.ADD, src=(c3,c2)) c5 = UOp(Ops.ADD, src=(c2,c3)) @@ -171,8 +171,8 @@ class TestPatternMatcher(unittest.TestCase): def test_src_repeat(self): matcher = PatternMatcher([(UPat(GroupOp.ALU, name="x", src=UPat(Ops.CONST)), lambda x: x.rtag())]) - c1 = UOp.const(dtypes.float, 1.0) - c2 = UOp.const(dtypes.float, 2.0) + c1 = UOp.const(1.0, dtypes.float) + c2 = UOp.const(2.0, dtypes.float) c3 = UOp(Ops.ADD, src=(c1,c2)) c4 = UOp(Ops.ADD, src=(c2,c3)) self.assertEqual(matcher.rewrite(c3), c3.rtag()) @@ -180,9 +180,9 @@ class TestPatternMatcher(unittest.TestCase): def test_allow_len(self): matcher = PatternMatcher([(UPat(Ops.MULACC, name="x", src=(UPat(Ops.CONST),), allow_any_len=True), lambda x: x.rtag())]) - c1 = UOp.const(dtypes.float, 1.0) - c2 = UOp.const(dtypes.float, 2.0) - c3 = UOp.const(dtypes.float, 3.0) + c1 = UOp.const(1.0, dtypes.float) + c2 = UOp.const(2.0, dtypes.float) + c3 = UOp.const(3.0, dtypes.float) c4 = UOp(Ops.EXP2, src=(c1,)) c5 = UOp(Ops.ADD, src=(c1,c2)) c6 = UOp(Ops.MULACC, src=(c1,c2,c3)) @@ -191,8 +191,8 @@ class TestPatternMatcher(unittest.TestCase): self.assertEqual(matcher.rewrite(c6), c6.rtag()) def test_deep_src_permutations(self): - c1 = UOp.const(dtypes.float, 1.0) - c2 = UOp.const(dtypes.float, 2.0) + c1 = UOp.const(1.0, dtypes.float) + c2 = UOp.const(2.0, dtypes.float) u1 = (c1 + c2) + c1 u2 = (c2 + c1) + c1 matcher = PatternMatcher([ diff --git a/test/null/test_simplify_valid_idx.py b/test/null/test_simplify_valid_idx.py index 436dee3d96..6760e4402d 100644 --- a/test/null/test_simplify_valid_idx.py +++ b/test/null/test_simplify_valid_idx.py @@ -23,7 +23,7 @@ def get_load_image_uop(image_shape:tuple[int, ...], valid:UOp, idx:tuple[UOp, UO UOp.param(0, dtypes.float, image_shape).index(idx[1].valid(valid), idx[0].valid(valid)), )) -def Special(expr, nmax): return UOp(Ops.SPECIAL, src=(UOp.const(dtypes.weakint, nmax),), arg=expr) +def Special(expr, nmax): return UOp(Ops.SPECIAL, src=(UOp.const(nmax),), arg=expr) def Variable(expr, nmin, nmax): return UOp.variable(expr, nmin, nmax) def Range(n, nmax): return UOp.range(nmax, n) @@ -207,7 +207,7 @@ class TestImageSimplification(unittest.TestCase): if svalid is not None: check_uop_against_string(self, off.src[1].get_valid(), svalid) else: - self.assertEqual(off.src[1].get_valid(), UOp.const(dtypes.bool, True), "svalid is None but valid is not True") + self.assertEqual(off.src[1].get_valid(), UOp.const(True), "svalid is None but valid is not True") def test_idx_gt_c(self): # (idx1 < c+1).ne(True) ? (..., idx1-1+c) : 0 can drop the valid @@ -455,7 +455,7 @@ class TestImageSimplification(unittest.TestCase): A1 = lidx0*32 + r0*32 + lidx1*4 - 99 valid = ((lidx1 < 1).ne(True)) & ((lidx0 + r0) < 3).ne(True) & ((lidx0 + r0) < 19) alu0 = gidx0 + (A1 % 32)*32 + (A1 // 32 % 16)*1024 - load = get_load_image_uop((1, 16384, 4), valid, (alu0, UOp.const(dtypes.weakint, 0))) + load = get_load_image_uop((1, 16384, 4), valid, (alu0, UOp.const(0))) try: self.check(load, None, "(gidx0+lidx0*1024+r0*1024+lidx1*128+-3168)", "0") except AssertionError: @@ -474,7 +474,7 @@ class TestImageSimplification(unittest.TestCase): A1 = lidx0*16 + r0*16 + lidx1*4 - 51 valid = ((lidx1 < 1).ne(True)) & ((lidx0 + r0) < 3).ne(True) & ((lidx0 + r0) < 11) alu0 = lidx2 + gidx0*4 + (A1 % 16)*64 + (A1 // 16 % 8)*1024 - load = get_load_image_uop((1, 8192, 4), valid, (alu0, UOp.const(dtypes.weakint, 0))) + load = get_load_image_uop((1, 8192, 4), valid, (alu0, UOp.const(0))) try: self.check(load, None, "(lidx2+gidx0*4+lidx0*1024+r0*1024+lidx1*256+-3264)", "0") except AssertionError: @@ -488,7 +488,7 @@ class TestImageSimplification(unittest.TestCase): gidx0 = Special("gidx0", 1064) r12 = Range(12, 3) valid = ((gidx0 < 645).ne(True)) & (gidx0 < 653) - idx = (r12*4 + (gidx0+3)%4 + (gidx0+3)//4*24 - 3888, UOp.const(dtypes.weakint, 0)) + idx = (r12*4 + (gidx0+3)%4 + (gidx0+3)//4*24 - 3888, UOp.const(0)) load = get_load_image_uop((1, 48, 4), valid, idx) self.check(load, None, "(r12*4+(gidx0+3)%4+(gidx0+3)//4*24+-3888)", "0") @@ -496,11 +496,11 @@ class TestImageSimplification(unittest.TestCase): # the fused index pass runs without symbolic, so committing a weak src must not leave a CAST that # symbolic later folds inside the index only: the gate's copy of the expression has to stay the same node f = UOp.variable("f", 0.0, 9.0, dtypes.float) - idx_y = (f + UOp.const(None, 1.0)).cast(dtypes.int) - load = get_load_image_uop((10, 10, 4), (UOp.const(None, -1) < idx_y) & (idx_y < UOp.const(None, 10)), + idx_y = (f + UOp.const(1.0)).cast(dtypes.int) + load = get_load_image_uop((10, 10, 4), (UOp.const(-1) < idx_y) & (idx_y < UOp.const(10)), (Special("gidx0", 10), idx_y)) off = graph_rewrite(load.sink(), pm_lower_index_dtype+indexing_simplify, ctx={}).src[0].src[0] - self.assertEqual(off.src[1].get_valid(), UOp.const(dtypes.bool, True)) + self.assertEqual(off.src[1].get_valid(), UOp.const(True)) class TestDropTrueGate(unittest.TestCase): def test_drop_true_gate_on_index(self): @@ -509,8 +509,8 @@ class TestDropTrueGate(unittest.TestCase): from tinygrad.uop.ops import graph_rewrite from tinygrad.uop.symbolic import sym buf = UOp.param(0, dtypes.int, (1,)) - idx = UOp.const(dtypes.weakint, 0) - true_gate = UOp.const(dtypes.bool, True) + idx = UOp.const(0) + true_gate = UOp.const(True) index_with_gate = UOp(Ops.INDEX, src=(buf, idx.valid(true_gate))) # apply the optimization result = graph_rewrite(index_with_gate, sym+indexing_simplify) @@ -526,7 +526,7 @@ class TestRangeShrink(unittest.TestCase): def test_range_shrink_single_guard(self): # range 0..203 guarded by r < 4 everywhere -> shrink to 0..3 r = Range(0, 204) - load = get_gated_load_uop(r < UOp.const(dtypes.weakint, 4), r) + load = get_gated_load_uop(r < UOp.const(4), r) ranges = self.get_ranges(load.sink()) self.assertEqual(len(ranges), 1) self.assertEqual(ranges[0].src[0].arg, 4) @@ -534,8 +534,8 @@ class TestRangeShrink(unittest.TestCase): def test_range_shrink_picks_max_guard(self): # two loads guard the same range with r < 4 and r < 8 -> shrink to max(4, 8) = 8 r = Range(0, 204) - load1 = get_gated_load_uop(r < UOp.const(dtypes.weakint, 4), r) - load2 = get_gated_load_uop(r < UOp.const(dtypes.weakint, 8), r) + load1 = get_gated_load_uop(r < UOp.const(4), r) + load2 = get_gated_load_uop(r < UOp.const(8), r) ranges = self.get_ranges(UOp.sink(load1, load2)) self.assertEqual(len(ranges), 1) self.assertEqual(ranges[0].src[0].arg, 8) @@ -543,7 +543,7 @@ class TestRangeShrink(unittest.TestCase): def test_range_no_shrink_guard_ge_max(self): # guard r < 300 with range max 204 -> no shrink (guard doesn't constrain) r = Range(0, 204) - load = get_gated_load_uop(r < UOp.const(dtypes.weakint, 300), r) + load = get_gated_load_uop(r < UOp.const(300), r) ranges = self.get_ranges(load.sink()) self.assertEqual(len(ranges), 1) self.assertEqual(ranges[0].src[0].arg, 204) @@ -551,7 +551,7 @@ class TestRangeShrink(unittest.TestCase): def test_range_no_shrink_when_unguarded_elsewhere(self): # one load guards r < 4, but another load uses r without a gate -> no shrink r = Range(0, 204) - load1 = get_gated_load_uop(r < UOp.const(dtypes.weakint, 4), r) + load1 = get_gated_load_uop(r < UOp.const(4), r) load2 = UOp(Ops.LOAD, src=(UOp.param(1, dtypes.float, (204,)).index(r),)) ranges = self.get_ranges(UOp.sink(load1, load2)) self.assertEqual(len(ranges), 1) @@ -560,7 +560,7 @@ class TestRangeShrink(unittest.TestCase): def test_range_no_shrink_when_used_in_reduce(self): # range used in both a gated load AND directly in the reduce expression -> no shrink r = Range(0, 204) - gated_load = get_gated_load_uop(r < UOp.const(dtypes.weakint, 4), r) + gated_load = get_gated_load_uop(r < UOp.const(4), r) red = (r.cast(dtypes.float) + gated_load).reduce(r, arg=Ops.ADD) ranges = self.get_ranges(red.sink()) self.assertEqual(len(ranges), 1) @@ -569,7 +569,7 @@ class TestRangeShrink(unittest.TestCase): def test_range_shrink_to_single_iteration(self): # guard r < 1 shrinks range to 1 -> single iteration, range eliminated entirely r = Range(0, 204) - load = get_gated_load_uop(r < UOp.const(dtypes.weakint, 1), r) + load = get_gated_load_uop(r < UOp.const(1), r) ranges = self.get_ranges(load.sink()) self.assertEqual(len(ranges), 0) @@ -577,7 +577,7 @@ class TestRangeShrink(unittest.TestCase): # emulates mask.where(x.pad_to(mask.shape), Invalid): range should shrink accordingly from tinygrad.dtype import Invalid r = Range(0, 204) - x = (r < 4).where(UOp.const(dtypes.float, 1), Invalid) + x = (r < 4).where(UOp.const(1, dtypes.float), Invalid) ranges = self.get_ranges(UOp.param(0, dtypes.float, (204,)).index(r).store((r < 4).where(x, Invalid)).sink()) self.assertEqual(len(ranges), 1) self.assertEqual(ranges[0].src[0].arg, 4) @@ -586,7 +586,7 @@ class TestRangeShrink(unittest.TestCase): # above, but flipped from tinygrad.dtype import Invalid r = Range(0, 204) - x = (r < 4).where(UOp.const(dtypes.float, 1), Invalid) + x = (r < 4).where(UOp.const(1, dtypes.float), Invalid) ranges = self.get_ranges(UOp.param(0, dtypes.float, (204,)).index(r).store((r >= 4).where(Invalid, x)).sink()) self.assertEqual(len(ranges), 1) self.assertEqual(ranges[0].src[0].arg, 4) diff --git a/test/null/test_tensor_uop_mixin.py b/test/null/test_tensor_uop_mixin.py index 22c0a14c17..983b5c6940 100644 --- a/test/null/test_tensor_uop_mixin.py +++ b/test/null/test_tensor_uop_mixin.py @@ -60,7 +60,7 @@ class TestTensorUOpClone(unittest.TestCase): t = _t(3, 4).float() self.assertIs(_strip_unique(t.clone().uop), _strip_unique(t.uop.clone())) def test_clone_deviceless_const(self): - u = UOp.const(dtypes.float, 2.0) + u = UOp.const(2.0, dtypes.float) self.assertIs(_strip_unique(Tensor(u).clone().uop), _strip_unique(u.clone())) class TestTensorUOpGradient(unittest.TestCase): @@ -382,7 +382,7 @@ class TestTensorUOpStack(unittest.TestCase): self.assertIs(_t(2, 3).uop.stack(w.uop).dtype, dtypes.float32) def test_stack_index_dtype(self): # index is outside the promotion lattice, equal dtypes bypass promotion - self.assertEqual(UOp.const(dtypes.weakint, 1).stack(UOp.const(dtypes.weakint, 2)).shape, (2,)) + self.assertEqual(UOp.const(1).stack(UOp.const(2)).shape, (2,)) class TestTensorUOpConv2d(unittest.TestCase): def test_conv2d_basic(self): diff --git a/test/null/test_transcendental_helpers.py b/test/null/test_transcendental_helpers.py index 504725587c..15a09bbe8a 100644 --- a/test/null/test_transcendental_helpers.py +++ b/test/null/test_transcendental_helpers.py @@ -10,7 +10,7 @@ class TestTranscendentalFunctions(unittest.TestCase): # TODO: Test constant input when constant folding is fixed (or maybe test both variants) # Load input value from a buffer to prevent constant folding input_buf = UOp.param(1, dtypes.double, (1,)) - loaded_value = input_buf.index(UOp.const(dtypes.int, 0)).load() + loaded_value = input_buf.index(UOp.const(0, dtypes.int)).load() def eval_payne_hanek_reduction(v:float) -> tuple[float, int]: return tuple(eval_uop(u, [(dtypes.float64, [v])]) for u in payne_hanek_reduction(loaded_value)) @@ -27,48 +27,48 @@ class TestTranscendentalFunctions(unittest.TestCase): np.testing.assert_equal(q, 4) def test_cody_waite_reduction(self): - r, q = (eval_uop(u) for u in cody_waite_reduction(UOp.const(dtypes.float64, 12 * math.pi + 0.1))) + r, q = (eval_uop(u) for u in cody_waite_reduction(UOp.const(12 * math.pi + 0.1, dtypes.float64))) np.testing.assert_allclose(r, 0.1) np.testing.assert_equal(q, 12) def test_frexp(self): for x in (1, -1): - mantissa, exponent = (eval_uop(u) for u in frexp(UOp.const(dtypes.float64, x))) + mantissa, exponent = (eval_uop(u) for u in frexp(UOp.const(x, dtypes.float64))) np.testing.assert_equal(mantissa, 0.5) np.testing.assert_equal(exponent, 1) for x in (2, -2): - mantissa, exponent = (eval_uop(u) for u in frexp(UOp.const(dtypes.float64, 2.0))) + mantissa, exponent = (eval_uop(u) for u in frexp(UOp.const(2.0, dtypes.float64))) np.testing.assert_equal(mantissa, 0.5) np.testing.assert_equal(exponent, 2) - mantissa, exponent = (eval_uop(u) for u in frexp(UOp.const(dtypes.float64, 5.0))) + mantissa, exponent = (eval_uop(u) for u in frexp(UOp.const(5.0, dtypes.float64))) np.testing.assert_equal(mantissa, 0.625) np.testing.assert_equal(exponent, 3) - mantissa, exponent = (eval_uop(u) for u in frexp(UOp.const(dtypes.float64, 1000.0))) + mantissa, exponent = (eval_uop(u) for u in frexp(UOp.const(1000.0, dtypes.float64))) np.testing.assert_allclose(mantissa, 0.9765625) np.testing.assert_equal(exponent, 10) def test_rintk(self): - np.testing.assert_allclose(eval_uop(rintk(UOp.const(dtypes.float, 0.0))), 0) - np.testing.assert_allclose(eval_uop(rintk(UOp.const(dtypes.float, 5.0))), 5) - np.testing.assert_allclose(eval_uop(rintk(UOp.const(dtypes.float, 5.5))), 6) - np.testing.assert_allclose(eval_uop(rintk(UOp.const(dtypes.float, 5.999))), 6) - np.testing.assert_allclose(eval_uop(rintk(UOp.const(dtypes.float, -5.0))), -5) - np.testing.assert_allclose(eval_uop(rintk(UOp.const(dtypes.float, -5.5))), -6) - np.testing.assert_allclose(eval_uop(rintk(UOp.const(dtypes.float, -5.999))), -6) + np.testing.assert_allclose(eval_uop(rintk(UOp.const(0.0, dtypes.float))), 0) + np.testing.assert_allclose(eval_uop(rintk(UOp.const(5.0, dtypes.float))), 5) + np.testing.assert_allclose(eval_uop(rintk(UOp.const(5.5, dtypes.float))), 6) + np.testing.assert_allclose(eval_uop(rintk(UOp.const(5.999, dtypes.float))), 6) + np.testing.assert_allclose(eval_uop(rintk(UOp.const(-5.0, dtypes.float))), -5) + np.testing.assert_allclose(eval_uop(rintk(UOp.const(-5.5, dtypes.float))), -6) + np.testing.assert_allclose(eval_uop(rintk(UOp.const(-5.999, dtypes.float))), -6) def test_pow2if(self): - np.testing.assert_allclose(eval_uop(pow2if(UOp.const(dtypes.int, 0), dtypes.float)), 1.0) - np.testing.assert_allclose(eval_uop(pow2if(UOp.const(dtypes.int, 1), dtypes.float)), 2.0) - np.testing.assert_allclose(eval_uop(pow2if(UOp.const(dtypes.int, 2), dtypes.float)), 4.0) - np.testing.assert_allclose(eval_uop(pow2if(UOp.const(dtypes.int, 10), dtypes.float)), 1024.0) - np.testing.assert_allclose(eval_uop(pow2if(UOp.const(dtypes.int, 63), dtypes.float)), 2**63) - np.testing.assert_allclose(eval_uop(pow2if(UOp.const(dtypes.int, -1), dtypes.float)), 0.5) - np.testing.assert_allclose(eval_uop(pow2if(UOp.const(dtypes.int, -2), dtypes.float)), 0.25) - np.testing.assert_allclose(eval_uop(pow2if(UOp.const(dtypes.int, -10), dtypes.float)), 2**-10) - np.testing.assert_allclose(eval_uop(pow2if(UOp.const(dtypes.int, -63), dtypes.float)), 2**-63) + np.testing.assert_allclose(eval_uop(pow2if(UOp.const(0, dtypes.int), dtypes.float)), 1.0) + np.testing.assert_allclose(eval_uop(pow2if(UOp.const(1, dtypes.int), dtypes.float)), 2.0) + np.testing.assert_allclose(eval_uop(pow2if(UOp.const(2, dtypes.int), dtypes.float)), 4.0) + np.testing.assert_allclose(eval_uop(pow2if(UOp.const(10, dtypes.int), dtypes.float)), 1024.0) + np.testing.assert_allclose(eval_uop(pow2if(UOp.const(63, dtypes.int), dtypes.float)), 2**63) + np.testing.assert_allclose(eval_uop(pow2if(UOp.const(-1, dtypes.int), dtypes.float)), 0.5) + np.testing.assert_allclose(eval_uop(pow2if(UOp.const(-2, dtypes.int), dtypes.float)), 0.25) + np.testing.assert_allclose(eval_uop(pow2if(UOp.const(-10, dtypes.int), dtypes.float)), 2**-10) + np.testing.assert_allclose(eval_uop(pow2if(UOp.const(-63, dtypes.int), dtypes.float)), 2**-63) if __name__ == '__main__': unittest.main() diff --git a/test/null/test_uop_graph.py b/test/null/test_uop_graph.py index 24d17b6c2e..8a9a4d456b 100644 --- a/test/null/test_uop_graph.py +++ b/test/null/test_uop_graph.py @@ -7,9 +7,9 @@ from tinygrad.uop.symbolic import sym from test.helpers import to_uops_list simple_pm = PatternMatcher([ - (UPat.cvar('x', dtypes.int), lambda x: UOp.const(dtypes.float, 1.0) + UOp.const(dtypes.float, 2.0)), - (UPat.cvar('x') + UPat.cvar('y'), lambda x,y: UOp.const(dtypes.float, x.arg+y.arg)), - (UPat.cvar('x') * UPat.cvar('y') * UPat.cvar('z'), lambda x,y,z: UOp.const(dtypes.float, x.arg*y.arg*z.arg)), + (UPat.cvar('x', dtypes.int), lambda x: UOp.const(1.0, dtypes.float) + UOp.const(2.0, dtypes.float)), + (UPat.cvar('x') + UPat.cvar('y'), lambda x,y: UOp.const(x.arg+y.arg, dtypes.float)), + (UPat.cvar('x') * UPat.cvar('y') * UPat.cvar('z'), lambda x,y,z: UOp.const(x.arg*y.arg*z.arg, dtypes.float)), ((UPat.var('x') + UPat.cvar('c1')) + UPat.cvar('c2'), lambda x,c1,c2: x + (c1.arg+c2.arg)), ]) @@ -20,22 +20,22 @@ def const_values(u:UOp): class TestGraphRewriteConst(unittest.TestCase): def test_gep_const(self): - v1 = UOp.const(dtypes.int, (0,1,2)) + v1 = UOp.const((0,1,2), dtypes.int) v2 = v1.index(1) ret = graph_rewrite(v2, sym) self.assertEqual(ret.dtype, dtypes.int) self.assertEqual(ret.arg, 1) def test_add_const(self): - v1 = UOp.const(dtypes.int, (0,1,2)) - v2 = UOp.const(dtypes.int, (5,6,7)) + v1 = UOp.const((0,1,2), dtypes.int) + v2 = UOp.const((5,6,7), dtypes.int) ret = graph_rewrite(v1+v2, sym) self.assertEqual(ret.op, Ops.STACK) self.assertEqual(const_values(ret), (5,7,9)) def test_add_const_lose_v(self): - v1 = UOp.const(dtypes.int, (0,1,2)) - v2 = UOp.const(dtypes.int, (2,1,0)) + v1 = UOp.const((0,1,2), dtypes.int) + v2 = UOp.const((2,1,0), dtypes.int) ret = graph_rewrite(v1+v2, sym) self.assertEqual(ret.op, Ops.STACK) self.assertEqual(const_values(ret), (2,2,2)) @@ -54,42 +54,42 @@ class TestModularWraparound(unittest.TestCase): @xfail_broken_const_wraparound def test_cast(self): t = self._test - t(UOp.const(dtypes.uint, 0xABCD17D6).cast(dtypes.uint8), 0xD6) - t(UOp.const(dtypes.uint, 0xABCD17D6).cast(dtypes.uint8).cast(dtypes.uint), 0xD6) + t(UOp.const(0xABCD17D6, dtypes.uint).cast(dtypes.uint8), 0xD6) + t(UOp.const(0xABCD17D6, dtypes.uint).cast(dtypes.uint8).cast(dtypes.uint), 0xD6) @xfail_broken_const_wraparound def test_mul(self): t = self._test - t(UOp.const(dtypes.uint, 0xABCD17D6) * 0xAABBCCDD, 1147018174) - t(UOp.const(dtypes.int, 0xABCD17D6) * 10, -1241321892) + t(UOp.const(0xABCD17D6, dtypes.uint) * 0xAABBCCDD, 1147018174) + t(UOp.const(0xABCD17D6, dtypes.int) * 10, -1241321892) @xfail_broken_const_wraparound def test_div(self): t = self._test - t(UOp.const(dtypes.uint, 0xABCD17D6) * 0xAABBCCDD // 11, 104274379) - t(UOp.const(dtypes.int, 0xABCD17D6) * 10 // 11, -112847444) + t(UOp.const(0xABCD17D6, dtypes.uint) * 0xAABBCCDD // 11, 104274379) + t(UOp.const(0xABCD17D6, dtypes.int) * 10 // 11, -112847444) @xfail_broken_const_wraparound def test_neg(self): t = self._test - t(-UOp.const(dtypes.uint8, 1), 0xFF) - t(-UOp.const(dtypes.uint16, 1), 0xFFFF) - t(-UOp.const(dtypes.uint32, 1), 0xFFFFFFFF) - t(-UOp.const(dtypes.uint64, 1), 0xFFFFFFFFFFFFFFFF) + t(-UOp.const(1, dtypes.uint8), 0xFF) + t(-UOp.const(1, dtypes.uint16), 0xFFFF) + t(-UOp.const(1, dtypes.uint32), 0xFFFFFFFF) + t(-UOp.const(1, dtypes.uint64), 0xFFFFFFFFFFFFFFFF) @xfail_broken_const_wraparound def test_neg_min_int(self): t = self._test - t(-UOp.const(dtypes.int8, -2**7), -2**7) - t(-UOp.const(dtypes.int16, -2**15), -2**15) - t(-UOp.const(dtypes.int32, -2**31), -2**31) - t(-UOp.const(dtypes.int64, -2**63), -2**63) + t(-UOp.const(-2**7, dtypes.int8), -2**7) + t(-UOp.const(-2**15, dtypes.int16), -2**15) + t(-UOp.const(-2**31, dtypes.int32), -2**31) + t(-UOp.const(-2**63, dtypes.int64), -2**63) @xfail_broken_const_wraparound def test_payne_hanek_reduction_bug(self): t = self._test - a = (UOp.const(dtypes.uint, 43748177600).cast(dtypes.uint) | 36).cast(dtypes.ulong) - b = 2536655455 * a + 4294967296 * UOp.const(dtypes.ulong, 25366554550) + a = (UOp.const(43748177600, dtypes.uint).cast(dtypes.uint) | 36).cast(dtypes.ulong) + b = 2536655455 * a + 4294967296 * UOp.const(25366554550, dtypes.ulong) c = (b + 2261737165) // 4611686018427387904 t(c, 0) @@ -103,62 +103,62 @@ class TestGraphRewrite(unittest.TestCase): # NOTE: this shows why we can't have a UOp in arg @unittest.expectedFailure def test_no_dedup_args(self): - a1 = UOp.variable("a1", UOp.const(dtypes.int, 0), UOp.const(dtypes.int, 11), dtypes.int) - a2 = UOp.variable("a2", UOp.const(dtypes.int, 0), UOp.const(dtypes.int, 11), dtypes.int) + a1 = UOp.variable("a1", UOp.const(0, dtypes.int), UOp.const(11, dtypes.int), dtypes.int) + a2 = UOp.variable("a2", UOp.const(0, dtypes.int), UOp.const(11, dtypes.int), dtypes.int) sink = a1.sink(a2) variables = [x for x in graph_rewrite(sink, PatternMatcher([])).toposort() if x.op is Ops.PARAM and x.addrspace is AddrSpace.ALU] self.assertEqual(len(variables), 1) def test_simple(self): - c1 = UOp.const(dtypes.float, 1.0) - c2 = UOp.const(dtypes.float, 2.0) + c1 = UOp.const(1.0, dtypes.float) + c2 = UOp.const(2.0, dtypes.float) nout = graph_rewrite(c1+c2, simple_pm) self.assertEqual(nout.op, Ops.CONST) self.assertEqual(nout.arg, 3.0) def test_depth_2_late(self): - c1 = UOp.const(dtypes.float, 1.0) - c2 = UOp.const(dtypes.float, 2.0) - c3 = UOp.const(dtypes.float, 3.0) + c1 = UOp.const(1.0, dtypes.float) + c2 = UOp.const(2.0, dtypes.float) + c3 = UOp.const(3.0, dtypes.float) nout = graph_rewrite(c1*c2*(c3+c3), simple_pm) self.assertEqual(nout.op, Ops.CONST) self.assertEqual(nout.arg, 12.0) def test_double(self): - c1 = UOp.const(dtypes.float, 1.0) - c2 = UOp.const(dtypes.float, 2.0) - c3 = UOp.const(dtypes.float, 3.0) + c1 = UOp.const(1.0, dtypes.float) + c2 = UOp.const(2.0, dtypes.float) + c3 = UOp.const(3.0, dtypes.float) nout = graph_rewrite(c1+c2+c3, simple_pm) self.assertEqual(nout.op, Ops.CONST) self.assertEqual(nout.arg, 6.0) def test_triple(self): - c1 = UOp.const(dtypes.float, 1.0) - c2 = UOp.const(dtypes.float, 2.0) - c3 = UOp.const(dtypes.float, 3.0) - c4 = UOp.const(dtypes.float, 4.0) + c1 = UOp.const(1.0, dtypes.float) + c2 = UOp.const(2.0, dtypes.float) + c3 = UOp.const(3.0, dtypes.float) + c4 = UOp.const(4.0, dtypes.float) nout = graph_rewrite(c1+c2+c3+c4, simple_pm) self.assertEqual(nout.op, Ops.CONST) self.assertEqual(nout.arg, 10.0) def test_diamond(self): - c1 = UOp.const(dtypes.float, 1.0) - c2 = UOp.const(dtypes.float, 2.0) - c3 = UOp.const(dtypes.float, 3.0) + c1 = UOp.const(1.0, dtypes.float) + c2 = UOp.const(2.0, dtypes.float) + c3 = UOp.const(3.0, dtypes.float) nout = graph_rewrite((c1+c2)+(c1+c3), simple_pm) self.assertEqual(nout.op, Ops.CONST) self.assertEqual(nout.arg, 7.0) def test_magic_4(self): - c1 = UOp.const(dtypes.int, 4.0) + c1 = UOp.const(4.0, dtypes.int) nout = graph_rewrite(c1, simple_pm) self.assertEqual(nout.op, Ops.CONST) self.assertEqual(nout.arg, 3.0) def test_depth_2_fold(self): v = UOp.variable("v", 0, 1, dtypes.float) - c1 = UOp.const(dtypes.float, 1.0) - c2 = UOp.const(dtypes.float, 2.0) + c1 = UOp.const(1.0, dtypes.float) + c2 = UOp.const(2.0, dtypes.float) nout = graph_rewrite(v+c1+c2, simple_pm) self.assertEqual(nout.op, Ops.ADD) self.assertEqual(nout.src[0].op, Ops.PARAM) @@ -191,8 +191,8 @@ class TestGraphRewrite(unittest.TestCase): class TestUOpGraph(unittest.TestCase): def test_add_constant_fold(self): - c1 = UOp.const(dtypes.float, 1.0) - c2 = UOp.const(dtypes.float, 2.0) + c1 = UOp.const(1.0, dtypes.float) + c2 = UOp.const(2.0, dtypes.float) out = c1+c2 uops = to_uops_list([out]) self.assertEqual(len(uops), 2) # +1 for SINK @@ -202,9 +202,9 @@ class TestUOpGraph(unittest.TestCase): def test_where_same_fold(self): v = UOp.variable('tmp', 0, 1) - c0 = UOp.const(dtypes.weakint, 0) + c0 = UOp.const(0) vc = v != c0 - c1 = UOp.const(dtypes.float, 1.0) + c1 = UOp.const(1.0, dtypes.float) out = vc.where(c1, c1) uops = to_uops_list([out]) self.assertEqual(len(uops), 2) # +1 for SINK @@ -213,9 +213,9 @@ class TestUOpGraph(unittest.TestCase): self.assertEqual(out.arg, 1.0) def test_where_const_fold(self): - bf = UOp.const(dtypes.bool, False) - c1 = UOp.const(dtypes.float, 1.0) - c2 = UOp.const(dtypes.float, 2.0) + bf = UOp.const(False) + c1 = UOp.const(1.0, dtypes.float) + c2 = UOp.const(2.0, dtypes.float) out = bf.where(c1, c2) uops = to_uops_list([out]) self.assertEqual(len(uops), 2) # +1 for SINK @@ -224,7 +224,7 @@ class TestUOpGraph(unittest.TestCase): self.assertEqual(out.arg, 2.0) def test_const_cast(self): - bf = UOp.const(dtypes.bool, False) + bf = UOp.const(False) out = bf.cast(dtypes.int) uops = to_uops_list([out]) self.assertEqual(len(uops), 2) # +1 for SINK @@ -233,7 +233,7 @@ class TestUOpGraph(unittest.TestCase): self.assertEqual(out.arg, 0) def test_const_bitcast(self): - bf = UOp.const(dtypes.float, 1.0) + bf = UOp.const(1.0, dtypes.float) out = bf.bitcast(dtypes.uint32) uops = to_uops_list([out]) self.assertEqual(len(uops), 2) # +1 for SINK @@ -243,7 +243,7 @@ class TestUOpGraph(unittest.TestCase): @unittest.expectedFailure def test_const_shape_change_bitcast(self): - bf = UOp.const(dtypes.uint8, 0x3F) + bf = UOp.const(0x3F, dtypes.uint8) out = bf.bitcast(dtypes.half) uops = to_uops_list([out]) self.assertEqual(len(uops), 2) # +1 for SINK @@ -251,7 +251,7 @@ class TestUOpGraph(unittest.TestCase): def test_devectorize_derives_lane_dtype(self): from tinygrad.codegen import do_devectorize # an Invalid lane derives bool while the value lane derives float: the lane rebuild must derive, not inherit - lhs = UOp.stack(UOp.invalid(), UOp.const(None, 1.0).cast(dtypes.float)) + lhs = UOp.stack(UOp.invalid(), UOp.const(1.0).cast(dtypes.float)) out = do_devectorize(lhs * lhs) invalid_lane_mul = next(u for u in out.src[0].toposort() if u.op is Ops.MUL) self.assertIs(invalid_lane_mul.dtype, dtypes.bool) @@ -259,7 +259,7 @@ class TestUOpGraph(unittest.TestCase): @unittest.skip("this test isn't valid uops") def test_noop_vectorize_fold(self): d0 = UOp.param(0, dtypes.float, (1,)) - idx = UOp.const(dtypes.int, 0) + idx = UOp.const(0, dtypes.int) ld = d0.load(idx, dtype=dtypes.float) vec = UOp(Ops.STACK, dtypes.float, (ld,)) x = vec.index(0) @@ -273,7 +273,7 @@ class TestUOpGraph(unittest.TestCase): d0 = UOp.param(0, dtypes.float, (1,)) d1 = UOp.param(1, dtypes.float, (1,)) d2 = UOp.param(2, dtypes.float, (1,)) - idx = UOp.const(dtypes.int, 0) + idx = UOp.const(0, dtypes.int) def _test_vec(geps, count=4): vec = UOp(Ops.STACK, dtypes.float, geps) out = d0.index(idx).store(vec) @@ -310,7 +310,7 @@ class TestUOpGraph(unittest.TestCase): def test_gep_vec_const_fold(self): for vec_size in [2, 4, 8]: - consts = [UOp.const(dtypes.float, float(i)) for i in range(vec_size)] + consts = [UOp.const(float(i), dtypes.float) for i in range(vec_size)] vec = UOp(Ops.STACK, src=tuple(consts)) with Context(SPEC=0): uops = to_uops_list([vec.index(i) for i in range(vec_size)]) @@ -320,7 +320,7 @@ class TestUOpGraph(unittest.TestCase): def test_cast_alu_fold(self): d0 = UOp.param(0, dtypes.bool, (1,)) d1 = UOp.param(1, dtypes.int, (1,)) - idx = UOp.const(dtypes.int, 0) + idx = UOp.const(0, dtypes.int) ld = d1.index(idx) alu = (ld<1).cast(dtypes.bool) out = d0.index(idx).store(alu) @@ -330,7 +330,7 @@ class TestUOpGraph(unittest.TestCase): def test_double_cast_fold(self): d0 = UOp.param(0, dtypes.float, (1,)) d1 = UOp.param(1, dtypes.int, (1,)) - idx = UOp.const(dtypes.int, 0) + idx = UOp.const(0, dtypes.int) ld = d1.index(idx) alu = ld.cast(dtypes.float).cast(dtypes.float) out = d0.index(idx).store(alu) @@ -339,8 +339,8 @@ class TestUOpGraph(unittest.TestCase): def test_depth_2_const_fold(self): v = UOp.variable("tmp", 0, 1, dtypes.int) - c2 = UOp.const(dtypes.int, 2) - c4 = UOp.const(dtypes.int, 4) + c2 = UOp.const(2, dtypes.int) + c4 = UOp.const(4, dtypes.int) vc = v+c2 out = vc+c4 uops = to_uops_list([out]) @@ -353,14 +353,14 @@ class TestUOpGraph(unittest.TestCase): def test_bitcast_to_same_dtype_fold(self): for dt in dtypes.ints + dtypes.floats + (dtypes.bool,): d0 = UOp.param(0, dt, (1,)) - v = d0.index(UOp.const(dtypes.int, 0)) + v = d0.index(UOp.const(0, dtypes.int)) uops = to_uops_list([v.bitcast(dt)]) self.assertEqual(len([x for x in uops if x.op is Ops.BITCAST and x.dtype is dt]), 0, f"dtype = {dt}") def test_sub_with_cast_folds(self): a = Variable("a", 0, 5) uops = to_uops_list([a.cast(dtypes.int)+(-a).cast(dtypes.int)]) - assert uops[0] == UOp.const(dtypes.int, 0) + assert uops[0] == UOp.const(0, dtypes.int) assert uops[-1].op == Ops.SINK def test_where_on_gated_load_fold(self): @@ -400,7 +400,7 @@ class TestUOpGraph(unittest.TestCase): ridx0 = UOp.range(100, 0) d0 = UOp.param(0, dtypes.float, (100,)) ld = d0.index(ridx0.valid(ridx0<50)) - w = ((ridx0<50) & (ridx0>30)).where(ld, UOp.const(dtypes.float, 0)).cast(dtypes.half) + w = ((ridx0<50) & (ridx0>30)).where(ld, UOp.const(0, dtypes.float)).cast(dtypes.half) out = UOp.param(1, dtypes.half, (100,)) uops = to_uops_list([out.index(ridx0).store(w)]) for u in uops: @@ -410,7 +410,7 @@ class TestUOpGraph(unittest.TestCase): ridx0 = UOp.range(100, 0) d0 = UOp.param(0, dtypes.float, (100,)) ld = d0.index(ridx0.valid(ridx0<50)) - w = ((ridx0<50) & (ridx0>30)).where(UOp.const(dtypes.float, 0), ld).cast(dtypes.half) + w = ((ridx0<50) & (ridx0>30)).where(UOp.const(0, dtypes.float), ld).cast(dtypes.half) out = UOp.param(1, dtypes.half, (100,)) uops = to_uops_list([out.index(ridx0).store(w)]) for u in uops: @@ -432,16 +432,16 @@ class TestUOpGraph(unittest.TestCase): # mnist indexing with split reduceop # Make sure we are not doign math on the loaded index, which would promote it to long c0 = UOp.param(0, dtypes.uchar, (128000,)) - c1 = UOp.range(UOp.const(dtypes.weakint, 512), 1, AxisType.WEAK) - c2 = UOp.range(UOp.const(dtypes.weakint, 250), 2, AxisType.WEAK) + c1 = UOp.range(UOp.const(512), 1, AxisType.WEAK) + c2 = UOp.range(UOp.const(250), 2, AxisType.WEAK) c3 = UOp.param(1, dtypes.int, (512,)) c4 = c3.index(c1) - c5 = UOp.range(UOp.const(dtypes.weakint, 240), 0, AxisType.REDUCE) - c6 = ((c2*UOp.const(dtypes.weakint, 240))+c5) + c5 = UOp.range(UOp.const(240), 0, AxisType.REDUCE) + c6 = ((c2*UOp.const(240))+c5) c7 = UOp.param(2, dtypes.uchar, (60000,)) c8 = c7.index(c6) c9 = ((c4<0).where((c4+60000), c4)!=c6.cast(dtypes.int)).where(0, c8.cast(dtypes.uint).cast(dtypes.uchar)).reduce(c5, arg=Ops.ADD) - c10 = c0.index(((c1*UOp.const(dtypes.weakint, 250))+c2)).store(c9).end(c1, c2) + c10 = c0.index(((c1*UOp.const(250))+c2)).store(c9).end(c1, c2) uops = to_uops_list([c10]) for u in uops: self.assertNotEqual(u.dtype, dtypes.long) @@ -449,19 +449,19 @@ class TestUOpGraph(unittest.TestCase): def test_load_idx_no_math_on_loaded(self): # test the (x+y) 7 + u = UOp.const(4, dtypes.int) > 7 self.assertFalse(u) def test_ssimplify(self): - self.assertEqual((8 % UOp.const(dtypes.int, 4)).ssimplify(), 0) - self.assertEqual((8 * UOp.const(dtypes.int, 4)).ssimplify(), 32) + self.assertEqual((8 % UOp.const(4, dtypes.int)).ssimplify(), 0) + self.assertEqual((8 * UOp.const(4, dtypes.int)).ssimplify(), 32) def test_ambiguous_less_than(self): u = UOp.variable("i", 1, 10) @@ -52,7 +52,7 @@ class TestUOpResolve(unittest.TestCase): self.assertFalse(resolve(u < -1, True)) def test_float_direct(self): - u = UOp.const(dtypes.float, 4.5) + 7 + u = UOp.const(4.5, dtypes.float) + 7 self.assertEqual(float(u), 11.5) def test_var_cmp_t(self): diff --git a/test/null/test_uop_symbolic.py b/test/null/test_uop_symbolic.py index 3a2912f96e..fa77ae4806 100644 --- a/test/null/test_uop_symbolic.py +++ b/test/null/test_uop_symbolic.py @@ -12,13 +12,13 @@ from tinygrad.uop.validate import uops_to_z3 def check_uop_against_string(self, v:UOp, s:str): sym_vars = {v.render():v for v in v.toposort() if v.op in (Ops.RANGE, Ops.SPECIAL, Ops.PARAM)} s_eval = eval(s, sym_vars) - if isinstance(s_eval, int) and v.dtype==dtypes.weakint: s_eval = UOp.const(dtypes.weakint, s_eval) - elif isinstance(s_eval, (bool, int, float)): s_eval = UOp.const(dtypes.from_py(s_eval), s_eval) + if isinstance(s_eval, int) and v.dtype==dtypes.weakint: s_eval = UOp.const(s_eval, dtypes.weakint) + elif isinstance(s_eval, (bool, int, float)): s_eval = UOp.const(s_eval) s_eval = graph_rewrite(s_eval, commutative, name="cannonicalize eval") self.assertIs(s_eval, v, f"eval did not match simplified: {s_eval} != {v.render()} for {s}") def Variable(name: str, min_val: ConstType, max_val: ConstType, dtype: DType=dtypes.weakint): return UOp.variable(name,min_val,max_val,dtype) -def uconst(val): return UOp.const(dtypes.weakint, val) +def uconst(val): return UOp.const(val) def usum(ops): return functools.reduce(lambda x,y: x+y, ops) def uand(ops): return functools.reduce(lambda x,y: x*y, ops) @@ -1017,12 +1017,12 @@ class TestSymbolic(unittest.TestCase): # TODO: copied from render, render does not support cast glbl = UOp.param(0, dtypes.int, (1,)) - uops = get_uops(UOp(Ops.STORE, src=(glbl.index(UOp.const(dtypes.int, 0)), expr)).sink()) + uops = get_uops(UOp(Ops.STORE, src=(glbl.index(UOp.const(0, dtypes.int)), expr)).sink()) rewritten_uop = [uop for uop in uops if uop.op is Ops.STORE][0].src[1] # the vars are now scalar PARAMs pvar = {u.expr: u for u in rewritten_uop.toposort() if u.op is Ops.PARAM} - self.assertEqual(rewritten_uop, (pvar['s'] NOOP rule. This rule matches patterns that EMERGE during simplification.""" def test_store_load_folding(self): # store(idx, load(idx)) -> NOOP, including emergent patterns like store(idx, load(idx) + 0) buf = UOp.param(0, dtypes.int, (1,)) - index = buf.index(UOp.const(dtypes.weakint, 0)) + index = buf.index(UOp.const(0)) # Direct: store(idx, load(idx)) -> NOOP self.assertEqual(graph_rewrite(index.store(index.load()), sym).op, Ops.NOOP) # Emergent: store(idx, load(idx) + 0) -> store(idx, load(idx)) -> NOOP - self.assertEqual(graph_rewrite(index.store(index.load() + UOp.const(dtypes.int, 0)), sym).op, Ops.NOOP) + self.assertEqual(graph_rewrite(index.store(index.load() + UOp.const(0, dtypes.int)), sym).op, Ops.NOOP) # Emergent: store(idx, load(idx) * 1) -> store(idx, load(idx)) -> NOOP - self.assertEqual(graph_rewrite(index.store(index.load() * UOp.const(dtypes.int, 1)), sym).op, Ops.NOOP) + self.assertEqual(graph_rewrite(index.store(index.load() * UOp.const(1, dtypes.int)), sym).op, Ops.NOOP) # 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) + self.assertEqual(graph_rewrite(index.store(index.load() + UOp.const(1, dtypes.int)), sym).op, Ops.STORE) class TestMoveWhereOnLoad(unittest.TestCase): def test_bool_index_preserves_dtype(self): diff --git a/test/null/test_uop_vmin_vmax.py b/test/null/test_uop_vmin_vmax.py index d43f700fa4..b426cea615 100644 --- a/test/null/test_uop_vmin_vmax.py +++ b/test/null/test_uop_vmin_vmax.py @@ -5,12 +5,12 @@ from tinygrad.dtype import dtypes, Invalid class TestVminVmaxProperties(unittest.TestCase): def test_vmin_vmax_constant(self): # vmin and vmax for a constant - uop = UOp.const(dtypes.int32, 42) + uop = UOp.const(42, dtypes.int32) self.assertEqual(uop.vmin, 42) self.assertEqual(uop.vmax, 42) def test_vmin_vmax_cmpne(self): - uop = UOp.const(dtypes.int32, 42) + uop = UOp.const(42, dtypes.int32) def test_bool(u, x): self.assertEqual(u.vmin, x) self.assertEqual(u.vmax, x) @@ -81,8 +81,8 @@ class TestVminVmaxProperties(unittest.TestCase): def test_vmin_vmax_multiplication_0_inf(self): # vmin and vmax for multiplication with a variable - x = UOp.const(dtypes.float, 0.0) - y = UOp.load(UOp.param(0, dtypes.float, (1,)), UOp.const(dtypes.int, 0), dtype=dtypes.float) + x = UOp.const(0.0, dtypes.float) + y = UOp.load(UOp.param(0, dtypes.float, (1,)), UOp.const(0, dtypes.int), dtype=dtypes.float) uop = x * y # TODO: these should be 0, but definitely should not be nan self.assertEqual(uop.vmin, -math.inf) @@ -167,7 +167,7 @@ class TestVminVmaxProperties(unittest.TestCase): self.assertNotEqual(i.vmin, i.vmax) def test_vmin_vmax_invalid_vconst(self): - x = UOp.const(dtypes.weakint, (0, 4, Invalid, Invalid)) + x = UOp.const((0, 4, Invalid, Invalid), dtypes.weakint) self.assertEqual((x.vmin, x.vmax), (0, 4)) class TestVminVmaxDivMod(unittest.TestCase): @@ -286,44 +286,44 @@ class TestVminVmaxDivMod(unittest.TestCase): class TestVminVmaxVConst(unittest.TestCase): def test_vmin_vmax_vconst_single_element(self): # vmin and vmax for a single-element vector constant - uop = UOp.const(dtypes.int32, (42,)) + uop = UOp.const((42,), dtypes.int32) self.assertEqual(uop.vmin, 42) self.assertEqual(uop.vmax, 42) def test_vmin_vmax_vconst_multiple_elements(self): # vmin and vmax for a multi-element vector constant - uop = UOp.const(dtypes.int32, (10, 20, -5, 7)) + uop = UOp.const((10, 20, -5, 7), dtypes.int32) self.assertEqual(uop.vmin, -5) self.assertEqual(uop.vmax, 20) def test_vmin_vmax_vconst_all_equal(self): # vmin and vmax for a vector where all elements are equal - uop = UOp.const(dtypes.int32, (7, 7, 7)) + uop = UOp.const((7, 7, 7), dtypes.int32) self.assertEqual(uop.vmin, 7) self.assertEqual(uop.vmax, 7) def test_vmin_vmax_vconst_with_negative_values(self): # vmin and vmax for a vector constant containing negative values - uop = UOp.const(dtypes.int32, (-10, -20, -5, -15)) + uop = UOp.const((-10, -20, -5, -15), dtypes.int32) self.assertEqual(uop.vmin, -20) self.assertEqual(uop.vmax, -5) def test_vmin_vmax_vconst_with_floats(self): # vmin and vmax for a vector constant of float values - uop = UOp.const(dtypes.float32, (1.5, -3.2, 0.0)) + uop = UOp.const((1.5, -3.2, 0.0), dtypes.float32) self.assertEqual(uop.vmin, -3.2) self.assertEqual(uop.vmax, 1.5) def test_vmin_vmax_vconst_with_bools(self): # vmin and vmax for a vector constant of bool values - uop = UOp.const(dtypes.bool, (True, False, False)) + uop = UOp.const((True, False, False)) self.assertIs(uop.vmin, False) self.assertIs(uop.vmax, True) def test_vmin_vmax_vector_with_gep(self): # vmin and vmax for a vector constant of bool values d1 = UOp.param(1, dtypes.int, (1,)) - idx = UOp.const(dtypes.int, 0) + idx = UOp.const(0, dtypes.int) val = UOp(Ops.LOAD, src=(d1.index(idx),)) uop = (val // 32) self.assertEqual(uop.vmin, -67108864) @@ -332,17 +332,17 @@ class TestVminVmaxVConst(unittest.TestCase): class TestConstFactor(unittest.TestCase): def test_const_factor_constant(self): # const_factor for a constant - uop = UOp.const(dtypes.int32, 42) + uop = UOp.const(42, dtypes.int32) self.assertEqual(uop.const_factor(), 42) def test_const_factor_addition(self): # const_factor for an addition of constants - uop = UOp.const(dtypes.int32, 30) + UOp.const(dtypes.int32, 12) + uop = UOp.const(30, dtypes.int32) + UOp.const(12, dtypes.int32) self.assertEqual(uop.const_factor(), 6) # GCD(30, 12) = 6 def test_const_factor_multiplication(self): # const_factor for a multiplication of constants - uop = UOp.const(dtypes.int32, 5) * UOp.const(dtypes.int32, 7) + uop = UOp.const(5, dtypes.int32) * UOp.const(7, dtypes.int32) self.assertEqual(uop.const_factor(), 5) # For multiplication, it's one of the factors def test_const_factor_with_variable(self): @@ -377,14 +377,14 @@ class TestConstFactor(unittest.TestCase): class TestDivides(unittest.TestCase): def test_divides_constant_exact(self): # Divides a constant by an exact divisor - uop = UOp.const(dtypes.int32, 42) + uop = UOp.const(42, dtypes.int32) result = uop.divides(7) self.assertIsNotNone(result) self.assertEqual(result.const_factor(), 6) # 42 / 7 = 6 def test_divides_constant_inexact(self): # Try to divide a constant by a non-exact divisor - uop = UOp.const(dtypes.int32, 42) + uop = UOp.const(42, dtypes.int32) result = uop.divides(5) self.assertIsNone(result) # 42 is not divisible by 5 diff --git a/test/null/test_uops.py b/test/null/test_uops.py index 0226024ff6..1e85d26201 100644 --- a/test/null/test_uops.py +++ b/test/null/test_uops.py @@ -12,16 +12,16 @@ from test.helpers import eval_uop, to_uops_list class TestDTypeFromUOp(unittest.TestCase): def test_broadcastable_promotion(self): - self.assertEqual(dtype_from_uop(Ops.ADD, (UOp.const(dtypes.float32, 1.0), UOp.const(dtypes.float16, 1.0)), None), dtypes.float32) - self.assertEqual(dtype_from_uop(Ops.MUL, (UOp.const(dtypes.int8, 1), UOp.const(dtypes.int32, 1)), None), dtypes.int32) + self.assertEqual(dtype_from_uop(Ops.ADD, (UOp.const(1.0, dtypes.float32), UOp.const(1.0, dtypes.float16)), None), dtypes.float32) + self.assertEqual(dtype_from_uop(Ops.MUL, (UOp.const(1, dtypes.int8), UOp.const(1, dtypes.int32)), None), dtypes.int32) def test_same_dtype_fast_path(self): - src = (UOp.const(dtypes.weakint, 1), UOp.const(dtypes.weakint, 2)) + src = (UOp.const(1), UOp.const(2)) self.assertEqual(dtype_from_uop(Ops.ADD, src, None), dtypes.weakint) def test_where_promotion(self): - cond = UOp.const(dtypes.bool, True) - self.assertEqual(dtype_from_uop(Ops.WHERE, (cond, UOp.const(dtypes.float32, 1.0), UOp.const(dtypes.float16, 1.0)), None), dtypes.float32) + cond = UOp.const(True) + self.assertEqual(dtype_from_uop(Ops.WHERE, (cond, UOp.const(1.0, dtypes.float32), UOp.const(1.0, dtypes.float16)), None), dtypes.float32) idx = UOp.range(4, 0) self.assertEqual(idx.valid(idx < 4).dtype, dtypes.weakint) @@ -37,33 +37,33 @@ class TestDTypeFromUOp(unittest.TestCase): self.assertEqual(UOp(Ops.CONST, arg=True).dtype, dtypes.bool) self.assertEqual(UOp(Ops.CONST, arg=Invalid).dtype, dtypes.bool) # an explicit (strong) const dtype is legal until the field is removed - self.assertEqual(UOp.const(dtypes.int32, 3).dtype, dtypes.int32) + self.assertEqual(UOp.const(3, dtypes.int32).dtype, dtypes.int32) def test_weak_dtype_rejected_by_program_spec(self): for weak, concrete, value in ((dtypes.weakint, dtypes.int32, 1), (dtypes.weakfloat, dtypes.float32, 1.0)): - with self.assertRaises(RuntimeError): type_verify(UOp.const(weak, value).sink(), spec_program) - type_verify(UOp.const(concrete, value).sink(), spec_program) + with self.assertRaises(RuntimeError): type_verify(UOp.const(value, weak).sink(), spec_program) + type_verify(UOp.const(value, concrete).sink(), spec_program) def test_invalid_dtype_and_consumers(self): invalid = UOp.invalid() self.assertIs(invalid.dtype, dtypes.bool) - self.assertIs(UOp.const(dtypes.float32, Invalid), invalid) + self.assertIs(UOp.const(Invalid, dtypes.float32), invalid) self.assertIs((moved:=invalid.reshape((1,))).cast(dtypes.float32), moved) scratch = Tensor.invalids(4, dtype=dtypes.float32) self.assertEqual((scratch.dtype, next(u.dtype for u in scratch.uop.toposort() if u.op is Ops.BUFFER), next(u.dtype for u in scratch.uop.toposort() if u.arg is Invalid)), (dtypes.float32, dtypes.float32, dtypes.bool)) - invalid, value = UOp.invalid(), UOp.const(dtypes.float32, 1) + invalid, value = UOp.invalid(), UOp.const(1, dtypes.float32) for u in (UOp(Ops.STACK, dtypes.float32, src=(value, invalid)), UOp(Ops.ADD, dtypes.float32, src=(value, invalid)), - UOp.const(dtypes.bool, True).where(value, invalid), UOp(Ops.CMPLT, src=(invalid, value)), UOp(Ops.CMPLT, src=(value, invalid)), + UOp.const(True).where(value, invalid), UOp(Ops.CMPLT, src=(invalid, value)), UOp(Ops.CMPLT, src=(value, invalid)), UOp.param(0, dtypes.float32, (4,)).index(invalid)): type_verify(u, spec_shared) gate, value = UOp.param(0, dtypes.bool, ()), UOp.param(1, dtypes.float, ()) - self.assertIs((out:=graph_rewrite(gate.where(value, UOp.invalid()), pm_remove_invalid)).src[2], UOp.const(dtypes.float, 0)) + self.assertIs((out:=graph_rewrite(gate.where(value, UOp.invalid()), pm_remove_invalid)).src[2], UOp.const(0, dtypes.float)) type_verify(out.sink(), spec_program) def test_remove_invalid_stack_lanes(self): - stack = UOp(Ops.STACK, dtypes.half, (UOp.const(dtypes.half, 1), UOp.invalid())) + stack = UOp(Ops.STACK, dtypes.half, (UOp.const(1, dtypes.half), UOp.invalid())) out = graph_rewrite(stack, pm_remove_invalid) - self.assertEqual(out.src, (UOp.const(dtypes.half, 1), UOp.const(dtypes.half, 0))) + self.assertEqual(out.src, (UOp.const(1, dtypes.half), UOp.const(0, dtypes.half))) type_verify(out.sink(), spec_program) class TestLowerIndexDtype(unittest.TestCase): @@ -72,7 +72,7 @@ class TestLowerIndexDtype(unittest.TestCase): # width the offset bounds select (this one needs long) buf = UOp.param(0, dtypes.float, (2**31+64,)) i = UOp.variable("i", 0, 2**28) - shrink = UOp(Ops.SHRINK, src=(buf, (i*24).valid(i < 2**28), UOp.const(dtypes.weakint, 4))) + shrink = UOp(Ops.SHRINK, src=(buf, (i*24).valid(i < 2**28), UOp.const(4))) lowered = graph_rewrite(shrink.sink(), pm_lower_index_dtype) self.assertTrue(all(u.dtype != dtypes.weakint for u in lowered.backward_slice_with_self), "lowering must resolve all weakint") sh = next(u for u in lowered.backward_slice_with_self if u.op is Ops.SHRINK) @@ -198,9 +198,9 @@ class TestGatedStoreRewrite(unittest.TestCase): def test_tiny_gate_store(self): gmem = UOp.param(0, dtypes.float, (8,)) gidx0 = UOp.special(4, 'gidx0') - gate = gidx0= 0) & (ld0 < 64))).load(dtype=dtypes.int)]) @@ -150,7 +150,7 @@ class TestValidateOOB(unittest.TestCase): with Context(CHECK_OOB=1, SPEC=2): buf_bool = UOp.param(0, dtypes.bool, (16,)) buf_int = UOp.param(1, dtypes.int, (8,)) - gidx = UOp(Ops.SPECIAL, src=(UOp.const(dtypes.weakint, 16),), arg="gidx0") + gidx = UOp(Ops.SPECIAL, src=(UOp.const(16),), arg="gidx0") ld_bool = buf_bool.index(gidx).load() with self.assertRaises(RuntimeError): to_uops_list([buf_int.index(gidx.valid(ld_bool)).load()]) # gidx 0..15, buf_int size 8 @@ -164,12 +164,12 @@ class TestValidateOOB(unittest.TestCase): sbuf = UOp.placeholder((8,), dtypes.uint, slot=0, addrspace=AddrSpace.LOCAL) # Define indices, valids and barrier - gidx = UOp(Ops.SPECIAL, src=(UOp.const(dtypes.int, 416),), arg="gidx0") - lidx = UOp(Ops.SPECIAL, src=(UOp.const(dtypes.int, 10),), arg="lidx0") + gidx = UOp(Ops.SPECIAL, src=(UOp.const(416, dtypes.int),), arg="gidx0") + lidx = UOp(Ops.SPECIAL, src=(UOp.const(10, dtypes.int),), arg="lidx0") gate = (gidx<400) & (lidx<8) - local_store = sbuf.index(lidx.valid(lidx<8)).store(UOp.const(dtypes.uint, 1)) + local_store = sbuf.index(lidx.valid(lidx<8)).store(UOp.const(1, dtypes.uint)) barrier = UOp(Ops.BARRIER, src=(local_store,)) if_barrier = UOp(Ops.IF, src=(gate, barrier)) @@ -187,7 +187,7 @@ class TestValidateOOB(unittest.TestCase): glbl0 = UOp.param(0, dtypes.int, (16,)) mask = UOp.param(0, dtypes.bool, (16,)) ridx = UOp.range(20, 0) - ld0 = UOp(Ops.LOAD, src=(glbl0.index(UOp.const(ridx, ridx<16&mask)))) + ld0 = UOp(Ops.LOAD, src=(glbl0.index(UOp.const(ridx<16&mask, ridx)))) to_uops_list([ld0]) if __name__ == "__main__": diff --git a/test/null/test_viz.py b/test/null/test_viz.py index 9facdd2a33..8d66e5f561 100644 --- a/test/null/test_viz.py +++ b/test/null/test_viz.py @@ -99,7 +99,7 @@ class TestViz(unittest.TestCase): assert x.arg <= 3 return x.replace(arg=x.arg+1) err_pm = PatternMatcher([(UPat.cvar("x"), count_3),]) - a = UOp.const(dtypes.int, 1) + a = UOp.const(1, dtypes.int) with save_viz() as viz: with self.assertRaises(AssertionError): exec_rewrite(a, [err_pm]) lst = viz.list_items() @@ -199,8 +199,8 @@ class TestViz(unittest.TestCase): self.assertEqual(ansistrip(a2["label"]), "CUSTOM\nx\nyzww\nw") def test_inf_loop(self): - a = UOp.const(dtypes.int, 3) - b = UOp.const(dtypes.int, 4) + a = UOp.const(3, dtypes.int) + b = UOp.const(4, dtypes.int) pm = PatternMatcher([ (UPat(Ops.CONST, arg=3, name="x"), lambda x: x.replace(arg=4)), (UPat(Ops.CONST, arg=4, name="x"), lambda x: x.replace(arg=3)), @@ -226,7 +226,7 @@ class TestViz(unittest.TestCase): def test_enter_calls_rewrite(self): pm = PatternMatcher([(UPat(Ops.CONST, arg=3, name="x"), lambda x: x.replace(arg=4))]) with save_viz() as viz: - inner = UOp.const(dtypes.int, 3) + inner = UOp.const(3, dtypes.int) call = UOp(Ops.CALL, src=(UOp(Ops.SINK, src=(inner,)),)) func = UOp(Ops.FUNCTION, src=(UOp(Ops.TUPLE, src=(call,)),)) graph_rewrite(func, TrackedPatternMatcher(pm.patterns), enter_calls=True) @@ -236,8 +236,8 @@ class TestViz(unittest.TestCase): def test_const_node_visibility(self): with save_viz() as viz: a = UOp.variable("a", 0, 10, dtype=dtypes.int) - z = UOp.const(a.dtype, 0) - y = UOp.const(dtypes.float, math.pi) + z = UOp.const(0, a.dtype) + y = UOp.const(math.pi, dtypes.float) alu = a*z ret = exec_rewrite(sink:=UOp.sink(alu, y), [sym]) lst = viz.list_items() @@ -253,7 +253,7 @@ class TestViz(unittest.TestCase): def test_const_reshape_expand_folded(self): # CONST->EXPAND should be folded into the ALU node, not shown as separate EXPAND nodes - c = UOp.const(dtypes.float, 1.0, shape=(3,4)) # creates CONST->EXPAND chain + c = UOp.const(1.0, dtypes.float, shape=(3,4)) # creates CONST->EXPAND chain a = UOp.variable("a", 0.0, 10.0, dtypes.float) alu = a + c with save_viz() as viz: @@ -267,13 +267,13 @@ class TestViz(unittest.TestCase): def test_stack_movement_not_folded_unless_all_const(self): a = UOp.variable("a", 0, 10, dtype=dtypes.int) - c = UOp.const(dtypes.int, 1) + c = UOp.const(1, dtypes.int) stack = a.stack(c) reshaped = stack.reshape((1, 2)) graph = uop_to_json(VizData(), reshaped) self.assertFalse(graph[id(stack)]["exclude"]) - const_stack = c.stack(UOp.const(dtypes.int, 2)) + const_stack = c.stack(UOp.const(2, dtypes.int)) const_reshaped = const_stack.reshape((1, 2)) const_graph = uop_to_json(VizData(), const_reshaped) self.assertTrue(const_graph[id(const_stack)]["exclude"]) @@ -401,7 +401,7 @@ class TestVizIntegration(unittest.TestCase): with save_viz() as viz: def test(root): return graph_rewrite(root, sym) - test(c:=UOp.const(dtypes.int, 1)) + test(c:=UOp.const(1, dtypes.int)) test(c+1) ls = viz.list_items() self.assertEqual(len(ls), 1) @@ -414,7 +414,7 @@ class TestVizIntegration(unittest.TestCase): @track_rewrites() def test(root): return graph_rewrite(root, sym) - test(c:=UOp.const(dtypes.int, 1)) + test(c:=UOp.const(1, dtypes.int)) test(c+1) ls = viz.list_items() self.assertEqual(len(ls), 2) @@ -425,7 +425,7 @@ class TestVizIntegration(unittest.TestCase): with save_viz() as viz: def default_test(root): return graph_rewrite(root, sym) tracked_test = track_rewrites()(default_test) - c = UOp.const(dtypes.int, 1) + c = UOp.const(1, dtypes.int) default_test(c+1) # goes to the default group tracked_test(c) # all rewrites after this go inside the second group. default_test(c+2) diff --git a/test/unit/test_assign.py b/test/unit/test_assign.py index 9f3f2c5de2..f7f3e72bbe 100644 --- a/test/unit/test_assign.py +++ b/test/unit/test_assign.py @@ -661,7 +661,7 @@ class TestAssign(unittest.TestCase): def test_assign_deviceless_const(self): s = Tensor.empty(4, device="CPU:1", dtype=dtypes.float) - s.assign(Tensor(UOp.const(dtypes.float, 2.0))) + s.assign(Tensor(UOp.const(2.0, dtypes.float))) np.testing.assert_equal(s.numpy(), [2, 2, 2, 2]) def test_nested_after_contiguous_store(self): diff --git a/test/unit/test_dtype_weak.py b/test/unit/test_dtype_weak.py index 03b2fc8cda..5ae7cca6b0 100644 --- a/test/unit/test_dtype_weak.py +++ b/test/unit/test_dtype_weak.py @@ -12,18 +12,18 @@ from tinygrad.engine.jit import JitError class TestWeakPromotion(unittest.TestCase): def test_rand_requires_concrete(self): with self.assertRaises(ValueError): Tensor.rand(2, dtype=dtypes.weakfloat) - with self.assertRaises(ValueError): Tensor.const(dtypes.weakfloat, 1.0).rand_like() - with self.assertRaises(ValueError): Tensor.const(dtypes.weakfloat, 1.0).randn_like() + with self.assertRaises(ValueError): Tensor.const(1.0).rand_like() + with self.assertRaises(ValueError): Tensor.const(1.0).randn_like() def test_reduce_strips_weakness(self): for weak, value, strong in ((dtypes.weakint, 1, dtypes.default_int), (dtypes.weakfloat, 1.0, dtypes.default_float)): - t = Tensor.const(weak, value).expand(3) + t = Tensor.const(value, weak).expand(3) for out in (t.sum(), t.max(), t.prod(), t.cumsum(0), t.cummax(0)[0]): self.assertEqual(out.dtype, strong) - self.assertEqual((Tensor.const(dtypes.weakfloat, 1.0).expand(3).sum() + Tensor([1], dtype=dtypes.float16)).dtype, dtypes.float32) + self.assertEqual((Tensor.const(1.0).expand(3).sum() + Tensor([1], dtype=dtypes.float16)).dtype, dtypes.float32) def test_materialize_at_default_dtype(self): for weak, value, strong in ((dtypes.weakfloat, 0.5, dtypes.default_float),): - t = Tensor.const(weak, value) + t = Tensor.const(value, weak) self.assertEqual(t.dtype, weak) self.assertEqual(t.data().itemsize, strong.itemsize) self.assertEqual(t.numpy().dtype.itemsize, strong.itemsize) @@ -33,7 +33,7 @@ class TestWeakPromotion(unittest.TestCase): self.assertEqual(t.contiguous().dtype, weak) def test_assign_into_weak_commits(self): - t = Tensor.const(dtypes.weakfloat, 0.5) + t = Tensor.const(0.5) t.assign(Tensor(1.0, dtype=dtypes.default_float)) self.assertEqual((t.dtype, t.item()), (dtypes.default_float, 1.0)) @@ -54,7 +54,7 @@ class TestWeakPromotion(unittest.TestCase): self.assertEqual((y._uop.base.op, y.dtype, x.dtype), (Ops.CONST, dtypes.weakint, dtypes.int8)) x, y = Tensor([1], dtype=dtypes.int8)._broadcasted(0.5) self.assertEqual((y._uop.base.op, y.dtype, x.dtype), (Ops.CONST, dtypes.weakfloat, dtypes.weakfloat)) - x, y = Tensor.const(dtypes.weakint, 1).reshape(1)._broadcasted(Tensor([1.0], dtype=dtypes.float32)) + x, y = Tensor.const(1).reshape(1)._broadcasted(Tensor([1.0], dtype=dtypes.float32)) self.assertEqual((x._uop.base.op, x._uop.base.arg, x.dtype, x.shape, y.dtype), (Ops.CONST, 1, dtypes.weakfloat, (1,), dtypes.float32)) @@ -69,34 +69,34 @@ class TestWeakPromotion(unittest.TestCase): # and a bare weak const UOp is the same spelling as the python scalar: both lift to the same node x = UOp.variable("x", 0.0, 1.0, dtypes.float32) self.assertIsInstance((x + 2).src[1].arg, float) - self.assertIs(x + UOp.const(None, 2), x + 2) + self.assertIs(x + UOp.const(2), x + 2) def test_index_dtype_ignores_weakness(self): with Context(SPEC=2): - idx = UOp.const(None, 0).cast(dtypes.int32) - weak = UOp.const(None, 1.0).expand((1,)) + idx = UOp.const(0).cast(dtypes.int32) + weak = UOp.const(1.0).expand((1,)) self.assertEqual(UOp(Ops.INDEX, dtypes.float32, (weak, idx)).dtype, dtypes.float32) with self.assertRaisesRegex(RuntimeError, "bad dtype"): UOp(Ops.INDEX, dtypes.int32, (weak, idx)) def test_store_weak_value_uses_destination_dtype(self): with Context(DEFAULT_FLOAT=dtypes.float16): - dst = UOp.param(0, dtypes.bfloat16, (1,)).index(UOp.const(None, 0).cast(dtypes.int32)) - gate = UOp.const(None, True) - out = graph_rewrite(dst.store(UOp.const(None, 5.0), gate), pm_lower_index_dtype, ctx={}) + dst = UOp.param(0, dtypes.bfloat16, (1,)).index(UOp.const(0).cast(dtypes.int32)) + gate = UOp.const(True) + out = graph_rewrite(dst.store(UOp.const(5.0), gate), pm_lower_index_dtype, ctx={}) # a bare weak CONST commits directly: the pass runs without symbolic, so a CAST here would survive it - self.assertEqual((out.src[1], out.src[2]), (UOp.const(dtypes.bfloat16, 5.0), gate)) + self.assertEqual((out.src[1], out.src[2]), (UOp.const(5.0, dtypes.bfloat16), gate)) def test_weak_srcs_commit_only_at_a_concrete_lub(self): - weak_lub = UOp(Ops.ADD, src=(UOp.const(None, 1), UOp.const(None, 1.0))) + weak_lub = UOp(Ops.ADD, src=(UOp.const(1), UOp.const(1.0))) self.assertIs(graph_rewrite(weak_lub, pm_lower_index_dtype, ctx={}), weak_lub) - concrete = UOp.const(None, 2.0).cast(dtypes.float16) - where = graph_rewrite(UOp(Ops.WHERE, src=(UOp.const(None, True), concrete, UOp.const(None, 1.0))), pm_lower_index_dtype, ctx={}) + concrete = UOp.const(2.0).cast(dtypes.float16) + where = graph_rewrite(UOp(Ops.WHERE, src=(UOp.const(True), concrete, UOp.const(1.0))), pm_lower_index_dtype, ctx={}) self.assertEqual(tuple(x.dtype for x in where.src), (dtypes.bool, dtypes.float16, dtypes.float16)) def test_weak_shift_lhs_commits_the_node(self): # a shift derives its lhs's dtype, so committing the lhs restates the root (WGSL's packed store writes `mask << shift_am`) - shl = graph_rewrite(UOp.const(None, 0xFFFF) << UOp.variable("x", 0, 16, dtypes.uint), symbolic_simple+pm_commit_weak) - self.assertEqual((shl.dtype, shl.src[0]), (dtypes.uint, UOp.const(dtypes.uint, 0xFFFF))) + shl = graph_rewrite(UOp.const(0xFFFF) << UOp.variable("x", 0, 16, dtypes.uint), symbolic_simple+pm_commit_weak) + self.assertEqual((shl.dtype, shl.src[0]), (dtypes.uint, UOp.const(0xFFFF, dtypes.uint))) @unittest.expectedFailure # TODO: a weak const defers to its consumer (JAX): these dtypes change once python scalars are weak consts def test_changed_rows(self): @@ -129,17 +129,17 @@ class TestWeakPromotion(unittest.TestCase): def test_weak_int_binop(self): v = UOp.variable("i", 0, 10, dtypes.weakint) self.assertEqual((v << 1).dtype, dtypes.weakint) - self.assertEqual(dtype_from_uop(Ops.SHL, (UOp.const(dtypes.int8, 1), UOp.const(dtypes.uint32, 1)), None), dtypes.int8) - self.assertEqual(UOp.const(dtypes.weakint, 1).alu(Ops.SHL, UOp.const(dtypes.uint, 1)).dtype, dtypes.weakint) + self.assertEqual(dtype_from_uop(Ops.SHL, (UOp.const(1, dtypes.int8), UOp.const(1, dtypes.uint32)), None), dtypes.int8) + self.assertEqual(UOp.const(1).alu(Ops.SHL, UOp.const(1, dtypes.uint)).dtype, dtypes.weakint) self.assertEqual((v & 3).dtype, dtypes.weakint) - with self.assertRaises(RuntimeError): Tensor.const(dtypes.weakfloat, 1.0) << Tensor.const(dtypes.weakfloat, 1.0) - with self.assertRaises(RuntimeError): UOp.const(dtypes.int32, 1).alu(Ops.SHL, UOp.const(dtypes.float64, 1)) + with self.assertRaises(RuntimeError): Tensor.const(1.0) << Tensor.const(1.0) + with self.assertRaises(RuntimeError): UOp.const(1, dtypes.int32).alu(Ops.SHL, UOp.const(1, dtypes.float64)) for op in (Ops.SHL, Ops.SHR): with self.assertRaises(RuntimeError): - UOp.const(dtypes.float32, 1).alu(op, UOp.const(dtypes.int32, 1)) + UOp.const(1, dtypes.float32).alu(op, UOp.const(1, dtypes.int32)) # float bitwise builds, the spec rejects it with Context(SPEC=1): - f32, wf = UOp.const(dtypes.float32, 1.0), UOp.const(dtypes.weakfloat, 1.0) + f32, wf = UOp.const(1.0, dtypes.float32), UOp.const(1.0) for bad in (f32.alu(Ops.AND, f32), UOp(Ops.AND, dtypes.float32, (f32, f32)), UOp(Ops.AND, dtypes.int32, (wf, wf))): with self.assertRaises(RuntimeError): type_verify([bad], spec_shared) @@ -178,7 +178,7 @@ class TestWeakPromotion(unittest.TestCase): class TestWeakStorageBoundary(unittest.TestCase): # weak has no storage: a weak assignment source casts when it defers to the destination, everything else raises def test_weak_source(self): - w05 = Tensor.const(dtypes.weakfloat, 0.5).reshape(1) + w05 = Tensor.const(0.5).reshape(1) dst = Tensor.zeros(2, dtype=dtypes.int8, device="CPU").contiguous().realize() with self.assertRaises(RuntimeError): dst.assign(w05.expand(2)) # weakfloat into int does not defer with self.assertRaises(RuntimeError): dst[0:1] = w05 @@ -199,7 +199,7 @@ class TestWeakMaterializationEntries(unittest.TestCase): def test_reads_commit_storage_raises(self): for weak, value, strong in ((dtypes.weakfloat, 0.5, dtypes.default_float),): def weak_val(): - return Tensor([True], device="CPU").where(Tensor.const(weak, value), Tensor.const(weak, value)) + return Tensor([True], device="CPU").where(Tensor.const(value, weak), Tensor.const(value, weak)) self.assertEqual(weak_val().dtype, weak) self.assertEqual(weak_val().to("CPU").dtype, weak) self.assertEqual(weak_val().data().format, strong.fmt) @@ -215,7 +215,7 @@ class TestWeakMaterializationEntries(unittest.TestCase): def test_weak_is_virtual(self): # NOTE: int64 lub uint64 is weakfloat, so this is device-ful weak from promotion, never from a cast to weak devful = Tensor([1], dtype=dtypes.int64, device="CPU") + Tensor([1], dtype=dtypes.uint64, device="CPU") - for t in (Tensor.const(dtypes.weakfloat, 0.5), devful): + for t in (Tensor.const(0.5), devful): self.assertTrue(t.uop.is_virtual) # realize is a no-op, so a weak input can never become the real buffer TinyJit needs with self.assertRaises(JitError): TinyJit(lambda x: (x+1).realize())(t) @@ -226,7 +226,7 @@ class TestWeakMaterializationEntries(unittest.TestCase): def test_empty_reads_commit(self): for weak, strong in ((dtypes.weakfloat, dtypes.default_float),): - empty = Tensor.const(weak, 0).reshape(1).shrink(((0, 0),)) + empty = Tensor.const(0, weak).reshape(1).shrink(((0, 0),)) self.assertEqual(empty.data().format, strong.fmt) self.assertEqual(empty.numpy().dtype.itemsize, strong.itemsize) self.assertEqual(empty.tolist(), []) diff --git a/test/unit/test_function.py b/test/unit/test_function.py index ca7014ea4c..6e4ce4d0d4 100644 --- a/test/unit/test_function.py +++ b/test/unit/test_function.py @@ -593,7 +593,7 @@ class TestFunctionTuple(unittest.TestCase): state = Tensor([10., 20., 30., 40.], device="CPU").contiguous().realize() @function(precompile=True, allow_implicit=True) def f(a:Tensor): - after = state.uop.after(state.uop.shrink(((0, 2),)).store(UOp.const(dtypes.float32, Invalid, shape=(2,)))) + after = state.uop.after(state.uop.shrink(((0, 2),)).store(UOp.const(Invalid, dtypes.float32, shape=(2,)))) return Tensor(after).contiguous() + a out = f(Tensor([1., 1., 1., 1.], device="CPU").contiguous().realize()) np.testing.assert_allclose(out.numpy(), [11., 21., 31., 41.]) diff --git a/test/unit/test_invalid_tensor.py b/test/unit/test_invalid_tensor.py index 8dfc9922f4..a23c0df55a 100644 --- a/test/unit/test_invalid_tensor.py +++ b/test/unit/test_invalid_tensor.py @@ -134,8 +134,8 @@ class TestInvalidTensor(unittest.TestCase): self._invalid_test_helper(out, [1.0, 2.0, None, None]) def test_uop_where_keeps_invalid_bare(self): - cond = UOp.const(dtypes.weakint, 0) < UOp.const(dtypes.weakint, 1) - idx = UOp(Ops.STACK, src=tuple(UOp.const(dtypes.weakint, x) for x in range(3))) + cond = UOp.const(0) < UOp.const(1) + idx = UOp(Ops.STACK, src=tuple(UOp.const(x) for x in range(3))) out = cond.where(idx, UOp.invalid()) self.assertIs(cond.op, Ops.CMPLT) self.assertIs(idx.op, Ops.STACK) diff --git a/test/unit/test_jit.py b/test/unit/test_jit.py index 7735776fec..8aaee9eec2 100644 --- a/test/unit/test_jit.py +++ b/test/unit/test_jit.py @@ -351,13 +351,13 @@ class TestJit(unittest.TestCase): @TinyJit def f(x:Tensor) -> Tensor: return (x + 1).realize() with self.assertRaises(JitError): - f(Tensor(UOp.const(dtypes.float, 2.0))).item() + f(Tensor(UOp.const(2.0, dtypes.float))).item() def test_jit_deviceless_compute_input(self): @TinyJit def f(x:Tensor) -> Tensor: return (x + 1).realize() with self.assertRaises(JitError): - f(Tensor(UOp.const(dtypes.float, 2.0) + UOp.const(dtypes.float, 1.0))).item() + f(Tensor(UOp.const(2.0, dtypes.float) + UOp.const(1.0, dtypes.float))).item() def test_jit_init_empty_alt(self): @TinyJit diff --git a/test/unit/test_metal_graph.py b/test/unit/test_metal_graph.py index 6f2951cfee..a85616ff7f 100644 --- a/test/unit/test_metal_graph.py +++ b/test/unit/test_metal_graph.py @@ -17,7 +17,7 @@ class TestMetalGraph(unittest.TestCase): buf.op = Ops.SLICE src = MagicMock() src.dtype = dtypes.uint8 - buf.src = (src, UOp.const(dtypes.weakint, offset)) + buf.src = (src, UOp.const(offset)) buf.dtype = dtypes.uint8 else: buf.op = Ops.BUFFER diff --git a/test/unit/test_multitensor.py b/test/unit/test_multitensor.py index 60f4e6701c..3a2cccac9a 100644 --- a/test/unit/test_multitensor.py +++ b/test/unit/test_multitensor.py @@ -60,8 +60,8 @@ class TestMultiTensor(unittest.TestCase): def test_shard_elementwise(self): self._test_shard_op(lambda t:(t+t).reshape(2, 2), [[2.,2.],[2.,2.]]) def test_alu_deviceless_const(self): s = Tensor([1.0, 2, 3, 4]).shard((f"{Device.DEFAULT}:0", f"{Device.DEFAULT}:1"), axis=0) - np.testing.assert_equal((s + Tensor(UOp.const(dtypes.float, 1.0))).numpy(), [2, 3, 4, 5]) - np.testing.assert_equal((s + Tensor(UOp.const(dtypes.float, 1.0)).reshape((1,)).expand((4,))).numpy(), [2, 3, 4, 5]) + np.testing.assert_equal((s + Tensor(UOp.const(1.0, dtypes.float))).numpy(), [2, 3, 4, 5]) + np.testing.assert_equal((s + Tensor(UOp.const(1.0, dtypes.float)).reshape((1,)).expand((4,))).numpy(), [2, 3, 4, 5]) def test_add_rank_expand_shard(self): # a sharded src keeps its own rank under implicit broadcast, its shard axis right-aligns into the output diff --git a/test/unit/test_tensor_data.py b/test/unit/test_tensor_data.py index d8cd0a3623..8874e93efe 100644 --- a/test/unit/test_tensor_data.py +++ b/test/unit/test_tensor_data.py @@ -66,8 +66,8 @@ class TestTensorData(unittest.TestCase): assert dat.shape == () def test_const_dtype_for_uop(self): - self.assertEqual(Tensor.const(dtypes.int8, UOp.const(dtypes.float32, 1.0)).dtype, dtypes.int8) - self.assertEqual(Tensor.const(dtypes.int32, UOp.variable("x", 1, 10).bind(5)).item(), 5) + self.assertEqual(Tensor.const(UOp.const(1.0, dtypes.float32), dtypes.int8).dtype, dtypes.int8) + self.assertEqual(Tensor.const(UOp.variable("x", 1, 10).bind(5), dtypes.int32).item(), 5) def test_data_float32(self): a = Tensor([[1,2.5],[3,4]], dtype=dtypes.float32) diff --git a/tinygrad/callify.py b/tinygrad/callify.py index 2c9345af9b..5ff272a4ec 100644 --- a/tinygrad/callify.py +++ b/tinygrad/callify.py @@ -61,7 +61,7 @@ def _make_buffer_view(src:UOp) -> UOp|None: buf = buf.src[0] if byte_offset % buf.dtype.itemsize != 0: return None offset = byte_offset // buf.dtype.itemsize - return UOp(Ops.SLICE, src.dtype, (buf, UOp.const(None, offset)), src.numel()) + return UOp(Ops.SLICE, src.dtype, (buf, UOp.const(offset)), src.numel()) def contiguous_mops_to_view(c:UOp, src:UOp): """MOPS(BUFFER) → SLICE when movement ops collapse to a contiguous range.""" diff --git a/tinygrad/codegen/__init__.py b/tinygrad/codegen/__init__.py index 230897630a..fffce829fa 100644 --- a/tinygrad/codegen/__init__.py +++ b/tinygrad/codegen/__init__.py @@ -84,7 +84,7 @@ def expand_wmma(ctx:dict[int, int], u:UOp): expander2 = PatternMatcher([ (UPat(Ops.REDUCE, name="r"), expand_reduce), (UPat(Ops.RANGE, name="r"), - lambda ctx, r: UOp.const(r.dtype, tuple(range(r.vmax+1))) \ + lambda ctx, r: UOp.const(tuple(range(r.vmax+1)), r.dtype) \ .reshape(tuple([r.vmax+1 if i == ctx[r.arg[0]] else 1 for i in range(len(ctx))])) if r.arg[0] in ctx else None), (UPat(Ops.WMMA, name="u"), expand_wmma), ])+pm_flatten_range+mop_cleanup @@ -126,7 +126,7 @@ def do_devectorize(b:UOp): if not all(x.shape == b.shape or x.base.arg is Invalid for x in b.src): return None src = [] for idx in itertools.product(*[range(x) for x in b.shape]): - idx_c = [UOp.const(None, i) for i in idx] + idx_c = [UOp.const(i) for i in idx] src.append(b.replace(dtype=None, src=tuple(x.base if x.base.arg is Invalid else x.index(*idx_c) for x in b.src))) return UOp.stack(*src).reshape(b.shape) if b.op is not Ops.STORE else UOp.group(*src) diff --git a/tinygrad/codegen/decomp/dtype.py b/tinygrad/codegen/decomp/dtype.py index 25bddacbfe..df9a24e47b 100644 --- a/tinygrad/codegen/decomp/dtype.py +++ b/tinygrad/codegen/decomp/dtype.py @@ -18,7 +18,7 @@ def reindex(idx:UOp, off:int, mul=2) -> UOp: # 4.3.1 is the relevant section in TAOCP def l2i(op: Ops, dt: DType, *uops:UOp): - zero = UOp.const(dt, 0) + zero = UOp.const(0, dt) if len(uops) == 2: a0, a1 = uops elif len(uops) == 3: a0, a1, b0 = uops # a shift's count is a single word elif len(uops) == 4: a0, a1, b0, b1 = uops @@ -57,10 +57,10 @@ def l2i(op: Ops, dt: DType, *uops:UOp): ua0, ua1, ub0, ub1 = a0.bitcast(dtypes.uint), a1.bitcast(dtypes.uint), b0.bitcast(dtypes.uint), b1.bitcast(dtypes.uint) a0, a1 = (a_neg:=a1 < zero).where((n:=l2i(Ops.NEG, dtypes.uint, ua0, ua1))[0], ua0), a_neg.where(n[1], ua1) b0, b1 = (b_neg:=b1 < zero).where((n:=l2i(Ops.NEG, dtypes.uint, ub0, ub1))[0], ub0), b_neg.where(n[1], ub1) - q, r = (z:=UOp.const(dtypes.uint, 0), z), (z, z) + q, r = (z:=UOp.const(0, dtypes.uint), z), (z, z) for i in range(63, -1, -1): - r = l2i(Ops.SHL, dtypes.uint, *r, UOp.const(dtypes.uint, 1), z) - r = (r[0] | l2i(Ops.SHR, dtypes.uint, a0, a1, UOp.const(dtypes.uint, i), z)[0] & 1), r[1] + r = l2i(Ops.SHL, dtypes.uint, *r, UOp.const(1, dtypes.uint), z) + r = (r[0] | l2i(Ops.SHR, dtypes.uint, a0, a1, UOp.const(i, dtypes.uint), z)[0] & 1), r[1] cond = l2i(Ops.CMPLT, dtypes.uint, *r, b0, b1).logical_not() diff = l2i(Ops.SUB, dtypes.uint, *r, b0, b1) q = ((q[0] | shl(cond.cast(dtypes.uint), i % 32), q[1]) if i < 32 else (q[0], q[1] | shl(cond.cast(dtypes.uint), i % 32))) @@ -158,7 +158,7 @@ pm_long_decomp = PatternMatcher([ (UPat(Ops.LOAD, tuple(l2i_dt.keys()), src=(UPat.var('idx'),), name='x'), lambda x,idx: x.replace(dtype=l2i_dt[x.dtype], src=(reindex(idx, x.tag[0]).replace(dtype=l2i_dt[x.dtype], tag=None),), tag=None) if x.tag is not None else None), (UPat(Ops.CONST, tag={(w, dt) for w in (0, 1) for dt in l2i_dt.values()}, name='x'), lambda x: - UOp.const(x.tag[1], truncate[x.tag[1]]((x.arg >> 32) if x.tag[0] == 1 else (x.arg & 0xFFFFFFFF)))) + UOp.const(truncate[x.tag[1]]((x.arg >> 32) if x.tag[0] == 1 else (x.arg & 0xFFFFFFFF)), x.tag[1])) ]) # float decomposition patterns - ctx is (fr, to) tuple diff --git a/tinygrad/codegen/decomp/op.py b/tinygrad/codegen/decomp/op.py index 8848ccf8ff..6efe4147d1 100644 --- a/tinygrad/codegen/decomp/op.py +++ b/tinygrad/codegen/decomp/op.py @@ -129,5 +129,5 @@ def get_late_rewrite_patterns(ops:tuple[Ops, ...], disable_fast_idiv:bool) -> Pa # some backends emit FDIV for RECIP, in that case: a*(1/b) -> a/b if Ops.FDIV in ops: pat += [(UPat.var("x").reciprocal(), lambda x: x.const_like(1).alu(Ops.FDIV, x))] - pat += [(UPat.var("a", dtypes.floats) * UPat(Ops.FDIV, dtypes.floats, src=(UPat.const(None, 1), UPat.var("b"))), lambda a,b: a.alu(Ops.FDIV, b))] + pat += [(UPat.var("a", dtypes.floats) * UPat(Ops.FDIV, dtypes.floats, src=(UPat.const(1), UPat.var("b"))), lambda a,b: a.alu(Ops.FDIV, b))] return PatternMatcher(pat) diff --git a/tinygrad/codegen/decomp/transcendental.py b/tinygrad/codegen/decomp/transcendental.py index 10ee275c1c..23b1723503 100644 --- a/tinygrad/codegen/decomp/transcendental.py +++ b/tinygrad/codegen/decomp/transcendental.py @@ -93,7 +93,7 @@ def payne_hanek_reduction(d:UOp) -> tuple[UOp, UOp]: def _shl_lazy(x:UOp, y:UOp): return (x.cast(dtypes.uint64) * pow2if(y, d.dtype).cast(dtypes.uint64)).cast(dtypes.uint32) def _shr_lazy(x:UOp, y:UOp): return (x.cast(dtypes.uint64) // pow2if(y, d.dtype).cast(dtypes.uint64)).cast(dtypes.uint32) - a = [_take(UOp.const(dtypes.uint32, 0), i) for i in range(4)] + a = [_take(UOp.const(0, dtypes.uint32), i) for i in range(4)] # (two_over_pi_f[Int(i) + n] << e) | (two_over_pi_f[Int(i) + n+1] >> (nbits - e)) # Note: e >= 1 for all numbers d >= 1.0. assume e != 0 hi = _shl_lazy(a[0], e) | _shr_lazy(a[1], offset) diff --git a/tinygrad/codegen/late/coalesce.py b/tinygrad/codegen/late/coalesce.py index 077b4a734b..436032e35f 100644 --- a/tinygrad/codegen/late/coalesce.py +++ b/tinygrad/codegen/late/coalesce.py @@ -141,12 +141,12 @@ def memory_coalescing(sink:UOp, ctx:Renderer) -> UOp: grouped_offsets = [[x for _,x in group] for _,group in itertools.groupby(enumerate(sorted(offsets.keys())), lambda x: x[1]-x[0])] for full_grp in grouped_offsets: while len(full_grp): - offset = (base+full_grp[0]) if isinstance(base, UOp) else UOp.const(None, full_grp[0]) + offset = (base+full_grp[0]) if isinstance(base, UOp) else UOp.const(full_grp[0]) length = [l for l in lengths if l <= len(full_grp) and (not must_divide or offset.divides(l) is not None)][0] grp = full_grp[:length] # NOTE: we apply the valid again after we determine the length offset = offset.valid(valid) if valid is not None else offset - idx = UOp(Ops.SHRINK, src=(buf, offset, UOp.const(None, len(grp)))) if len(grp) > 1 else buf.index(offset) + idx = UOp(Ops.SHRINK, src=(buf, offset, UOp.const(len(grp)))) if len(grp) > 1 else buf.index(offset) if op == Ops.STORE: datas = [] for i,g in enumerate(grp): diff --git a/tinygrad/codegen/late/regalloc.py b/tinygrad/codegen/late/regalloc.py index e3a456a439..85cf7cfb1b 100644 --- a/tinygrad/codegen/late/regalloc.py +++ b/tinygrad/codegen/late/regalloc.py @@ -52,7 +52,7 @@ class LinearScanRegallocContext: # the value of a BUFFER is its 64bit address, XMM registers need 16 bytes sz = 16 if v.cons[0].size == 16 else (8 if self.vdef(v).op is Ops.BUFFER else self.vdef(v).dtype.itemsize) offset = self.stack_size + (sz - self.stack_size % sz) % sz - self.spills[v] = UOp.const(dtypes.int32, offset) + self.spills[v] = UOp.const(offset, dtypes.int32) self.stack_size = offset + sz r = alloc(cons if cons is not None else v.cons, i) self.insert_before.setdefault(i, []).append((v, r)) @@ -84,7 +84,7 @@ class LinearScanRegallocContext: # allocate stack array if u.op is Ops.BUFFER: - self.locals[u] = UOp.const(dtypes.int32, self.stack_size) + self.locals[u] = UOp.const(self.stack_size, dtypes.int32) self.stack_size += u.max_numel() * u.dtype.itemsize # loop prologue, avoid loading inside the loop diff --git a/tinygrad/codegen/opt/postrange.py b/tinygrad/codegen/opt/postrange.py index dd4e5f03e7..1cf597b0f7 100644 --- a/tinygrad/codegen/opt/postrange.py +++ b/tinygrad/codegen/opt/postrange.py @@ -198,7 +198,7 @@ class Scheduler: for b in self.bufs: if rng in (i:=b.src[1].get_idx()).backward_slice_with_self: nb = b.replace(src=(b.src[0], i.valid(valid&b.src[1].get_valid()))) - replaces[b] = nb if b in store_targets else valid.where(nb, UOp.const(b.dtype, Invalid)) + replaces[b] = nb if b in store_targets else valid.where(nb, UOp.const(Invalid, b.dtype)) self.ast = self.ast.substitute(replaces, f"padto {rng.arg[:-1]} {opt.arg}") elif opt.op is OptOps.SWAP: try: @@ -302,7 +302,7 @@ class Scheduler: # TODO: remove tc_upcast_axes from the arg # do the reduce_axes always disappear? i think they don't # they need to be moved into the WMMA srcs - tc_uop = UOp.wmma(srcs[0], srcs[1], UOp.const(tc.dtype_out, (0.0,)*tc.elements_per_thread[2]), + tc_uop = UOp.wmma(srcs[0], srcs[1], UOp.const((0.0,)*tc.elements_per_thread[2], tc.dtype_out), tc.dims, self.ren.target.device, tc.threads, tc_upcast_axes=tc_upcast_axes) # preserve extra reduces diff --git a/tinygrad/llm/gguf.py b/tinygrad/llm/gguf.py index 6c486e29fc..49772dabf1 100644 --- a/tinygrad/llm/gguf.py +++ b/tinygrad/llm/gguf.py @@ -37,7 +37,7 @@ def ggml_data_to_tensor(t: Tensor, n: int, ggml_type: int) -> Tensor: def q_to_uint8(t: Tensor, b: int) -> Tensor: # TODO: rewrite with arange? - shift_tensor, bitmask = Tensor.const(t.dtype, tuple(2**(i*b) for i in range(8//b))), 0xff >> (8 - b) + shift_tensor, bitmask = Tensor.const(tuple(2**(i*b) for i in range(8//b)), t.dtype), 0xff >> (8 - b) return t.unsqueeze(-1).div(shift_tensor, rounding_mode="trunc").bitwise_and(bitmask).transpose(-1, -2).flatten(-2) if (nelements_nbytes := _GGML_QUANT.get(ggml_type)) is not None: @@ -74,7 +74,7 @@ def ggml_data_to_tensor(t: Tensor, n: int, ggml_type: int) -> Tensor: d = blocks[:, :2].bitcast(dtypes.float16).cast(dtypes.float32).reshape((-1, 1, 1, 1)) scale_words = blocks[:, 66:98].bitcast(dtypes.uint32) db = d * (scale_words.rshift(28).cast(dtypes.float32) + 0.5).reshape((-1, 8, 1, 1)) * 0.5 - sign_idx = scale_words.unsqueeze(-1).rshift(Tensor.const(dtypes.uint32, (0, 7, 14, 21))).bitwise_and(0x7F).reshape((-1, 32)).cast(dtypes.int32) + sign_idx = scale_words.unsqueeze(-1).rshift(Tensor.const((0, 7, 14, 21), dtypes.uint32)).bitwise_and(0x7F).reshape((-1, 32)).cast(dtypes.int32) even_signs = Tensor([i | (0x80 if i.bit_count() % 2 else 0) for i in range(128)], dtype=dtypes.uint8, device=t.device) signs = (q_to_uint8(even_signs[sign_idx].reshape((-1, 32, 1)), 1) == 0).where(1.0, -1.0).reshape((-1, 8, 4, 8)) grid = _ggml_iq_grid(t.device, _ggml.iq3xxs_grid, (256, 4))[blocks[:, 2:66]].reshape((-1, 8, 4, 8)) @@ -95,7 +95,7 @@ def ggml_data_to_tensor(t: Tensor, n: int, ggml_type: int) -> Tensor: return (db * _ggml_iq_grid(t.device, _ggml.iq2s_grid, (1024, 8))[q].reshape((-1, 16, 2, 8)) * signs).flatten(-3) if ggml_type == 23: d = blocks[:, :2].bitcast(dtypes.float16).cast(dtypes.float32).reshape((-1, 1, 1)) - scale_shifts = Tensor.const(dtypes.uint16, (0, 2, 4, 6, 8, 10, 12, 14)) + scale_shifts = Tensor.const((0, 2, 4, 6, 8, 10, 12, 14), dtypes.uint16) iq4_xs_lut = Tensor(list(_ggml.kvalues_iq4nl), dtype=dtypes.float32, device=t.device) scales_l = Tensor.stack((sl:=blocks[:, 4:8]).bitwise_and(0xF), sl.rshift(4), dim=2).reshape((-1, 8)) scales_h = blocks[:, 2:4].bitcast(dtypes.uint16).unsqueeze(-1).rshift(scale_shifts).bitwise_and(0x03).reshape((-1, 8)).cast(dtypes.uint8) diff --git a/tinygrad/mixin/creation.py b/tinygrad/mixin/creation.py index 61026ac5a8..31d7157d43 100644 --- a/tinygrad/mixin/creation.py +++ b/tinygrad/mixin/creation.py @@ -9,7 +9,7 @@ if TYPE_CHECKING: class CreationMixin(DTypeMixin, MovementMixin): @staticmethod - def const(dtype, b): raise NotImplementedError + def const(b, dtype=None): raise NotImplementedError def const_like(self, b: ConstType) -> Self: return self._wrap_uop(self._uop.const_like(b)) @@ -78,7 +78,7 @@ class CreationMixin(DTypeMixin, MovementMixin): from tinygrad.uop.ops import UOp new_shape = argfix(shape) dt = to_dtype(dtype) if dtype is not None else fill_value.dtype if isinstance(fill_value, UOp) else dtypes.from_py(fill_value) - val = cls.const(dt, fill_value) + val = cls.const(fill_value, dt) val = val.reshape((1,)*len(new_shape)).expand(new_shape) if not buffer: return val ret = val.empty_like(dt if dtype is not None else None, device) diff --git a/tinygrad/mixin/elementwise.py b/tinygrad/mixin/elementwise.py index 31a5d7b25a..3bc2e6347d 100644 --- a/tinygrad/mixin/elementwise.py +++ b/tinygrad/mixin/elementwise.py @@ -24,7 +24,7 @@ class ElementwiseMixin(CreationMixin): out_dtype = least_upper_dtype(x.dtype, y.dtype) # keep weak CONST weak, might lift weakint -> weakfloat def promote(t): - if t.dtype in dtypes.weaks and t._uop.base.op is Ops.CONST: return t._wrap_uop(t._uop.const(weak_dtype(out_dtype), t._uop.base.arg, t.shape)) + if t.dtype in dtypes.weaks and t._uop.base.op is Ops.CONST: return t._wrap_uop(t._uop.const(t._uop.base.arg, weak_dtype(out_dtype), t.shape)) return t.cast(out_dtype) return promote(x), promote(y) diff --git a/tinygrad/mixin/op.py b/tinygrad/mixin/op.py index 90f5d8f4fd..b79ef7274c 100644 --- a/tinygrad/mixin/op.py +++ b/tinygrad/mixin/op.py @@ -159,7 +159,7 @@ class OpMixin(ElementwiseMixin, ReduceMixin): per_dim.append((idx >= s) & (idx < e) & (((e-1-idx) if m['stride'] < 0 else (idx-s)) % st == 0)) vb = vb.flip(tuple(d for d, m in enumerate(mops) if m['stride'] < 0)) vb = vb.pad(tuple((m['boundary'][0], self.shape[d] - m['boundary'][1]) for d, m in enumerate(mops))) - return (type(self).uprod(*per_dim) if per_dim else type(self).const(dtypes.bool, True)).where(vb, self) + return (type(self).uprod(*per_dim) if per_dim else type(self).const(True)).where(vb, self) @classmethod def arange(cls, start, stop=None, step=1, dtype:DTypeLike|None=None) -> Self: @@ -1411,7 +1411,7 @@ class OpMixin(ElementwiseMixin, ReduceMixin): ret = self for dim in range(dims): ret = ret.transpose(0, dim) - ret = sum(type(self).const(ret.dtype, tuple(float(m[k]) for m in mat)).reshape((len(mat),)+(1,)*(ret.ndim-1)) * ret[k] + ret = sum(type(self).const(tuple(float(m[k]) for m in mat), ret.dtype).reshape((len(mat),)+(1,)*(ret.ndim-1)) * ret[k] for k in range(len(mat[0]))) assert not isinstance(ret, int), "sum over empty winograd matrix" ret = ret.transpose(0, dim) @@ -1898,7 +1898,7 @@ class OpMixin(ElementwiseMixin, ReduceMixin): # https://keccak.team/keccak_specs_summary.html def ctensor(l: Sequence[PyConst], dtype: DType = dtypes.uint64): - return type(self).const(dtype, tuple(l)) + return type(self).const(tuple(l), dtype) rot_offsets = [44, 43, 21, 14, 28, 20, 3, 45, 61, 1, 6, 25, 8, 18, 27, 36, 10, 15, 56, 62, 55, 39, 41, 2] rot_offsets_v0, rot_offsets_v1 = ctensor([0] + [1 << v for v in rot_offsets]), ctensor([1] + [1 << (64 - v) for v in rot_offsets]) @@ -1918,7 +1918,7 @@ class OpMixin(ElementwiseMixin, ReduceMixin): lbe = (data.shape[1] - 1) * 200 + rate - data_pad if data_pad == 1: mb = [(lbe, 0), (1, dsbyte ^ 0x80), (200 - rate, 0)] else: mb = [(lbe, 0), (1, dsbyte), (data_pad - 2, 0), (1, 0x80), (200 - rate, 0)] - pad_mask = type(self).cat(*(type(self).const(dtypes.uint8, v).expand(l) for l, v in mb if l > 0)).unsqueeze(0) + pad_mask = type(self).cat(*(type(self).const(v, dtypes.uint8).expand(l) for l, v in mb if l > 0)).unsqueeze(0) data = (data.flatten(1) ^ pad_mask).reshape(*data.shape[:2], 200).bitcast(dtypes.uint64) diff --git a/tinygrad/renderer/cstyle.py b/tinygrad/renderer/cstyle.py index 1b0287b7de..dfb18ff382 100644 --- a/tinygrad/renderer/cstyle.py +++ b/tinygrad/renderer/cstyle.py @@ -522,7 +522,7 @@ class HIPRenderer(CStyleLanguage): lambda x: x.replace(src=(x.src[0].bitcast(dtypes.uint64), x.src[1].bitcast(dtypes.uint64), x.src[2])) if x.src[0].max_numel() == 8 and x.src[0].dtype in dtypes.fp8_ocp else None), # bfloat16 constant casting - (UPat.cvar('x', dtypes.bfloat16), lambda x: cast_float_to_bf16(UOp.const(dtypes.float, x.arg))), + (UPat.cvar('x', dtypes.bfloat16), lambda x: cast_float_to_bf16(UOp.const(x.arg, dtypes.float))), ]) def asm(self, prg:UOp, lin:UOp) -> bytes: diff --git a/tinygrad/renderer/isa/x86.py b/tinygrad/renderer/isa/x86.py index 5f3597e18d..2436d8bcf4 100644 --- a/tinygrad/renderer/isa/x86.py +++ b/tinygrad/renderer/isa/x86.py @@ -166,7 +166,7 @@ def scratch_buffer(elem_dt:DType, count:int, slot:int) -> UOp: def gated_load(ctx, addr:UOp, alt:UOp, gate:UOp, x:UOp): local = scratch_buffer(addr.src[0].dtype.scalar(), x.max_numel(), next(ctx)) - local_idx = local.index(UOp.const(dtypes.int32, 0), dtype=dtypes.uint64) + local_idx = local.index(UOp.const(0, dtypes.int32), dtype=dtypes.uint64) # the selected address is a 64bit value, the AFTER orders the load after the scratch store and carries the element dtype for the encoder sel = gate.where(addr.replace(dtype=dtypes.uint64), local_idx) ptr = UOp(Ops.AFTER, addr.dtype, (sel, (local_idx if x.max_numel() == 1 else local).store(alt))) @@ -174,7 +174,7 @@ def gated_load(ctx, addr:UOp, alt:UOp, gate:UOp, x:UOp): def gated_store(addr:UOp, gate:UOp, val:UOp): local = scratch_buffer(addr.src[0].dtype.scalar(), val.max_numel(), -1) - sel = gate.where(addr.replace(dtype=dtypes.uint64), local.index(UOp.const(dtypes.int32, 0), dtype=dtypes.uint64)) + sel = gate.where(addr.replace(dtype=dtypes.uint64), local.index(UOp.const(0, dtypes.int32), dtype=dtypes.uint64)) return UOp(Ops.AFTER, addr.dtype, (sel,)).store(val) # legalize the new style graph for isel. NOTE: this runs after the spec is verified, some of these rewrites violate it @@ -195,7 +195,7 @@ pre_isel_matcher = PatternMatcher([ # if gate in scalar int cmove is not a comparison need to add one to set the flag # NOTE: the 0 is int so the bool gate zero-extends and compares as int (a byte compare renders different kernels) (UPat.var("m", dtypes.bool).where(UPat.var("a"), UPat.var("b")), - lambda m,a,b: m.ne(UOp.const(dtypes.int, 0)).where(a,b) if m.op not in GroupOp.Comparison else None), + lambda m,a,b: m.ne(UOp.const(0, dtypes.int)).where(a,b) if m.op not in GroupOp.Comparison else None), ]) # ***** X86 registers ***** @@ -224,7 +224,7 @@ def base(x:UOp, i:int) -> UOp: return s.src[0] if (s:=x.src[i]).op is Ops.INDEX def lane(x:UOp, i:int) -> int: return s.src[1].arg if (s:=x.src[i]).op is Ops.INDEX else 0 def to_int(dt:DType): return {dtypes.float16: dtypes.int16, dtypes.float32: dtypes.int32, dtypes.float64: dtypes.int64}[dt] def def_reg(dt:DType, reg:Register|None=None) -> UOp: return UOp(Ops.INS, dt, arg=X86Ops.DEFINE, tag=None if reg is None else (reg,)) -def imm(dt:DType, v:int) -> UOp: return UOp.const(dt, truncate[dt](v)).rtag() +def imm(dt:DType, v:int) -> UOp: return UOp.const(truncate[dt](v), dt).rtag() def to_imm(c:UOp) -> UOp|None: if c.op is not Ops.CONST: return None if c.dtype is dtypes.int64: return imm(dtypes.int32, c.arg) if not c.overflows(dtypes.int32) else None @@ -370,7 +370,7 @@ isel_matcher = PatternMatcher([ (UPat.cvar("x", dtypes.int64s), lambda x: x.ins(X86Ops.MOVABS, src=(imm(x.dtype, x.arg),)) if not x.tag else None), (UPat.cvar("x", dtypes.ints+(dtypes.bool,)), lambda x: x.ins(X86Ops.MOVi, src=(imm(x.dtype, x.arg),)) if not x.tag else None), (UPat.cvar("x", dtypes.floats), lambda x: - UOp.const(dt:=to_int(x.dtype), struct.unpack(dt.fmt, struct.pack(x.dtype.fmt, x.arg))[0]).bitcast(x.dtype) if not x.tag else None), + UOp.const(struct.unpack((dt:=to_int(x.dtype)).fmt, struct.pack(x.dtype.fmt, x.arg))[0], dt).bitcast(x.dtype) if not x.tag else None), # conditional moves that use masks NOTE: these currently assume a mask producing cmp exists (UPat.var("m").where(UPat.var("a", dtypes.int8s+dtypes.int16s+dtypes.int32s+(dtypes.int64,)), UPat.var("b")), lambda m,a,b: a.ins(X86Ops.VPBLENDVB, src=(b, a, m.replace(dtype=m.src[0].dtype))) if a.max_numel() > 1 else None), @@ -380,7 +380,7 @@ isel_matcher = PatternMatcher([ a.ins(X86Ops.VBLENDVPD, src=(b, a, m.replace(dtype=m.src[0].dtype)))), # in this case we have a mask producing comparison whose user expects a bool, so we convert to bool (UPat(GroupOp.Comparison, dtypes.bool, (UPat.var("y", (dtypes.float32, dtypes.float64)), UPat()), name="x"), lambda y,x: - UOp(Ops.AND, src=(x.replace(dtype=y.dtype).bitcast(dt:=to_int(y.dtype)), UOp.const(dt, 1))).f(Ops.NOOP, dtype=dtypes.bool)), + UOp(Ops.AND, src=(x.replace(dtype=y.dtype).bitcast(dt:=to_int(y.dtype)), UOp.const(1, dt))).f(Ops.NOOP, dtype=dtypes.bool)), # conditional moves that use flags (UPat(Ops.CMPLT, src=(UPat(dtype=dtypes.sints), UPat()), name="m").where(UPat.var("a"), UPat.var("b")), lambda m,a,b: a.ins(X86Ops.CMOVL, src=(b, a, cmp(m)))), diff --git a/tinygrad/renderer/llvmir.py b/tinygrad/renderer/llvmir.py index 81eb6453d0..6ae37820d8 100644 --- a/tinygrad/renderer/llvmir.py +++ b/tinygrad/renderer/llvmir.py @@ -274,7 +274,7 @@ exit: %packed = phi i32 [%packed_bf8, %do_bf8], [%packed_fp8, %do_fp8]\n %trunc src=(x.src[0].bitcast(dtypes.uint32), x.src[1].bitcast(dtypes.uint32), x.src[2])) if x.src[0].dtype == dtypes.int8 and x.src[0].max_numel() == 16 else None), (UPat(Ops.WMMA, name="x", dtype=dtypes.half), lambda x: UOp(Ops.STACK, src=tuple(x.replace( - src=(x.src[0], x.src[1], UOp(Ops.STACK, src=tuple(x.src[2].index(j//2) if j%2 == 0 else UOp.const(x.src[2].dtype, 0.0) + src=(x.src[0], x.src[1], UOp(Ops.STACK, src=tuple(x.src[2].index(j//2) if j%2 == 0 else UOp.const(0.0, x.src[2].dtype) for j in range(x.max_numel()*2)))), arg=(*x.arg[:4], None)).index(i*2) for i in range(x.max_numel()))) if x.max_numel() == 8 else None), diff --git a/tinygrad/renderer/nir.py b/tinygrad/renderer/nir.py index 542e579870..fbac922a10 100644 --- a/tinygrad/renderer/nir.py +++ b/tinygrad/renderer/nir.py @@ -122,7 +122,7 @@ class NIRRenderer(Renderer): extra_matcher = PatternMatcher([ # handle negative unsigned CONST - (UPat.cvar("x", dtypes.uints), lambda x: UOp.const(x.dtype, x.dtype.max+x.arg+1) if x.arg < 0 else None), + (UPat.cvar("x", dtypes.uints), lambda x: UOp.const(x.dtype.max+x.arg+1, x.dtype) if x.arg < 0 else None), # from ptx (UPat.var('x', dtype=dtypes.bool) uint8 diff --git a/tinygrad/renderer/wgsl.py b/tinygrad/renderer/wgsl.py index d8620aed45..71c969629a 100644 --- a/tinygrad/renderer/wgsl.py +++ b/tinygrad/renderer/wgsl.py @@ -6,7 +6,7 @@ from tinygrad.helpers import strip_parens def _mask(dt:DType): return 0xFF if dt.itemsize == 1 else 0xFFFF def sign_extend(val:UOp, sext_am:int): - return (UOp.where((val >> (sext_am - 1)) > 0, UOp.const(dtypes.uint32, 0xffffffff) << sext_am, UOp.const(dtypes.uint32, 0)) \ + return (UOp.where((val >> (sext_am - 1)) > 0, UOp.const(0xffffffff, dtypes.uint32) << sext_am, UOp.const(0, dtypes.uint32)) \ | val.bitcast(dtypes.uint32)).bitcast(dtypes.int) # store for char: buf[idx/4] <- (var << (idx%4)*8)) @@ -17,7 +17,7 @@ def packed_store(bidx:UOp, var:UOp, gate:UOp|None=None): if var.dtype == dtypes.bool: var = var.cast(dtypes.int32) new_v, wmask = (var & mask).cast(dtypes.uint32) << shift_am, ((mask << shift_am) ^ 0xFFFFFFFF).cast(dtypes.uint32) idx = UOp(Ops.INDEX, src=(bidx.src[0], div_idx)) - buf = UOp.load(idx, *((UOp.const(dtypes.uint32, 0), gate) if gate is not None else ()), dtype=dtypes.uint32) + buf = UOp.load(idx, *((UOp.const(0, dtypes.uint32), gate) if gate is not None else ()), dtype=dtypes.uint32) return UOp.store(idx, (buf & wmask) | new_v, *((gate,) if gate is not None else ())) # load for char: sign_extend(buf[idx/4] >> ((idx%4)*8)) diff --git a/tinygrad/runtime/ops_cpu.py b/tinygrad/runtime/ops_cpu.py index 173e907712..af0bf2b406 100644 --- a/tinygrad/runtime/ops_cpu.py +++ b/tinygrad/runtime/ops_cpu.py @@ -28,20 +28,20 @@ def wait_prog(): return (v:=UOp.param(0, dtypes.uint32, (1,), volatile=True).after(l:=UOp.loop(0))[0].load()).end(l, v < val.cast(dtypes.uint32)) def timestamp_prog(): - if WIN: val = UOp.const(dtypes.uint64, 0) + if WIN: val = UOp.const(0, dtypes.uint64) else: fn, ts = UOp.param(1, dtypes.uint64, (1,)), UOp.placeholder((2,), dtypes.uint64, slot=0, addrspace=AddrSpace.REG) - call = fn[0].load().call(UOp.const(dtypes.int, 6 if OSX else 1), ts[0], ret_dtype=dtypes.void) # clock_gettime(CLOCK_MONOTONIC, &ts) + call = fn[0].load().call(UOp.const(6 if OSX else 1, dtypes.int), ts[0], ret_dtype=dtypes.void) # clock_gettime(CLOCK_MONOTONIC, &ts) val = ts.after(call)[0].load() * 1_000_000_000 + ts.after(call)[1].load() return UOp.param(0, dtypes.uint64, (1,))[0].store(val) def quit_prog(): fn = UOp.param(0, dtypes.uint64, (1 if WIN else 3,)) - if WIN: return fn[0].load().call(UOp.const(dtypes.uint64, 0), ret_dtype=dtypes.void) # ExitThread(0) + if WIN: return fn[0].load().call(UOp.const(0, dtypes.uint64), ret_dtype=dtypes.void) # ExitThread(0) sem = UOp.param(1, dtypes.uint64, (1,)) close = fn[2].load().call(sem[0], ret_dtype=dtypes.void) # sem_close(sem) - return fn.after(close)[0].load().call(UOp.const(dtypes.uint64, 0), ret_dtype=dtypes.void) # pthread_exit(0) + return fn.after(close)[0].load().call(UOp.const(0, dtypes.uint64), ret_dtype=dtypes.void) # pthread_exit(0) def worker_prog(): ring = UOp.param(0, dtypes.uint64, (RING_SLOTS * CMD_SIZE,), volatile=True) @@ -56,7 +56,7 @@ def worker_prog(): return entry[0].call(*entry[1:], ret_dtype=dtypes.void).end(cur) def host_wait(ctx, dst:UOp, val:UOp) -> UOp: - return (cur:=dst.after(loop:=UOp.loop(next(ctx))).index(UOp.const(dtypes.int, 0)).load()).end(loop, cur < val) + return (cur:=dst.after(loop:=UOp.loop(next(ctx))).index(UOp.const(0, dtypes.int)).load()).end(loop, cur < val) pm_host_opsel = PatternMatcher([(UPat(Ops.INS, arg="wait", src=(UPat(name="dst"), UPat(name="val"))), host_wait)]) @@ -64,7 +64,7 @@ def encode_host_queue(q:UOp) -> UOp: # TODO: subset of hcq2 for now spins, (store,) = partition(graph_rewrite(q, pm_host_opsel, ctx=itertools.count(), walk=True, name="host opsel").src, lambda u: u.op is Ops.END) assert store.op is Ops.INS and store.arg == "store", f"host queue cannot encode {store.op} {store.arg}" - return store.src[0].after(*spins).index(UOp.const(dtypes.int, 0)).store(store.src[1]) + return store.src[0].after(*spins).index(UOp.const(0, dtypes.int)).store(store.src[1]) class CPUComputeQueue(HWQueue): def __init__(self, dev): diff --git a/tinygrad/runtime/support/hcq2.py b/tinygrad/runtime/support/hcq2.py index ce2f68ffc4..3d5fdf567f 100644 --- a/tinygrad/runtime/support/hcq2.py +++ b/tinygrad/runtime/support/hcq2.py @@ -42,7 +42,7 @@ def unwrap_mstack(u): return unwrap_mstack(u.src[0]) if u.op in {Ops.MSELECT, Ops.SLICE} else (u,) def make_patches(buf:UOp, patches:Sequence[tuple[sint, UOp]]) -> UOp: - offsets = UOp(Ops.STACK, dtypes.int, tuple(UOp.const(dtypes.int, off // buf.dtype.itemsize) for off,_ in patches)) + offsets = UOp(Ops.STACK, dtypes.int, tuple(UOp.const(off // buf.dtype.itemsize, dtypes.int) for off,_ in patches)) values = UOp(Ops.STACK, buf.dtype, tuple(val.cast(buf.dtype) for _,val in patches)) return buf.index(offsets).store(values) @@ -127,7 +127,7 @@ def _build_wait_cmds(dep_lanes:list[tuple[tuple, int, int]], devices:tuple[str, for (ddevs, dqueue, dtag), lanes in deps.items(): sig = UOp.mstack(*[make_signal(d if dl is None else ddevs[dl], queue=dqueue, sentinel=dl is None) for dl, d in zip(lanes, devices)]) val = UOp.mstack(*[make_signal_value(d if dl is None else ddevs[dl], queue=dqueue) for dl, d in zip(lanes, devices)]) - waits.append(UOp(Ops.INS, arg="wait", src=(sig, val.index(UOp.const(dtypes.int, 0)) + dtag))) + waits.append(UOp(Ops.INS, arg="wait", src=(sig, val.index(UOp.const(0, dtypes.int)) + dtag))) return waits, {dtag for _, _, dtag in deps} def _build_finalizers(batch:list[tuple[UOp, tuple[str, ...]]], batch_info:list[tuple[tuple[str, ...], str]], @@ -138,7 +138,7 @@ def _build_finalizers(batch:list[tuple[UOp, tuple[str, ...]]], batch_info:list[t for b in itertools.chain.from_iterable(_get_call_bufs_by_lane(call, devices)): for bd in to_tuple(b.device): dev_bufs[bd][id(b)] = b - zero, n, submits, bumps, waited = UOp.const(dtypes.int, 0), len(batch_info), [], [], set() + zero, n, submits, bumps, waited = UOp.const(0, dtypes.int), len(batch_info), [], [], set() for _, devgroup in itertools.groupby(sorted(dev_bufs), key=lambda d: d.split(":")[0]): devs = tuple(devgroup) @@ -283,7 +283,7 @@ def make_addr_table(call:UOp, gaddrs:list[UOp], name:str) -> tuple[dict[UOp, UOp slots = {g:i for i,g in enumerate(order)} table = UOp.placeholder((len(order),), dtypes.uint64, next(UOp.unique_num), device=call.arg.aux.device).rtag(name) - reads = {g: table.after(*g.src[0].src[1:] if g.src[0].op is Ops.AFTER else ()).index(UOp.const(dtypes.int, slots[bare[g]])).load() for g in gaddrs} + reads = {g: table.after(*g.src[0].src[1:] if g.src[0].op is Ops.AFTER else ()).index(UOp.const(slots[bare[g]], dtypes.int)).load() for g in gaddrs} return reads, (table.after(make_patches(table, [(i * table.dtype.itemsize, addr) for addr, i in slots.items()])),) if slots else () def make_blob_bufs(call:UOp, blobs:list[UOp]) -> tuple[dict[UOp, UOp], tuple[UOp, ...]]: @@ -333,7 +333,7 @@ pm_replace_params = PatternMatcher([ def resolve_getaddr_slice(bv:UOp, g:UOp) -> UOp: base = bv.src[0].after(*g.src[0].src[1:] if g.src[0].op is Ops.AFTER else ()) itemsize = bv.src[0].dtype.itemsize if bv.src[0].without_after.op in (Ops.BUFFER, Ops.SLICE, Ops.MSTACK, Ops.MSELECT) else bv.dtype.itemsize - return UOp(Ops.GETADDR, src=(base,), arg=g.arg) + UOp.const(dtypes.uint64, bv.src[1].arg * itemsize) + return UOp(Ops.GETADDR, src=(base,), arg=g.arg) + UOp.const(bv.src[1].arg * itemsize, dtypes.uint64) pm_early_simplify = PatternMatcher([ (UPat(Ops.GETADDR, src=(UPat.any(sl:=UPat(Ops.SLICE, name="bv"), sl.after(allow_any_len=True)),), name="g"), resolve_getaddr_slice), @@ -355,7 +355,7 @@ def pack_hcq_placeholders(call:UOp) -> UOp|None: sizes[b.tag] = offs[b] + b.max_numel() counts = collections.Counter(b.tag for b in bufs) bases = {b.tag:UOp.placeholder((sizes[b.tag],), b.dtype, next(UOp.unique_num), device=b.device).rtag(b.tag) for b in bufs if counts[b.tag] > 1} - subs = {b:UOp(Ops.SLICE, b.dtype, (bases[b.tag], UOp.const(None, offs.get(b, 0))), b.max_numel()) for b in bufs if b.tag in bases} + subs = {b:UOp(Ops.SLICE, b.dtype, (bases[b.tag], UOp.const(offs.get(b, 0))), b.max_numel()) for b in bufs if b.tag in bases} return call.replace(src=(call.src[0].substitute(subs, walk=True), *call.src[1:])) if subs else None pm_pack_placeholders = PatternMatcher([ (UPat(Ops.CALL, src=(UPat(Ops.CUSTOM_FUNCTION, arg="hcq"),), name="call", allow_any_len=True), pack_hcq_placeholders)]) @@ -428,7 +428,7 @@ def resolve_getaddr(buf:UOp, g:UOp) -> UOp: devs, b = g.arg, buf.buffer bufs = tuple(cast(Buffer, x.buffer) for x in buf.src) if buf.op is Ops.MSTACK else tuple(b.bufs if isinstance(b, MultiBuffer) else (b,)*len(devs)) assert len(bufs) == len(devs), f"can't resolve {len(bufs)} buffers on {len(devs)} devices" - addrs = tuple(UOp.const(dtypes.uint64, x.get_buf(d).va_addr) for x, d in zip(bufs, devs)) + addrs = tuple(UOp.const(x.get_buf(d).va_addr, dtypes.uint64) for x, d in zip(bufs, devs)) return addrs[0] if len(addrs) == 1 else UOp(Ops.STACK, src=addrs) pm_resolve_patches = PatternMatcher([ diff --git a/tinygrad/schedule/indexing.py b/tinygrad/schedule/indexing.py index b9ab1ce1f9..04a5e67c23 100644 --- a/tinygrad/schedule/indexing.py +++ b/tinygrad/schedule/indexing.py @@ -52,8 +52,8 @@ class IndexingContext: range_idx: Iterator[int] = field(default_factory=itertools.count) def new_range(self, s:sint, axistype:AxisType=AxisType.WEAK) -> UOp: if isinstance(s, UOp) and s.op is Ops.RANGE: return s - # if a range has a 1 src, it's the same as UOp.const(dtypes.weakint, 0) - return UOp.range(s, next(self.range_idx), axistype) if resolve(s!=1) else UOp.const(None, 0) + # if a range has a 1 src, it's the same as UOp.const(0) + return UOp.range(s, next(self.range_idx), axistype) if resolve(s!=1) else UOp.const(0) def broadcast_rngs(x:UOp, src:UOp, rngs:tuple[UOp, ...]) -> tuple[UOp, ...]: if x.op not in GroupOp.Broadcastable: return rngs @@ -100,8 +100,8 @@ def create_bufferize_and_index_based_on_ranges(ctx:IndexingContext, x:UOp): def convert_pad_to_where_to_keep_behavior_local(ctx:IndexingContext, x:UOp): if x not in ctx.range_map: return None bx = create_bufferize_and_index_based_on_ranges(ctx, x) - valid: UOp = UOp.const(None, True).uprod([r.get_valid() for r in ctx.range_map[x][0]]) - return valid.where(bx.src[0], UOp.const(x.dtype, 0)) + valid: UOp = UOp.const(True).uprod([r.get_valid() for r in ctx.range_map[x][0]]) + return valid.where(bx.src[0], UOp.const(0, x.dtype)) def convert_reduce_to_reduce_with_ranges(ctx:IndexingContext, x:UOp): if x.arg[1] == 0: return None @@ -148,7 +148,7 @@ def _apply_reshape(in_shape:tuple[sint,...], out_shape:tuple[sint, ...], urngs:U for s,src in list(zip(out_shape, urngs.src))[::-1]: axes_in.append(acc*src) acc *= s - combined_axes = UOp.const(None, 0).usum(axes_in) + combined_axes = UOp.const(0).usum(axes_in) axes_out:list[UOp] = [] for s in in_shape[::-1]: axes_out.append(combined_axes % s) @@ -248,7 +248,7 @@ def run_rangeify(tsink:UOp, debug:bool=False) -> tuple[UOp, IndexingContext]: # we compare the ranges without their valids if all_all_same or (PCONTIG and all_same(local_rngs)): # the new valid is the OR of all the children valids - minimum_valid = UOp.const(None, False).usum(valids) + minimum_valid = UOp.const(False).usum(valids) _out_rngs.append(graph_rewrite(local_rngs[0].valid(minimum_valid), symbolic, name="minimum_valid")) else: _out_rngs.append(rctx.new_range(x.shape[i])) diff --git a/tinygrad/schedule/memory.py b/tinygrad/schedule/memory.py index 3f163617d9..a16acb09f9 100644 --- a/tinygrad/schedule/memory.py +++ b/tinygrad/schedule/memory.py @@ -56,7 +56,7 @@ def memory_plan_rewrite(linear:UOp, held_bufs:set[UOp]|None=None) -> UOp: arenas = {key: UOp.new_buffer(key[0], sz, dtypes.int8) for key, sz in arena_sizes.items()} replace_map:dict[UOp, UOp] = {} for buf_uop, offset in offsets.items(): - replace_map[buf_uop] = UOp(Ops.SLICE, buf_uop.dtype, (arenas[_key(buf_uop)], UOp.const(None, offset)), buf_uop.max_numel()) + replace_map[buf_uop] = UOp(Ops.SLICE, buf_uop.dtype, (arenas[_key(buf_uop)], UOp.const(offset)), buf_uop.max_numel()) if DEBUG >= 1 and (omem:=sum(nbytes.values()) / 1e6) != (nmem:=sum(arena_sizes.values()) / 1e6): print(f"memory reduced from {omem:.2f} MB -> {nmem:.2f} MB, {len(first_appearance)} -> {len(arenas)} bufs") diff --git a/tinygrad/schedule/rangeify.py b/tinygrad/schedule/rangeify.py index 26116e2af1..ff15b77b24 100644 --- a/tinygrad/schedule/rangeify.py +++ b/tinygrad/schedule/rangeify.py @@ -445,7 +445,7 @@ class LocalAddBufferContext: opts:tuple|None = None def debuf(ctx:LocalAddBufferContext, buf:UOp): - param = UOp(Ops.PARAM, src=(UOp.const(dtypes.int, prod(buf.max_shape)),), + param = UOp(Ops.PARAM, src=(UOp.const(prod(buf.max_shape), dtypes.int),), arg=ParamArg(ctx.dg, buf.dtype, addrspace=buf.addrspace, device=buf.device)) ret = param.reshape(buf.max_shape) # if the buffer has symbolic shape, shrink the max-sized view to the actual shape diff --git a/tinygrad/tensor.py b/tinygrad/tensor.py index 05c06549e1..71b4fa8f9d 100644 --- a/tinygrad/tensor.py +++ b/tinygrad/tensor.py @@ -71,11 +71,11 @@ class Tensor(RandMixin): # create a UOp from the different types of inputs if data is None: - data = UOp.const(_dtype, 0.0) + data = UOp.const(0.0, _dtype) elif isinstance(data, get_args(ConstType)): - data = UOp.const(_dtype, data) + data = UOp.const(data, _dtype) elif is_numpy_ndarray(data) and data.shape == (): - data = UOp.const(_dtype or _from_np_dtype(data.dtype), data.item()) + data = UOp.const(data.item(), _dtype or _from_np_dtype(data.dtype)) elif not isinstance(data, UOp): if _dtype in dtypes.weaks: raise RuntimeError(f"cannot create storage for weak dtype {_dtype}") if isinstance(data, bytes): @@ -120,7 +120,7 @@ class Tensor(RandMixin): @classmethod def _wrap_uop(cls, u:UOp) -> Tensor: return cls(u) @staticmethod - def const(dtype:DType, b:ConstLike) -> Tensor: return Tensor(UOp.const(dtype, b)) + def const(b:ConstLike, dtype:DType|None=None) -> Tensor: return Tensor(UOp.const(b, dtype)) def is_param_(self, is_param:bool=True) -> Tensor: self.is_param = is_param @@ -516,7 +516,7 @@ class Tensor(RandMixin): ref_frames = [x.contiguous() for x in ref_frames or []] assert frame_pos.op is Ops.BIND, "frame_pos must be a bound Variable" srcs = (out:=Tensor.empty(*shape, device=self.device, dtype=self.dtype), self.contiguous(), state.contiguous(), *ref_frames) - fn = UOp(Ops.CUSTOM_FUNCTION, src=(frame_pos.src[0], *[UOp.const(dtypes.int, s) for s in shape]), arg="encdec") + fn = UOp(Ops.CUSTOM_FUNCTION, src=(frame_pos.src[0], *[UOp.const(s, dtypes.int) for s in shape]), arg="encdec") return Tensor(out.uop.after(fn.call(*[s.uop for s in srcs], frame_pos))) P = ParamSpec("P") diff --git a/tinygrad/uop/ops.py b/tinygrad/uop/ops.py index ac996f5767..b0743d93e3 100644 --- a/tinygrad/uop/ops.py +++ b/tinygrad/uop/ops.py @@ -98,8 +98,8 @@ def shape_to_shape_arg(arg:tuple[sint, ...]) -> UOp: for x in arg: if isinstance(x, UOp) and not dtypes.is_int(x.dtype): raise RuntimeError(f"shape must be int, got {x.dtype} in {arg}") if len(arg) == 0: return UOp(Ops.STACK) - elif len(arg) == 1: return UOp.const(dtypes.weakint, arg[0]) - else: return UOp(Ops.STACK, src=tuple(UOp.const(None, x) if isinstance(x, int) else x for x in arg)) + elif len(arg) == 1: return UOp.const(arg[0], dtypes.weakint) + else: return UOp(Ops.STACK, src=tuple(UOp.const(x) if isinstance(x, int) else x for x in arg)) def consumer_map_from_toposort(lst:Iterable[UOp]): ret: dict[UOp, dict[UOp, None]] = {} @@ -560,7 +560,7 @@ class UOp(RandMixin, metaclass=UOpMetaClass): if len(srcs) == 1 and isinstance(srcs[0], UOp): return srcs[0] return UOp(Ops.GROUP, src=tuple([x for x in srcs if x is not None])) def index(self, *srcs:UOp|int|None, **kwargs): - new_srcs: list[UOp] = [UOp.const(None, x) if isinstance(x, int) else x for x in srcs if x is not None] + new_srcs: list[UOp] = [UOp.const(x) if isinstance(x, int) else x for x in srcs if x is not None] if len(new_srcs) == 1 and new_srcs[0].op is Ops.CONST and self.op is Ops.STACK: return self.src[new_srcs[0].arg] return UOp(Ops.INDEX, src=(self,)+tuple(new_srcs), **kwargs) def __getitem__(self, idx): @@ -582,13 +582,13 @@ class UOp(RandMixin, metaclass=UOpMetaClass): @classmethod def _wrap_uop(cls, u:UOp) -> UOp: return u def const_like(self, b:ConstLike, dtype:DType|None=None): - return UOp.const(dtype or self.dtype, b, shape=self._shape) + return UOp.const(b, dtype or self.dtype, shape=self._shape) def vconst_like(self, b:ConstLike, dtype:DType|None=None): # for use after movement ops have been removed - return UOp.const(dtype or self.dtype, b).broadcast(self.max_numel()) + return UOp.const(b, dtype or self.dtype).broadcast(self.max_numel()) def ufix(self, x): if isinstance(x, UOp): return x - return UOp.const(None, x) + return UOp.const(x) def broadcast(self, count:int): if count == 1: return self return UOp(Ops.STACK, src=(self,)*count) @@ -608,12 +608,12 @@ class UOp(RandMixin, metaclass=UOpMetaClass): for idx in itertools.product(*[range(int(r.vmax)+1) for r in rngs])]) def alu(self, op, *src:UOp, **kwargs): return UOp(op, src=(self, *src), **kwargs) @staticmethod - def const(dtype:DType|None, b:ConstLike, shape:tuple[sint, ...]|None=None): + def const(b:ConstLike, dtype:DType|None=None, shape:tuple[sint, ...]|None=None): if dtype is None: dtype = dtypes.from_py(b) if isinstance(b, UOp): return b.cast(dtype) # NOTE: it always has to be STACK now, even if they are all the same if isinstance(b, tuple): - stk = [UOp.const(dtype, c) for c in b] + stk = [UOp.const(c, dtype) for c in b] ret = UOp.stack(*stk) else: ret = UOp(Ops.CONST, dtype, arg=dtype.const(b), src=()) @@ -640,7 +640,7 @@ class UOp(RandMixin, metaclass=UOpMetaClass): ret = UOp(Ops.REDUCE, src=(self.permute(perm),), arg=(op, len(reduce_axis))) return ret.reshape(tuple(s for i,s in enumerate(self.shape) if i not in axis)) if axis != reduce_axis else ret @staticmethod - def invalid(): return UOp.const(None, Invalid) + def invalid(): return UOp.const(Invalid) def valid(self, cond): return cond.where(self, self.const_like(Invalid)) def get_idx(self) -> UOp: @@ -648,7 +648,7 @@ class UOp(RandMixin, metaclass=UOpMetaClass): return self.src[1] if self.op is Ops.WHERE and self.src[2].arg is Invalid else self def get_valid(self) -> UOp: if self.op is Ops.STACK: return UOp.stack(*(x.get_valid() for x in self.src)) - return self.src[0] if self.op is Ops.WHERE and self.src[2].arg is Invalid else UOp.const(None, self.arg is not Invalid) + return self.src[0] if self.op is Ops.WHERE and self.src[2].arg is Invalid else UOp.const(self.arg is not Invalid) def reduce(self, *src:UOp, **kwargs): arg = kwargs.pop('arg', None) if isinstance(arg, Ops): arg = (arg, 0) @@ -1359,7 +1359,7 @@ class UPat(OpMixin): @functools.cache def cvar(name:str|None=None, dtype:DType|tuple[DType, ...]|None=None, arg=None): return UPat(Ops.CONST, dtype, name=name, arg=arg) @staticmethod - def const(dtype:DType|tuple[DType, ...]|None, b:ConstType): return UPat(Ops.CONST, dtype=dtype, arg=b) + def const(b:ConstType, dtype:DType|tuple[DType, ...]|None=None): return UPat(Ops.CONST, dtype=dtype, arg=b) # lil helper def f(self, op, **kwargs): return UPat(op, src=(self,), **kwargs) @@ -1383,7 +1383,7 @@ class UPat(OpMixin): def after(self, *src:UPat, **kwargs): return UPat(Ops.AFTER, self.match_dtype, (self,)+src, **kwargs) def end(self, *src:UPat, **kwargs): return UPat(Ops.END, src=(self,)+src, **kwargs) - def const_like(self, b:ConstLike): return UPat.const(self.match_dtype, cast(ConstType, b)) + def const_like(self, b:ConstLike): return UPat.const(cast(ConstType, b), self.match_dtype) def _broadcasted(self, y, reverse=False) -> tuple[UPat, UPat]: y = self.ufix(y) return (y, self) if reverse else (self, y) @@ -1748,7 +1748,7 @@ def _rebuild_dtype(n:UOp, new_src:tuple[UOp,...]) -> DType: all(a.dtype is b.dtype or b.base.arg is Invalid for a,b in zip(n.src, new_src)): return n.dtype return dtype_from_uop(n.op, new_src, n.arg) or n.dtype -def sint_to_uop(x:sint, dtype=dtypes.weakint) -> UOp: return UOp.const(dtype, x) +def sint_to_uop(x:sint, dtype=dtypes.weakint) -> UOp: return UOp.const(x, dtype) def to_max_shape(shape:tuple[sint, ...]) -> tuple[int, ...]: return tuple(int(x.vmax) if isinstance(x, UOp) else x for x in shape) def select_dtype(u:UOp): @@ -1786,7 +1786,7 @@ def lower_weak_srcs(ctx:dict[UOp, UOp]|None, u:UOp) -> UOp|None: def commit_weak(s:UOp, dt:DType) -> UOp: # a bare weak CONST commits directly (its number must fit), a weak non-const src takes the demand cast - return UOp.const(dt, s.arg) if s.op is Ops.CONST else s.cast(dt) + return UOp.const(s.arg, dt) if s.op is Ops.CONST else s.cast(dt) def commit_weak_srcs(u:UOp) -> UOp|None: if (dt:=least_upper_dtype(*(s.dtype for s in u.src))) in dtypes.weaks: return None @@ -1835,8 +1835,8 @@ pm_unbind = PatternMatcher([(UPat(Ops.BIND, name="x"), do_unbind)]) # ctx is source UOp for which we are finding a contiguous view for. used in contiguous_view_offset pm_contiguous_view_offset = PatternMatcher([ - (UPat(Ops.INDEX, src=(UPat(),)), lambda: UOp.const(None, 0)), - (UPat(Ops.INDEX, src=(UPat(), UPat(Ops.RANGE))), lambda: UOp.const(None, 0)), + (UPat(Ops.INDEX, src=(UPat(),)), lambda: UOp.const(0)), + (UPat(Ops.INDEX, src=(UPat(), UPat(Ops.RANGE))), lambda: UOp.const(0)), (UPat(Ops.INDEX, src=(UPat(), UPat(Ops.RANGE)+UPat.cvar('c'))), lambda c: c), (UPat(Ops.INDEX, src=(UPat(), UPat.cvar('c'))), lambda ctx, c: c if resolve(ctx.numel() == 1, False) else None), ]) diff --git a/tinygrad/uop/render.py b/tinygrad/uop/render.py index bf8b32a4f1..b88c97f13b 100644 --- a/tinygrad/uop/render.py +++ b/tinygrad/uop/render.py @@ -77,7 +77,7 @@ def render_marg(ctx,x:UOp): sugar = {Ops.SINK, Ops.END, Ops.STORE, Ops.LOAD, Ops.SQRT, Ops.INDEX, Ops.REDUCE, Ops.AFTER, Ops.THREEFRY, Ops.RECIPROCAL, Ops.EXP2, Ops.LOG2, Ops.SIN, Ops.CONTIGUOUS, Ops.BARRIER, Ops.DETACH} pm_pyrender_extra = PatternMatcher([ - (UPat(Ops.CONST, src=(), name="x"), lambda x: f"UOp.const({x.dtype}, {x.arg})"), + (UPat(Ops.CONST, src=(), name="x"), lambda x: f"UOp.const({x.arg}, {x.dtype})"), (UPat((Ops.CAST, Ops.BITCAST), name="x"), lambda ctx,x: f"{ctx[x.src[0]]}.{x.op.name.lower()}({x.dtype})"), (UPat(Ops.SPECIAL, src=(UPat(Ops.CONST),), name="x"), lambda x: f"UOp.special({x.src[0].arg}, {repr(x.arg)}, dtype={x.dtype})"), (UPat(Ops.BUFFER, src=(UPat(),), name="x"), lambda x: diff --git a/tinygrad/uop/spec.py b/tinygrad/uop/spec.py index 5298129d0a..b61b971f14 100644 --- a/tinygrad/uop/spec.py +++ b/tinygrad/uop/spec.py @@ -11,7 +11,7 @@ def validate_index(uidx:UOp, gate:UOp|None=None): if len(uidx.src) != 2: return True # skip for non final index. TODO: check more complex index with shape buf,idx = uidx.src if idx.op is Ops.CONST and idx.arg is Invalid: return True - if gate is None: gate = UOp.const(None, True) + if gate is None: gate = UOp.const(True) # TODO: check for overflow if not CHECK_OOB or is_image_shape(buf._shape): return True diff --git a/tinygrad/uop/symbolic.py b/tinygrad/uop/symbolic.py index 0b35760e85..0c39015d09 100644 --- a/tinygrad/uop/symbolic.py +++ b/tinygrad/uop/symbolic.py @@ -94,7 +94,7 @@ pm_data_invalid = PatternMatcher([ pm_remove_invalid = PatternMatcher([ (invalid_gate.named("w"), lambda cond,x,i,w: w.replace(src=(cond,x,w.const_like(0)))), - (UPat(Ops.STACK, name="s"), lambda s: s.replace(src=tuple(UOp.const(s.dtype, 0) if x.arg is Invalid else x for x in s.src)) + (UPat(Ops.STACK, name="s"), lambda s: s.replace(src=tuple(UOp.const(0, s.dtype) if x.arg is Invalid else x for x in s.src)) if any(x.arg is Invalid for x in s.src) else None), ]) @@ -112,11 +112,11 @@ symbolic_simple = pm_data_invalid + PatternMatcher([ (UPat(Ops.ADD, dtype=dtypes.weakint, name="x"), fold_add_divmod_recombine), (UPat.var("x", dtype=dtypes.bool) & UPat.cvar("c"), lambda x,c: x if c.arg else c), (UPat.var("x", dtype=dtypes.bool) | UPat.cvar("c"), lambda x,c: c if c.arg else x), - (UPat.var("x", dtype=dtypes.bool) != UPat.const(dtypes.bool, False), lambda x: x), # x != False -> x + (UPat.var("x", dtype=dtypes.bool) != UPat.const(False, dtypes.bool), lambda x: x), # x != False -> x (UPat(GroupOp.Idempotent, src=(UPat.var("x"), UPat.var("x"))), lambda x: x), (UPat.var("x", dtype=dtypes.bool).logical_not().logical_not(), lambda x: x), - (UPat.var("x", dtype=dtypes.bool).where(UPat.const(dtypes.bool, True), UPat.const(dtypes.bool, False)), lambda x: x), - (UPat.var("x", dtype=dtypes.bool).where(UPat.const(dtypes.bool, False), UPat.const(dtypes.bool, True)), lambda x: x.logical_not()), + (UPat.var("x", dtype=dtypes.bool).where(UPat.const(True, dtypes.bool), UPat.const(False, dtypes.bool)), lambda x: x), + (UPat.var("x", dtype=dtypes.bool).where(UPat.const(False, dtypes.bool), UPat.const(True, dtypes.bool)), lambda x: x.logical_not()), # CAST(bool -> int) != const — CAST(True)=1, CAST(False)=0, so fold based on const value (UPat.var("x", dtype=dtypes.bool).cast(dtypes.ints+(dtypes.weakint,)) != UPat.cvar("c"), lambda x,c: x if c.arg == 0 else x.logical_not() if c.arg == 1 else x.const_like(True)), @@ -381,7 +381,7 @@ def reduce_mul_chain(r:UOp) -> UOp|None: def drop_and_clauses(cond:UOp, x:UOp, i:UOp) -> UOp|None: keep, drop = partition(cond.split_uop(Ops.AND), lambda c: any(r in x.ranges for r in c.ranges)) - return UOp.const(None, True).uprod(*keep).where(x, i) if drop else None + return UOp.const(True).uprod(*keep).where(x, i) if drop else None pm_drop_and_clauses = PatternMatcher([(invalid_gate, drop_and_clauses)]) # move conditions from where to load's valid, drop clauses already in load @@ -396,7 +396,7 @@ def where_on_load(cond:UOp, buf:UOp, idx:UOp, or_cast:UOp) -> UOp|None: if len(keep) == len(where_clauses): return None idx = buf.index(idx.get_idx().valid(load_valid.uprod(*moved))) ret_idx = idx.cast(or_cast.dtype) if or_cast.op is Ops.CAST else idx - return UOp.const(None, True).uprod(*keep).where(ret_idx, ret_idx.const_like(0)) + return UOp.const(True).uprod(*keep).where(ret_idx, ret_idx.const_like(0)) # where after gated load becomes alt value, TODO: this is sort of duplicated with rules in devectorizer pm_move_where_on_load = PatternMatcher([ From 7f4dbb8090fc60ebae93599fb46517db3f6b299c Mon Sep 17 00:00:00 2001 From: chenyu Date: Fri, 31 Jul 2026 11:18:59 -0400 Subject: [PATCH 22/44] remove shape= from UOp.const [PR] (#17331) inlined to const_like --- test/backend/test_custom_kernel.py | 2 +- test/null/test_uop_graph.py | 18 +++++++++--------- test/null/test_viz.py | 2 +- test/unit/test_function.py | 4 ++-- tinygrad/mixin/elementwise.py | 2 +- tinygrad/uop/ops.py | 13 +++++-------- 6 files changed, 19 insertions(+), 22 deletions(-) diff --git a/test/backend/test_custom_kernel.py b/test/backend/test_custom_kernel.py index 0d9345ec73..2150a23612 100644 --- a/test/backend/test_custom_kernel.py +++ b/test/backend/test_custom_kernel.py @@ -341,7 +341,7 @@ class TestCustomKernel(unittest.TestCase): def test_partial_invalid_store_keeps_uncovered_reads(self): x = Tensor([10., 20., 30., 40.]) - after = x.uop.after(x.uop.shrink(((0, 2),)).store(UOp.const(Invalid, dtypes.float, shape=(2,)))) + after = x.uop.after(x.uop.shrink(((0, 2),)).store(Invalid)) self.assertEqual(Tensor(after).contiguous().tolist(), [10., 20., 30., 40.]) def test_multi_after_invalid_store_dep_removed(self): diff --git a/test/null/test_uop_graph.py b/test/null/test_uop_graph.py index 8a9a4d456b..9ec5b0ddd3 100644 --- a/test/null/test_uop_graph.py +++ b/test/null/test_uop_graph.py @@ -682,36 +682,36 @@ class TestUOpGetItem(unittest.TestCase): class TestUOpBroadcast(unittest.TestCase): def test_broadcast_row(self): - a = UOp.const(1, dtypes.float, shape=(4, 8)) - b = UOp.const(2, dtypes.float, shape=(4, 1)) + a = UOp.const(1, dtypes.float).expand((4, 8)) + b = UOp.const(2, dtypes.float).expand((4, 1)) c = a + b self.assertEqual(c.shape, (4, 8)) self.assertEqual(c.op, Ops.ADD) def test_broadcast_col(self): - a = UOp.const(1, dtypes.float, shape=(4, 8)) - b = UOp.const(2, dtypes.float, shape=(1, 8)) + a = UOp.const(1, dtypes.float).expand((4, 8)) + b = UOp.const(2, dtypes.float).expand((1, 8)) c = a + b self.assertEqual(c.shape, (4, 8)) self.assertEqual(c.op, Ops.ADD) def test_broadcast_lower_dim(self): - a = UOp.const(1, dtypes.float, shape=(4, 8)) - b = UOp.const(2, dtypes.float, shape=(8,)) + a = UOp.const(1, dtypes.float).expand((4, 8)) + b = UOp.const(2, dtypes.float).expand((8,)) c = a * b self.assertEqual(c.shape, (4, 8)) self.assertEqual(c.op, Ops.MUL) def test_broadcast_scalar(self): - a = UOp.const(1, dtypes.float, shape=(4, 8)) + a = UOp.const(1, dtypes.float).expand((4, 8)) c = a * 2 self.assertEqual(c.shape, (4, 8)) self.assertEqual(c.op, Ops.MUL) def test_broadcast_symbolic_same_shape(self): t = Variable("t", 1, 10) - a = UOp.const(1, dtypes.float, shape=(1, 1, t)) - b = UOp.const(2, dtypes.float, shape=(1, 1, t)) + a = UOp.const(1, dtypes.float).expand((1, 1, t)) + b = UOp.const(2, dtypes.float).expand((1, 1, t)) c = a + b self.assertEqual(c.op, Ops.ADD) diff --git a/test/null/test_viz.py b/test/null/test_viz.py index 8d66e5f561..ddcfa15bdc 100644 --- a/test/null/test_viz.py +++ b/test/null/test_viz.py @@ -253,7 +253,7 @@ class TestViz(unittest.TestCase): def test_const_reshape_expand_folded(self): # CONST->EXPAND should be folded into the ALU node, not shown as separate EXPAND nodes - c = UOp.const(1.0, dtypes.float, shape=(3,4)) # creates CONST->EXPAND chain + c = UOp.const(1.0, dtypes.float).expand((3,4)) # creates CONST->EXPAND chain a = UOp.variable("a", 0.0, 10.0, dtypes.float) alu = a + c with save_viz() as viz: diff --git a/test/unit/test_function.py b/test/unit/test_function.py index 6e4ce4d0d4..e333d8f9a6 100644 --- a/test/unit/test_function.py +++ b/test/unit/test_function.py @@ -2,7 +2,7 @@ import numpy as np import unittest from tinygrad.function import function from tinygrad import Tensor, GlobalCounters, Device -from tinygrad.dtype import dtypes, Invalid +from tinygrad.dtype import Invalid from tinygrad.uop.ops import UOp, Ops, KernelInfo, ProgramInfo class TestFunction(unittest.TestCase): @@ -593,7 +593,7 @@ class TestFunctionTuple(unittest.TestCase): state = Tensor([10., 20., 30., 40.], device="CPU").contiguous().realize() @function(precompile=True, allow_implicit=True) def f(a:Tensor): - after = state.uop.after(state.uop.shrink(((0, 2),)).store(UOp.const(Invalid, dtypes.float32, shape=(2,)))) + after = state.uop.after(state.uop.shrink(((0, 2),)).store(Invalid)) return Tensor(after).contiguous() + a out = f(Tensor([1., 1., 1., 1.], device="CPU").contiguous().realize()) np.testing.assert_allclose(out.numpy(), [11., 21., 31., 41.]) diff --git a/tinygrad/mixin/elementwise.py b/tinygrad/mixin/elementwise.py index 3bc2e6347d..f73e0a0bea 100644 --- a/tinygrad/mixin/elementwise.py +++ b/tinygrad/mixin/elementwise.py @@ -24,7 +24,7 @@ class ElementwiseMixin(CreationMixin): out_dtype = least_upper_dtype(x.dtype, y.dtype) # keep weak CONST weak, might lift weakint -> weakfloat def promote(t): - if t.dtype in dtypes.weaks and t._uop.base.op is Ops.CONST: return t._wrap_uop(t._uop.const(t._uop.base.arg, weak_dtype(out_dtype), t.shape)) + if t.dtype in dtypes.weaks and t._uop.base.op is Ops.CONST: return t._wrap_uop(t._uop.const_like(t._uop.base.arg, weak_dtype(out_dtype))) return t.cast(out_dtype) return promote(x), promote(y) diff --git a/tinygrad/uop/ops.py b/tinygrad/uop/ops.py index b0743d93e3..b182e88345 100644 --- a/tinygrad/uop/ops.py +++ b/tinygrad/uop/ops.py @@ -582,7 +582,8 @@ class UOp(RandMixin, metaclass=UOpMetaClass): @classmethod def _wrap_uop(cls, u:UOp) -> UOp: return u def const_like(self, b:ConstLike, dtype:DType|None=None): - return UOp.const(b, dtype or self.dtype, shape=self._shape) + ret = UOp.const(b, dtype or self.dtype) + return ret._mop(Ops.EXPAND, arg=self._shape) if self._shape and ret._shape != self._shape else ret def vconst_like(self, b:ConstLike, dtype:DType|None=None): # for use after movement ops have been removed return UOp.const(b, dtype or self.dtype).broadcast(self.max_numel()) @@ -608,16 +609,12 @@ class UOp(RandMixin, metaclass=UOpMetaClass): for idx in itertools.product(*[range(int(r.vmax)+1) for r in rngs])]) def alu(self, op, *src:UOp, **kwargs): return UOp(op, src=(self, *src), **kwargs) @staticmethod - def const(b:ConstLike, dtype:DType|None=None, shape:tuple[sint, ...]|None=None): + def const(b:ConstLike, dtype:DType|None=None): if dtype is None: dtype = dtypes.from_py(b) if isinstance(b, UOp): return b.cast(dtype) # NOTE: it always has to be STACK now, even if they are all the same - if isinstance(b, tuple): - stk = [UOp.const(c, dtype) for c in b] - ret = UOp.stack(*stk) - else: - ret = UOp(Ops.CONST, dtype, arg=dtype.const(b), src=()) - return ret._mop(Ops.EXPAND, arg=shape) if shape is not None and shape != () and ret.shape != shape else ret + if isinstance(b, tuple): return UOp.stack(*[UOp.const(c, dtype) for c in b]) + return UOp(Ops.CONST, dtype, arg=dtype.const(b), src=()) @staticmethod def range(end:sint, axis_id, axis_type=AxisType.WEAK, *arg, dtype=dtypes.weakint, src=(), **kwargs): return UOp(Ops.RANGE, src=(sint_to_uop(end, dtype),)+src, arg=(axis_id, axis_type)+arg, **kwargs) From 1095bbe409f5ed3cbeca74aa3c2ca09bef634309 Mon Sep 17 00:00:00 2001 From: nimlgen <138685161+nimlgen@users.noreply.github.com> Date: Fri, 31 Jul 2026 18:43:23 +0300 Subject: [PATCH 23/44] hcq2: fix ib reuse (#17330) * hcq2: initialize IB reuse counters at link * x --- extra/hcq2/ops_amd2.py | 3 ++- tinygrad/runtime/support/hcq2.py | 1 + 2 files changed, 3 insertions(+), 1 deletion(-) diff --git a/extra/hcq2/ops_amd2.py b/extra/hcq2/ops_amd2.py index 640a519688..38a76c257a 100644 --- a/extra/hcq2/ops_amd2.py +++ b/extra/hcq2/ops_amd2.py @@ -158,7 +158,8 @@ def pm4_submit(ctx, lin): ib = UOp.placeholder((size_dw + 2,), dtypes.uint32, next(UOp.unique_num), device=devs, volatile=True).rtag("cmdbuf") done_idx, submit_idx = UOp.const(size_dw + 0, dtypes.int), UOp.const(size_dw + 1, dtypes.int) - submitted = (counter:=ib.after(make_patches(ib, [((size_dw + i) * 4, UOp.const(0, dtypes.uint32)) for i in range(2)])).index(submit_idx)).load() + init_counters = make_patches(ib, [((size_dw + i) * 4, UOp.const(0, dtypes.uint32)) for i in range(2)]).rtag("link") + submitted = (counter:=ib.after(init_counters).index(submit_idx)).load() completed = ib.after(loop:=UOp.loop(0)).index(done_idx).load() ib_free = completed.end(loop, completed != submitted) diff --git a/tinygrad/runtime/support/hcq2.py b/tinygrad/runtime/support/hcq2.py index 3d5fdf567f..3952f30e9e 100644 --- a/tinygrad/runtime/support/hcq2.py +++ b/tinygrad/runtime/support/hcq2.py @@ -252,6 +252,7 @@ def is_value_known_at_link(val:UOp) -> bool: return not val.variables() and not runtime_reads and all(b.op is not Ops.PARAM or b.tag is not None for b in addressed_bufs) def is_link_patch(p:UOp, jit:bool) -> bool: + if p.tag == "link": return True store = p.src[0] if (is_binary_patch:=(p.op is Ops.END and p.src[0].op is Ops.STORE)) else p if not jit: return store.buf_uop.tag == "program" return is_binary_patch or (store.op is Ops.STORE and is_value_known_at_link(store.src[1])) From b95bd5b2a54229898268afa09aa0a8805d652fc5 Mon Sep 17 00:00:00 2001 From: Christopher Milan Date: Fri, 31 Jul 2026 12:39:04 -0400 Subject: [PATCH 24/44] don't reset chestnut in benchmark (#17333) --- .github/workflows/benchmark.yml | 2 -- 1 file changed, 2 deletions(-) diff --git a/.github/workflows/benchmark.yml b/.github/workflows/benchmark.yml index 167dced5ae..b4427d9d05 100644 --- a/.github/workflows/benchmark.yml +++ b/.github/workflows/benchmark.yml @@ -515,8 +515,6 @@ jobs: run: | echo "CACHEDB=/tmp/staging.db" >> $GITHUB_ENV rm -f /tmp/staging.db /tmp/staging.db-shm /tmp/staging.db-wal - - name: reset chestnut - run: python3 extra/usbgpu/debug.py -rn - name: openpilot compile3 big_driving_supercombo run: BENCHMARK_LOG=usbgpu_openpilot_big_driving_supercombo PICKLE_OOB=1 PYTHONPATH="." GMMU=0 DEV=USB+AMD:LLVM ASSERT_MIN_STEP_TIME=50 python3 examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/10926f2c0911821ca0e72439c1c3bf3ec11f0a08789aa14b7ee8f25379b2afa4 openpilot.pkl - name: openpilot load_pickle big_driving_supercombo From a11ee26bb81a123f9207ee94d8e1d5fc1c4e9eac Mon Sep 17 00:00:00 2001 From: qazal <77887910+Qazalin@users.noreply.github.com> Date: Sat, 1 Aug 2026 01:16:30 +0800 Subject: [PATCH 25/44] viz: prep for faster cli DEBUG=3 (#17334) * move data * split --- tinygrad/viz/serve.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/tinygrad/viz/serve.py b/tinygrad/viz/serve.py index 11e3858b26..d9c32c0a31 100755 --- a/tinygrad/viz/serve.py +++ b/tinygrad/viz/serve.py @@ -191,8 +191,8 @@ def get_full_rewrite(data:VizData, ctx:TrackedGraphRewrite, depth:int|None=None) "diff":list(difflib.unified_diff(pystr(u0).splitlines(), pystr(u1).splitlines())), "upat":(upat_loc, match_repr), "_sink":new_sink} if not ctx.bottom_up: next_sink = new_sink -def get_sink_at(upats:tuple[str, ...], viz_data:VizData, ctx:TrackedGraphRewrite, depth:int|None=None) -> UOp|None: - for s in get_full_rewrite(viz_data, ctx, depth=depth): +def get_sink_at(upats:tuple[str, ...], viz_data:VizData, kernel_idx:int, lin_idx:int, depth:int|None=None) -> UOp|None: + for s in get_full_rewrite(viz_data, ctx:=viz_data.trace.rewrites[kernel_idx][lin_idx], depth=depth): if (s["upat"] is not None and any(n in s["upat"][1] for n in upats)) or len(ctx.matches) == 0: return s["_sink"] return None @@ -617,15 +617,15 @@ def get_render(viz_data:VizData, query:str) -> dict: data = viz_data.ctxs[i]["steps"][j]["_data"] if fmt == "graph-rewrites": return {"value":get_full_rewrite(viz_data, viz_data.trace.rewrites[i][j]), "content_type":"text/event-stream"} if fmt == "uops": - if (sink:=get_sink_at(("do_linearize",), viz_data, viz_data.trace.rewrites[i][data])) is None: return {"src":"No linear found"} + if (sink:=get_sink_at(("do_linearize",), viz_data, i, data)) is None: return {"src":"No linear found"} return {"src":sink.arg} if sink.op is Ops.REWRITE_ERROR else {"src":get_stdout(lambda: print_uops(list(unwrap(sink).src[1].src)))} if fmt == "code": - if (sink:=get_sink_at(("do_render",), viz_data, viz_data.trace.rewrites[i][data], depth=1)) is None: return {"src":"No source found"} + if (sink:=get_sink_at(("do_render",), viz_data, i, data, depth=1)) is None: return {"src":"No source found"} return {"src":sink.arg} if sink.op is Ops.REWRITE_ERROR else {"src":sink.src[2].arg, "lang":"cpp"} if fmt == "asm": ret:dict = {} renderer, idx = data - if (sink:=get_sink_at(("do_compile","do_assemble"), viz_data, viz_data.trace.rewrites[i][idx], depth=1)) is None: return {"src":"No binary found"} + if (sink:=get_sink_at(("do_compile","do_assemble"), viz_data, i, idx, depth=1)) is None: return {"src":"No binary found"} if sink.op is Ops.REWRITE_ERROR: return {"src":sink.arg} lib:bytes = sink.src[3].arg if renderer.target.arch.startswith("gfx"): From ad247504879791e3039c60701e3df00438b5b48b Mon Sep 17 00:00:00 2001 From: chenyu Date: Fri, 31 Jul 2026 13:24:31 -0400 Subject: [PATCH 26/44] const(value, dtype) -> const(value).cast(dtype) in tests (#17335) --- test/backend/test_call.py | 4 +- test/backend/test_linearizer_dumb.py | 2 +- test/backend/test_pickle.py | 2 +- test/backend/test_renderer_failures.py | 12 +-- test/backend/test_tensor.py | 10 +- test/backend/test_uops.py | 18 ++-- test/helpers.py | 2 +- test/null/test_const_folding.py | 7 +- test/null/test_gpudims.py | 2 +- test/null/test_gradient.py | 8 +- test/null/test_graph_rewrite.py | 86 ++++++++--------- test/null/test_helpers.py | 8 +- test/null/test_linearizer_failures.py | 2 +- test/null/test_microbenchmarks.py | 18 ++-- test/null/test_pattern_matcher.py | 63 +++++++------ test/null/test_schedule.py | 4 +- test/null/test_simplify_valid_idx.py | 4 +- test/null/test_symbolic_failures.py | 23 ++--- test/null/test_tensor_uop_mixin.py | 2 +- test/null/test_transcendental_helpers.py | 44 ++++----- test/null/test_uop_graph.py | 114 +++++++++++------------ test/null/test_uop_repr.py | 28 +++--- test/null/test_uop_resolve.py | 14 +-- test/null/test_uop_symbolic.py | 39 ++++---- test/null/test_uop_vmin_vmax.py | 38 ++++---- test/null/test_uops.py | 55 +++++------ test/null/test_validate_oob.py | 14 +-- test/null/test_viz.py | 23 ++--- test/unit/test_assign.py | 2 +- test/unit/test_jit.py | 4 +- test/unit/test_multitensor.py | 4 +- test/unit/test_tensor_data.py | 2 +- 32 files changed, 330 insertions(+), 328 deletions(-) diff --git a/test/backend/test_call.py b/test/backend/test_call.py index 53d8bb0ff2..851424c8e4 100644 --- a/test/backend/test_call.py +++ b/test/backend/test_call.py @@ -6,11 +6,11 @@ from tinygrad.renderer.cstyle import CStyleLanguage from tinygrad.uop.ops import KernelInfo def call_out_kernel(F:UOp, C:UOp) -> UOp: - call = F[0].load().call(UOp.const(3, dtypes.int), C[0], ret_dtype=dtypes.void) + call = F[0].load().call(UOp.const(3).cast(dtypes.int), C[0], ret_dtype=dtypes.void) return C.after(call)[1].store(C.after(call)[0].load() + 1).sink(arg=KernelInfo(name="call_out")) def call_ret_kernel(F:UOp, C:UOp) -> UOp: - val = F[0].load().call(UOp.const(21, dtypes.int), ret_dtype=dtypes.int) + val = F[0].load().call(UOp.const(21).cast(dtypes.int), ret_dtype=dtypes.int) return C[0].store(val * 2).sink(arg=KernelInfo(name="call_ret")) @unittest.skipUnless(isinstance(Device["CPU"].renderer, CStyleLanguage), "TODO: CALL is rendered in C style only") diff --git a/test/backend/test_linearizer_dumb.py b/test/backend/test_linearizer_dumb.py index 6604171484..317f0abd3b 100644 --- a/test/backend/test_linearizer_dumb.py +++ b/test/backend/test_linearizer_dumb.py @@ -22,7 +22,7 @@ class TestLinearizerFailure(unittest.TestCase): c8 = UOp.range(UOp.const(16), 2007, AxisType.GROUP_REDUCE) c9 = UOp.param(2, dtypes.uchar, (47040000,)) c10 = c9.index((((c3*UOp.const(4704000))+c2)+(c6*UOp.const(784))).valid(UOp.const(True))) - c11 = c5.alu(Ops.CMPNE, ((((c3*UOp.const(6000))+c6)+((c7*UOp.const(16))+c8)).alu(Ops.CMPLT, UOp.const(59999)).where(UOp.const(0, dtypes.int), UOp.const(1, dtypes.int)).reduce(c7, c8, arg=Ops.ADD)+UOp.const(-1, dtypes.int))).where(UOp.const(0, dtypes.uchar), c10).reduce(c6, arg=Ops.ADD) + c11 = c5.alu(Ops.CMPNE, ((((c3*UOp.const(6000))+c6)+((c7*UOp.const(16))+c8)).alu(Ops.CMPLT, UOp.const(59999)).where(UOp.const(0).cast(dtypes.int), UOp.const(1).cast(dtypes.int)).reduce(c7, c8, arg=Ops.ADD)+UOp.const(-1).cast(dtypes.int))).where(UOp.const(0).cast(dtypes.uchar), c10).reduce(c6, arg=Ops.ADD) c12 = c0.index((((c1*UOp.const(7840))+(c2*UOp.const(10)))+c3).valid(UOp.const(True))).store(c11).end(c1, c2, c3) ast = c12.sink(arg=KernelInfo(name='test', axis_types=(), dont_use_locals=False, applied_opts=(Opt(op=OptOps.GROUP, axis=1, arg=16),), opts_to_apply=None)) _ = to_program(ast, Device["METAL"].renderer) diff --git a/test/backend/test_pickle.py b/test/backend/test_pickle.py index 435270bfea..cbb0aa3782 100644 --- a/test/backend/test_pickle.py +++ b/test/backend/test_pickle.py @@ -13,7 +13,7 @@ class TestPickle(unittest.TestCase): def test_pickle_pattern_matcher(self): pm = PatternMatcher([(UPat.cvar('x'), lambda x: x*2)]) - sink = UOp.const(2, dtypes.int) + sink = UOp.const(2) tt = pm.rewrite(sink) pm_str = pickle.dumps(pm) pm2 = pickle.loads(pm_str) diff --git a/test/backend/test_renderer_failures.py b/test/backend/test_renderer_failures.py index 8d3c09d4b0..826928bece 100644 --- a/test/backend/test_renderer_failures.py +++ b/test/backend/test_renderer_failures.py @@ -24,7 +24,7 @@ def _setup_and_test_alu(alu_op:Ops, input_val:ConstType, *alu_src_uops:UOp): dtype = alu_src_uops[0].dtype a = UOp.param(0, dtype, (1,)) b = UOp.param(1, dtype, (1,)) - idx = UOp.const(0, dtypes.int) + idx = UOp.const(0) ld = b.index(idx).load() alu = ld.alu(alu_op, *alu_src_uops) store = UOp.store(a.index(idx), alu) @@ -35,7 +35,7 @@ class TestRendererFailures(unittest.TestCase): def test_gated_store_with_alu(self): a = UOp.param(0, dtypes.int, (4,)) gate_alu = (lidx0:=UOp.special(4, 'lidx0')).ne(0) - gated_alu_store = UOp(Ops.STORE, src=(a.index(lidx0.valid(gate_alu)), UOp.const(1, dtypes.int))) + gated_alu_store = UOp(Ops.STORE, src=(a.index(lidx0.valid(gate_alu)), UOp.const(1).cast(dtypes.int))) sink = UOp(Ops.SINK, src=(gated_alu_store,), arg=KernelInfo()) ret = _test_uop_result([], sink, local_size=[4, 1, 1])[0] np.testing.assert_equal(ret, [0, 1, 1, 1]) @@ -45,7 +45,7 @@ class TestRendererFailures(unittest.TestCase): a = UOp.param(0, dtypes.int, (8,)) gate_alu_0 = (lidx0:=UOp.special(4, 'lidx0')).ne(0) gate_alu_1 = (lidx1:=UOp.special(2, 'lidx1')).ne(0) - gated_alu_store = UOp(Ops.STORE, src=(a.index((lidx0+lidx1*4).valid(gate_alu_0&gate_alu_1)), UOp.const(1, dtypes.int))) + gated_alu_store = UOp(Ops.STORE, src=(a.index((lidx0+lidx1*4).valid(gate_alu_0&gate_alu_1)), UOp.const(1).cast(dtypes.int))) sink = UOp(Ops.SINK, src=(gated_alu_store,), arg=KernelInfo()) ret = _test_uop_result([], sink, local_size=[4, 2, 1])[0] np.testing.assert_equal(ret, [0, 0, 0, 0, 0, 1, 1, 1]) @@ -54,7 +54,7 @@ class TestRendererFailures(unittest.TestCase): class TestCStyleFailures(unittest.TestCase): def test_inline_const_alu(self): # CPU doesn't use the max function - ret = _setup_and_test_alu(Ops.MAX, 1, UOp.const(dtypes.int.min+1, dtypes.int)) + ret = _setup_and_test_alu(Ops.MAX, 1, UOp.const(dtypes.int.min+1).cast(dtypes.int)) self.assertEqual(ret[0], 1) def _test_src_strip_paren(self, op: Ops, should_strip_paren:bool=True): @@ -80,7 +80,7 @@ class TestWGSLFailures(unittest.TestCase): def test_multiply_infinity(self): # multiplying a positive constant by infinity should return infinity # WGSL pipelines do not handle this reliably, some of which return zero, unless infinity always comes from a read on a dynamic buffer - ret = _setup_and_test_alu(Ops.MUL, 5.0, UOp.const(float("inf"), dtypes.float32)) + ret = _setup_and_test_alu(Ops.MUL, 5.0, UOp.const(float("inf")).cast(dtypes.float32)) self.assertEqual(ret[0], float("inf")) # WGSL has a specific select(alt, val, gate) ternary operator instead of gate?val:alt @@ -104,7 +104,7 @@ class TestPTXFailures(unittest.TestCase): def test_gated_store_with_if(self): a = UOp.param(0, dtypes.int, (4,)) gate_alu = (lidx0:=UOp.special(4, 'lidx0')).ne(0) - val = UOp.const(1, dtypes.int) + val = UOp.const(1).cast(dtypes.int) if_uop = UOp(Ops.IF, src=(gate_alu,)) gated_alu_store = UOp(Ops.STORE, src=(a.index(lidx0, if_uop), val)) sink = UOp(Ops.SINK, src=(gated_alu_store,), arg=KernelInfo()) diff --git a/test/backend/test_tensor.py b/test/backend/test_tensor.py index 421343706c..c44a2b9e08 100644 --- a/test/backend/test_tensor.py +++ b/test/backend/test_tensor.py @@ -24,13 +24,13 @@ class TestTinygrad(unittest.TestCase): self.assertEqual(Tensor(3.14).shape, ()) def test_deviceless_const_construct_device_repr(self): - t = Tensor(UOp.const(2.0, dtypes.float)) + t = Tensor(UOp.const(2.0).cast(dtypes.float)) self.assertIsNone(t.uop.device) self.assertIsNone(t.device) self.assertIn(" UOp: - if op is Ops.CONST: uops.append(UOp.const(arg, dtype)) + if op is Ops.CONST: uops.append(UOp.const(arg).cast(dtype)) elif op is Ops.PARAM: uops.append(UOp.param(arg, dtype, shape=(1,))) else: uops.append(UOp(op, dtype, tuple(src), arg)) return uops[-1] @@ -43,7 +43,7 @@ def _test_single_value_const(vals, op, dts): buf_store = uop(uops, Ops.PARAM, output_dtype, (), 0) loads = (uop(uops, Ops.CONST, dtype, [], a) for a,dtype in zip(vals, dts)) alu = uop(uops, op, output_dtype, loads) - out = buf_store[UOp.const(0, dtypes.int32)].store(alu) + out = buf_store[UOp.const(0).cast(dtypes.int32)].store(alu) buf = Buffer(Device.DEFAULT, 1, output_dtype).allocate() run_uops([out], [buf]) return np.frombuffer(buf.as_memoryview(), _to_np_dtype(output_dtype))[0] @@ -221,12 +221,12 @@ class TestAssembly(unittest.TestCase): def test_bitshift_left(self): g1 = UOp.param(0, dtypes.int32, shape=(3,)) out = UOp.param(1, dtypes.int32, shape=(2,)) - c1 = UOp.const(2, dtypes.int) - c2 = UOp.const(3, dtypes.int) + c1 = UOp.const(2) + c2 = UOp.const(3) l1 = g1.index(c1) a1 = UOp(Ops.MUL, src=(l1, c1)) a2 = UOp(Ops.MUL, src=(l1, c2)) - uops = to_uops_list([out.index(UOp.const(0, dtypes.int)).store(a1), out.index(UOp.const(1, dtypes.int)).store(a2)], + uops = to_uops_list([out.index(UOp.const(0)).store(a1), out.index(UOp.const(1)).store(a2)], ren=Device[Device.DEFAULT].renderer) Device[Device.DEFAULT].renderer.render(uops) ops = [x.op for x in uops] @@ -249,16 +249,16 @@ class TestAssembly(unittest.TestCase): def test_mulacc_shl(self): g1 = UOp.param(0, dtypes.int32, shape=(2,)) - c1 = UOp.const(0, dtypes.int) - c2 = UOp.const(1, dtypes.int) - expr = g1.index(c1) * UOp.const(4096, dtypes.int) + g1.index(c2) + c1 = UOp.const(0) + c2 = UOp.const(1) + expr = g1.index(c1) * UOp.const(4096) + g1.index(c2) uops = to_uops_list([expr], ren=Device[Device.DEFAULT].renderer) Device[Device.DEFAULT].renderer.render(uops) self.assertIn(Ops.MULACC, [x.op for x in uops]) def test_use_cmpeq(self): g = UOp.param(0, dtypes.uint32, shape=(8,)) - c = UOp.const(7, dtypes.uint) + c = UOp.const(7) comp = g.index(c).ne(c).ne(True) uops = to_uops_list([comp], ren=Device[Device.DEFAULT].renderer) Device[Device.DEFAULT].renderer.render(uops) diff --git a/test/helpers.py b/test/helpers.py index 88c92787dd..1e2b893d16 100644 --- a/test/helpers.py +++ b/test/helpers.py @@ -90,7 +90,7 @@ def eval_uop(uop:UOp, inputs:list[tuple[DType, list[Any]]]|None=None, vals:tuple bufs.append(buf:=allocator.alloc(len(data) * buf_dt.itemsize)) allocator._copyin(buf, memoryview(struct.pack(str(len(data)) + (buf_dt.fmt or ""), *data))) g = UOp.param(0, uop.dtype, (1,)) - prg = to_program(UOp.store(g.index(UOp.const(0, dtypes.int)), uop).sink(arg=KernelInfo()), PythonRenderer(Target("PYTHON"))) + prg = to_program(UOp.store(g.index(UOp.const(0)), uop).sink(arg=KernelInfo()), PythonRenderer(Target("PYTHON"))) prog = dev.runtime(prg.to_elf()) prog(out_buf:=allocator.alloc(uop.dtype.itemsize), *bufs, vals=vals) return out_buf.cast(uop.dtype.fmt or "").tolist()[0] diff --git a/test/null/test_const_folding.py b/test/null/test_const_folding.py index e2a7bbfcd2..373174528c 100644 --- a/test/null/test_const_folding.py +++ b/test/null/test_const_folding.py @@ -40,10 +40,9 @@ class TestWeakConstFolding(unittest.TestCase): self.assertEqual((out.op, out.dtype, out.arg), (Ops.CONST, dtypes.weakint, 2**41)) def test_float_unaries(self): - for dtype in (dtypes.weakfloat,): - for op in (Ops.SIN, Ops.LOG2, Ops.EXP2, Ops.SQRT, Ops.RECIPROCAL): - out = UOp.const(4, dtype).alu(op).simplify() - self.assertEqual((out.op, out.dtype), (Ops.CONST, dtypes.weakfloat)) + for op in (Ops.SIN, Ops.LOG2, Ops.EXP2, Ops.SQRT, Ops.RECIPROCAL): + out = UOp.const(4.0).alu(op).simplify() + self.assertEqual((out.op, out.dtype), (Ops.CONST, dtypes.weakfloat)) def test_weakfloat_math(self): out = (UOp.const(1.25) + UOp.const(2.5)).simplify() diff --git a/test/null/test_gpudims.py b/test/null/test_gpudims.py index dbca42b954..06c50e4655 100644 --- a/test/null/test_gpudims.py +++ b/test/null/test_gpudims.py @@ -107,7 +107,7 @@ class TestGroupedDims(unittest.TestCase): def test_global_prod_max(self): g, l = UOp.range(256, 0, AxisType.GLOBAL), UOp.range(256, 1, AxisType.LOCAL) - sink = UOp.param(0, dtypes.float, (512,)).index(g + l).store(UOp.const(1.0, dtypes.float)).end(g, l).sink(arg=KernelInfo()) + sink = UOp.param(0, dtypes.float, (512,)).index(g + l).store(UOp.const(1.0)).end(g, l).sink(arg=KernelInfo()) class R(Renderer): global_max, local_max, global_prod_max = (256, 256, 256), (128, 128, 128), (128, 128, 128) specials = [u for u in add_gpudims(R(Target()), sink).toposort() if u.op is Ops.SPECIAL] self.assertGreater(len([s for s in specials if "lidx" in s.arg]), 1) diff --git a/test/null/test_gradient.py b/test/null/test_gradient.py index ad39dfa7c9..7153ef4835 100644 --- a/test/null/test_gradient.py +++ b/test/null/test_gradient.py @@ -14,10 +14,10 @@ class TestGradient(unittest.TestCase): def _test_one_input_function(self, f:Callable, jf:Callable|None=None): if jf is None: jf = f x = UOp.variable('x', -math.inf, math.inf, dtype=dtypes.float) - gx = compute_gradient(f(x), UOp.const(1.0, dtypes.float), set([x]))[x] + gx = compute_gradient(f(x), UOp.const(1.0), set([x]))[x] for val in [-5., -2.0, 0.0, 2.0, 5.]: - tg_out = gx.substitute({x: x.const_like(val)}).ssimplify() + tg_out = gx.substitute({x: UOp.const(val)}).ssimplify() tx = torch.tensor([val], dtype=torch.float, requires_grad=True) torch_out = torch.autograd.grad(jf(tx), tx)[0].item() self._cmp_nan_okay(tg_out, torch_out) @@ -26,13 +26,13 @@ class TestGradient(unittest.TestCase): if jf is None: jf = f x = UOp.variable('x', -math.inf, math.inf, dtype=dtypes.float) y = UOp.variable('y', -math.inf, math.inf, dtype=dtypes.float) - grads = compute_gradient(f(x, y), UOp.const(1.0, dtypes.float), set([x, y])) + grads = compute_gradient(f(x, y), UOp.const(1.0), set([x, y])) gx, gy = grads[x], grads[y] for valx in [-5., -2.0, 0.0, 2.0, 5.]: for valy in [-5., -2.0, 0.0, 2.0, 5.]: # Substitute the values into the gradient expressions - substitutions = {x: x.const_like(valx), y: y.const_like(valy)} + substitutions = {x: UOp.const(valx), y: UOp.const(valy)} tg_out_x = gx.substitute(substitutions).ssimplify() tg_out_y = gy.substitute(substitutions).ssimplify() diff --git a/test/null/test_graph_rewrite.py b/test/null/test_graph_rewrite.py index 56d19f5965..6b57882d26 100644 --- a/test/null/test_graph_rewrite.py +++ b/test/null/test_graph_rewrite.py @@ -32,30 +32,30 @@ def evaluate_uop(uop, variables): class TestArithmeticSimplifications(unittest.TestCase): def test_full_graph_rewrite_division_by_zero(self): - optimized_div_uop = apply_rewrite(UOp.const(10.0, dtypes.float32) / UOp.const(0.0, dtypes.float32)) + optimized_div_uop = apply_rewrite(UOp.const(10.0) / UOp.const(0.0)) self.assertEqual(optimized_div_uop.op, Ops.CONST) self.assertTrue(math.isinf(optimized_div_uop.arg) or math.isnan(optimized_div_uop.arg)) def test_full_graph_rewrite_redundant_operations(self): - optimized_uop = apply_rewrite((UOp.const(10.0, dtypes.float32) + UOp.const(0.0, dtypes.float32)) * UOp.const(1.0, dtypes.float32)) + optimized_uop = apply_rewrite((UOp.const(10.0) + UOp.const(0.0)) * UOp.const(1.0)) self.assertEqual(optimized_uop.op, Ops.CONST) self.assertEqual(optimized_uop.arg, 10.0) def test_full_graph_rewrite_large_graph(self): - prev_uop = UOp.const(0, dtypes.int32) + prev_uop = UOp.const(0) for i in range(1, 101): - prev_uop += UOp.const(i, dtypes.int32) + prev_uop += UOp.const(i) optimized_uop = apply_rewrite(prev_uop) self.assertEqual(optimized_uop.op, Ops.CONST) self.assertEqual(optimized_uop.arg, sum(range(1, 101))) def test_full_graph_rewrite_division_by_one(self): - optimized_uop = apply_rewrite(UOp.const(42.0, dtypes.float32) / UOp.const(1.0, dtypes.float32)) + optimized_uop = apply_rewrite(UOp.const(42.0) / UOp.const(1.0)) self.assertEqual(optimized_uop.op, Ops.CONST) self.assertEqual(optimized_uop.arg, 42.0) def test_full_graph_rewrite_modulo_by_one(self): - optimized_uop = apply_rewrite(UOp.const(42, dtypes.int32) % UOp.const(1, dtypes.int32)) + optimized_uop = apply_rewrite(UOp.const(42) % UOp.const(1)) self.assertEqual(optimized_uop.op, Ops.CONST) self.assertEqual(optimized_uop.arg, 0) @@ -63,17 +63,17 @@ class TestArithmeticSimplifications(unittest.TestCase): class TestFoldingAndReduction(unittest.TestCase): @unittest.skip("reduce is removed now") def test_full_graph_rewrite_constant_reduction_folding(self): - const1 = UOp.const(5, dtypes.int32) - const2 = UOp.const(10, dtypes.int32) - const3 = UOp.const(20, dtypes.int32) + const1 = UOp.const(5) + const2 = UOp.const(10) + const3 = UOp.const(20) optimized_sink = apply_rewrite((const1 + const2 + const3).reduce(Ops.ADD)) expected_sum = 5 + 10 + 20 self.assertEqual(optimized_sink.arg, expected_sum) @unittest.skip("reduce is removed now") def test_full_graph_rewrite_reduction_with_unused_range(self): - const1 = UOp.const(15, dtypes.int32) - const2 = UOp.const(25, dtypes.int32) + const1 = UOp.const(15) + const2 = UOp.const(25) rng = UOp.range(10, idx=0) optimized_sink = apply_rewrite((const1 + const2).reduce(Ops.ADD, rng)) expected_sum = 10 * (15 + 25) @@ -89,7 +89,7 @@ class TestFoldingAndReduction(unittest.TestCase): @unittest.skip("currently failing") def test_full_graph_rewrite_simple_reduction_folding(self): simple_range = UOp.range(4, idx=0) - add_uop = simple_range + UOp.const(1, dtypes.int32) + add_uop = simple_range + UOp.const(1) optimized_sink = apply_rewrite(add_uop.reduce(Ops.ADD, simple_range)) expected_sum = sum(i + 1 for i in range(4)) self.assertEqual(optimized_sink.arg, expected_sum) @@ -128,9 +128,9 @@ class TestModuloAndDivisionFolding(unittest.TestCase): def test_graph_rewrite_div_folding_bug(self): lhs = UOp(Ops.ADD, src=( - UOp(Ops.STACK, arg=None, src=(UOp(Ops.SPECIAL, src=(UOp.const(32, dtypes.int),), arg='lidx0'),)*4), - UOp.const((0, 256, 512, 768), dtypes.int))) - rhs = UOp.const((2,)*4, dtypes.int) + UOp(Ops.STACK, arg=None, src=(UOp(Ops.SPECIAL, src=(UOp.const(32),), arg='lidx0'),)*4), + UOp.const((0, 256, 512, 768)))) + rhs = UOp.const((2,)*4) unopt = lhs 0, @@ -182,27 +182,27 @@ class TestEdgeCasesAndSpecialOperations(unittest.TestCase): class TestGEPAndVectorizeRewrite(unittest.TestCase): def test_gep_single_element_extraction(self): # GEP on a vector dtype to extract a single element - base_vector = UOp.const((1.0, 2.0, 3.0, 4.0), dtypes.float32) + base_vector = UOp.const((1.0, 2.0, 3.0, 4.0)) self.assertEqual(apply_rewrite(base_vector.index(2)).arg, 3.0) def test_gep_tuple_extraction(self): # GEP on a vector dtype to extract multiple elements as a vector - base_vector = UOp.const((1.0, 2.0, 3.0, 4.0), dtypes.float32) + base_vector = UOp.const((1.0, 2.0, 3.0, 4.0)) self.assertEqual(list(apply_rewrite_values(UOp.stack(*[base_vector.index(i) for i in (2, 3)]))), [3.0, 4.0]) def test_gep_on_const_stack(self): # GEP on a const STACK to extract a single element - const_stack = UOp.const((1.0, 2.0, 3.0, 4.0), dtypes.float32) + const_stack = UOp.const((1.0, 2.0, 3.0, 4.0)) self.assertEqual(apply_rewrite(const_stack.index(2)).arg, 3.0) def test_gep_tuple_on_const_stack(self): # GEP on a const STACK using a tuple to extract multiple elements - const_stack = UOp.const((7.0, 8.0, 9.0, 10.0), dtypes.float32) + const_stack = UOp.const((7.0, 8.0, 9.0, 10.0)) self.assertEqual(list(apply_rewrite_values(UOp.stack(*[const_stack.index(i) for i in (1, 3)]))), [8.0, 10.0]) def test_vectorize_multiple_elements(self): # Vectorizing multiple elements using GEP - base_vector = UOp.const((5.0, 10.0, 15.0, 20.0), dtypes.float32) + base_vector = UOp.const((5.0, 10.0, 15.0, 20.0)) vectorized_uop = UOp(Ops.STACK, src=tuple(base_vector.index(i) for i in range(4))) self.assertEqual(list(apply_rewrite_values(vectorized_uop)), [5.0, 10.0, 15.0, 20.0]) @@ -213,7 +213,7 @@ from tinygrad.uop.symbolic import symbolic_simple class TestBottomUpRewrite(unittest.TestCase): def test_const_folding(self): - a = UOp.const(5, dtypes.int) + a = UOp.const(5) ret = (a*3) + (a*7) gt = graph_rewrite(ret, symbolic_simple) ret = graph_rewrite(ret, symbolic_simple, bottom_up=True) @@ -305,7 +305,7 @@ class TestRecurse(unittest.TestCase): graph_rewrite(a, pm, bottom_up=True) def test_inf_loop(self): - a = UOp.const(3, dtypes.int) + a = UOp.const(3) pm = PatternMatcher([ (UPat(Ops.CONST, arg=3, name="x"), lambda x: x.replace(arg=4)), (UPat(Ops.CONST, arg=4, name="x"), lambda x: x.replace(arg=3)), @@ -314,7 +314,7 @@ class TestRecurse(unittest.TestCase): graph_rewrite(a, pm) def test_inf_loop_bottom_up(self): - a = UOp.const(3, dtypes.int) + a = UOp.const(3) pm = PatternMatcher([ (UPat(Ops.CONST, arg=3, name="x"), lambda x: x.replace(arg=4)), (UPat(Ops.CONST, arg=4, name="x"), lambda x: x.replace(arg=3)), @@ -325,8 +325,8 @@ class TestRecurse(unittest.TestCase): def bidir_append(ctx, x, b): ctx.append((x.arg if x.op is Ops.CONST else "+", b)) class TestBidirectional(unittest.TestCase): def test_simple(self): - a = UOp.const(1, dtypes.int) - b = UOp.const(2, dtypes.int) + a = UOp.const(1) + b = UOp.const(2) c = a + b pm = PatternMatcher([ (UPat(GroupOp.All, name="x"), lambda ctx,x: bidir_append(ctx, x, False)) ]) bpm = PatternMatcher([ (UPat(GroupOp.All, name="x"), lambda ctx,x: bidir_append(ctx, x, True)) ]) @@ -336,11 +336,11 @@ class TestBidirectional(unittest.TestCase): class TestStopEarly(unittest.TestCase): def test_stop_early(self): - a = UOp.const(3, dtypes.int) - b = UOp.const(4, dtypes.int) + a = UOp.const(3) + b = UOp.const(4) c = a+b - cn = UOp.const(7, dtypes.int) - d = UOp.const(2, dtypes.int) + cn = UOp.const(7) + d = UOp.const(2) def visit_const(c:UOp): print(f"visit {c.arg}") assert c.arg not in (3,4) @@ -376,7 +376,7 @@ class TestWalkRewrite(unittest.TestCase): def test_walk_topdown_no_fixed_point(self): """A bouncing pattern applies once and stops instead of looping.""" - a = UOp.const(3, dtypes.int) + a = UOp.const(3) pm = PatternMatcher([ (UPat(Ops.CONST, arg=3, name="x"), lambda x: x.replace(arg=4)), (UPat(Ops.CONST, arg=4, name="x"), lambda x: x.replace(arg=3)), @@ -384,7 +384,7 @@ class TestWalkRewrite(unittest.TestCase): with self.assertRaises(RuntimeError): graph_rewrite(a, pm, bottom_up=True) ret = graph_rewrite(a, pm, walk=True) - self.assertIs(ret, UOp.const(4, dtypes.int)) + self.assertIs(ret, UOp.const(4)) def test_walk_topdown_rewrites_children(self): a = UOp.variable('a', 0, 10) @@ -421,8 +421,8 @@ class TestWalkRewrite(unittest.TestCase): ctx.append(x.arg if x.op is Ops.CONST else x.op) return None pm = PatternMatcher([(UPat(GroupOp.All, name="x"), track_visit)]) - a = UOp.const(1, dtypes.int) - b = UOp.const(2, dtypes.int) + a = UOp.const(1) + b = UOp.const(2) graph_rewrite(a + b, pm, ctx=visited, walk=True) self.assertEqual(visited, [1, 2, Ops.ADD]) @@ -454,13 +454,13 @@ class TestWalkRewrite(unittest.TestCase): def test_walk_bottomup_no_fixed_point(self): """Bottom-up walk also applies once per node, no fixed-point iteration.""" - a = UOp.const(3, dtypes.int) + a = UOp.const(3) pm = PatternMatcher([ (UPat(Ops.CONST, arg=3, name="x"), lambda x: x.replace(arg=4)), (UPat(Ops.CONST, arg=4, name="x"), lambda x: x.replace(arg=3)), ]) ret = graph_rewrite(a, pm, bottom_up=True, walk=True) - self.assertIs(ret, UOp.const(4, dtypes.int)) + self.assertIs(ret, UOp.const(4)) def test_walk_bottomup_visit_order(self): """Bottom-up walk fires bpm before descending (pre-order).""" @@ -469,8 +469,8 @@ class TestWalkRewrite(unittest.TestCase): ctx.append(x.arg if x.op is Ops.CONST else x.op) return None pm = PatternMatcher([(UPat(GroupOp.All, name="x"), track_visit)]) - a = UOp.const(1, dtypes.int) - b = UOp.const(2, dtypes.int) + a = UOp.const(1) + b = UOp.const(2) graph_rewrite(a + b, pm, ctx=visited, bottom_up=True, walk=True) # bpm fires on each node before children: +, 1, 2 self.assertEqual(visited, [Ops.ADD, 1, 2]) @@ -497,8 +497,8 @@ class TestWalkRewrite(unittest.TestCase): return None bpm = PatternMatcher([(UPat(GroupOp.All, name="x"), bpm_visit)]) pm = PatternMatcher([(UPat(GroupOp.All, name="x"), pm_visit)]) - a = UOp.const(1, dtypes.int) - b = UOp.const(2, dtypes.int) + a = UOp.const(1) + b = UOp.const(2) graph_rewrite(a + b, pm, ctx=visited, bpm=bpm, walk=True) # bpm fires pre-order, pm fires post-order self.assertEqual(visited, [ @@ -518,14 +518,14 @@ class TestWalkRewrite(unittest.TestCase): return None bpm = PatternMatcher([(UPat(GroupOp.All, name="x"), bpm_match)]) pm = PatternMatcher([(UPat(GroupOp.All, name="x"), pm_match)]) - a = UOp.const(1, dtypes.int) - b = UOp.const(2, dtypes.int) + a = UOp.const(1) + b = UOp.const(2) ret = graph_rewrite(a + b, pm, ctx=visited, bpm=bpm, walk=True) # bpm matches const(1) and short-circuits it, so pm never fires on const(1) self.assertNotIn((1, "pm"), visited) # but pm still fires on const(2) and the rebuilt ADD self.assertIn((2, "pm"), visited) - self.assertIs(ret, UOp.const(10, dtypes.int) + b) + self.assertIs(ret, UOp.const(10) + b) if __name__ == '__main__': unittest.main() diff --git a/test/null/test_helpers.py b/test/null/test_helpers.py index de4c04c6f6..b311e27c53 100644 --- a/test/null/test_helpers.py +++ b/test/null/test_helpers.py @@ -297,10 +297,10 @@ class TestPolyN(unittest.TestCase): from tinygrad.dtype import dtypes from tinygrad.uop.ops import UOp from test.helpers import eval_uop - np.testing.assert_allclose(eval_uop(polyN(UOp.const(1.0, dtypes.float), [1.0, -2.0, 1.0])), 0.0) - np.testing.assert_allclose(eval_uop(polyN(UOp.const(2.0, dtypes.float), [1.0, -2.0, 1.0])), 1.0) - np.testing.assert_allclose(eval_uop(polyN(UOp.const(3.0, dtypes.float), [1.0, -2.0, 1.0])), 4.0) - np.testing.assert_allclose(eval_uop(polyN(UOp.const(4.0, dtypes.float), [1.0, -2.0, 1.0])), 9.0) + np.testing.assert_allclose(eval_uop(polyN(UOp.const(1.0).cast(dtypes.float), [1.0, -2.0, 1.0])), 0.0) + np.testing.assert_allclose(eval_uop(polyN(UOp.const(2.0).cast(dtypes.float), [1.0, -2.0, 1.0])), 1.0) + np.testing.assert_allclose(eval_uop(polyN(UOp.const(3.0).cast(dtypes.float), [1.0, -2.0, 1.0])), 4.0) + np.testing.assert_allclose(eval_uop(polyN(UOp.const(4.0).cast(dtypes.float), [1.0, -2.0, 1.0])), 9.0) class TestTimeToStr(unittest.TestCase): def test_seconds(self): self.assertEqual(" 10.01s ", time_to_str(10.01)) diff --git a/test/null/test_linearizer_failures.py b/test/null/test_linearizer_failures.py index a5df41f53f..620cd8d0f3 100644 --- a/test/null/test_linearizer_failures.py +++ b/test/null/test_linearizer_failures.py @@ -16,7 +16,7 @@ class TestLinearizerFailures(unittest.TestCase): c6 = c4.index(((((((c5//UOp.const(8))%UOp.const(8))*UOp.const(8))+(c5%UOp.const(8)))+(((c2*UOp.const(40))+(c5//UOp.const(64)))*UOp.const(64)))+(c1*UOp.const(81920)))) c7 = UOp.param(2, dtypes.float, (64,)) c8 = c7.index(c3) - c9 = ((((c6+(c8*UOp.const(-1.0, dtypes.float)))*(c6+(c8*UOp.const(-1.0, dtypes.float)))).reduce(c5, arg=Ops.ADD)*UOp.const(0.000390625, dtypes.float))+UOp.const(1e-05, dtypes.float)).sqrt().reciprocal() + c9 = ((((c6+(c8*UOp.const(-1.0)))*(c6+(c8*UOp.const(-1.0)))).reduce(c5, arg=Ops.ADD)*UOp.const(0.000390625))+UOp.const(1e-05)).sqrt().reciprocal() c10 = c0.index(c3).store(c9).end(c1, c2) ast = c10.sink(arg=KernelInfo()) to_program(ast, renderer=Device[Device.DEFAULT].renderer) diff --git a/test/null/test_microbenchmarks.py b/test/null/test_microbenchmarks.py index c09fa997a4..964b63771c 100644 --- a/test/null/test_microbenchmarks.py +++ b/test/null/test_microbenchmarks.py @@ -1,5 +1,5 @@ import unittest, time -from tinygrad import dtypes, Tensor, UOp, getenv +from tinygrad import Tensor, UOp, getenv from tinygrad.helpers import Profiling PYPROFILE = getenv("PYPROFILE") @@ -27,29 +27,29 @@ class TestBench(unittest.TestCase): print(f"{self._testMethodName:30s} {et*1e6/self.N:.2f} us") def test_uop_instant_creation(self): - for i in range(self.N): UOp.const(100+i, dtypes.int) + for i in range(self.N): UOp.const(100+i) def test_uop_list_creation(self): - [UOp.const(100+i, dtypes.int) for i in range(self.N)] + [UOp.const(100+i) for i in range(self.N)] def test_uop_add_2n(self): - a = UOp.const(2, dtypes.int) + a = UOp.const(2) for _ in range(self.N): a = a + a def test_uop_toposort(self): - a = UOp.const(0, dtypes.int) - for i in range(self.N): a = a + UOp.const(100+i, dtypes.int) + a = UOp.const(0) + for i in range(self.N): a = a + UOp.const(100+i) self.start_time() self.assertEqual(len(a.toposort()), 2*self.N+1) def test_uop_toposort_2n(self): - a = UOp.const(0, dtypes.int) + a = UOp.const(0) for _ in range(self.N): a = a + a self.start_time() self.assertEqual(len(a.toposort()), self.N+1) def test_uop_simplify(self): - a = UOp.const(2, dtypes.int) + a = UOp.const(2) for _ in range(self.N): (a+a).simplify() def test_uop_simplify_complex(self): @@ -68,7 +68,7 @@ class TestBench(unittest.TestCase): for _ in range(self.N): expr.simplify() def test_uop_chain_free(self): - a = UOp.const(2, dtypes.int) + a = UOp.const(2) for _ in range(self.N): a = a + a self.start_time() del a diff --git a/test/null/test_pattern_matcher.py b/test/null/test_pattern_matcher.py index 3ed8deb8e5..01378bfba3 100644 --- a/test/null/test_pattern_matcher.py +++ b/test/null/test_pattern_matcher.py @@ -5,9 +5,9 @@ from tinygrad.uop.ops import PatternMatcher, UPat class TestPatternMatcher(unittest.TestCase): def test_simple_match(self): - matcher = PatternMatcher([(UPat(Ops.CONST, name="x", dtype=dtypes.float), lambda x: x.rtag())]) - c1 = UOp.const(1.0, dtypes.float) - c2 = UOp.const(1, dtypes.int) + matcher = PatternMatcher([(UPat(Ops.CONST, name="x", dtype=dtypes.weakfloat), lambda x: x.rtag())]) + c1 = UOp.const(1.0) + c2 = UOp.const(1) self.assertEqual(matcher.rewrite(c1), c1.rtag()) self.assertEqual(matcher.rewrite(c2), None) @@ -61,7 +61,7 @@ class TestPatternMatcher(unittest.TestCase): def test_uop(self): matcher = PatternMatcher([(UPat(Ops.CONST, name="x"), lambda x: x.rtag())]) - c1 = UOp.const(1.0, dtypes.float) + c1 = UOp.const(1.0) c2 = UOp(Ops.ADD, src=(c1, c1)) self.assertEqual(matcher.rewrite(c1), c1.rtag()) self.assertEqual(matcher.rewrite(c2), None) @@ -70,7 +70,7 @@ class TestPatternMatcher(unittest.TestCase): matcher = PatternMatcher([(UPat((Ops.CONST, Ops.CAST), name="x"), lambda x: x.rtag())]) c1 = UOp.const(False) c2 = UOp(Ops.CAST, arg=dtypes.int, src=(c1,)) - c3 = UOp.const(1.0, dtypes.float) + c3 = UOp.const(1.0) c4 = UOp(Ops.ADD, src=(c3, c3)) self.assertEqual(matcher.rewrite(c1), c1.rtag()) self.assertEqual(matcher.rewrite(c2), c2.rtag()) @@ -82,11 +82,11 @@ class TestPatternMatcher(unittest.TestCase): (UPat(Ops.CONST, arg=False, name="x"), lambda x: x.rtag()), (UPat(Ops.MAX, name="x"), lambda x: x.rtag()), ]) - c1 = UOp.const(0.0, dtypes.float) + c1 = UOp.const(0.0) c2 = UOp.const(False) c3 = UOp(Ops.MAX, src=(c1, c1)) c4 = UOp(Ops.MUL, src=(c1, c1)) - c5 = UOp.const(-1, dtypes.int) + c5 = UOp.const(-1) self.assertEqual(matcher.rewrite(c1), c1.rtag()) self.assertEqual(matcher.rewrite(c2), c2.rtag()) self.assertEqual(matcher.rewrite(c3), c3.rtag()) @@ -98,9 +98,9 @@ class TestPatternMatcher(unittest.TestCase): (UPat(Ops.MUL, src=[UPat(Ops.CONST, name="c"), UPat(Ops.CONST, arg=2)], name="x"), lambda x,c: x.rtag() if c.arg in {1, -1} else None) ]) - y1 = UOp.const(1, dtypes.int) - y2 = UOp.const(2, dtypes.int) - y3 = UOp.const(-1, dtypes.int) + y1 = UOp.const(1) + y2 = UOp.const(2) + y3 = UOp.const(-1) c1 = UOp(Ops.MUL, src=(y1, y2)) c2 = UOp(Ops.MUL, src=(y2, y2)) c3 = UOp(Ops.MUL, src=(y3, y2)) @@ -114,26 +114,27 @@ class TestPatternMatcher(unittest.TestCase): def test_dup_name(self): matcher = PatternMatcher([(UPat(GroupOp.ALU, name="x", src=(UPat(Ops.CONST, name="y"), UPat(Ops.CONST, name="y"))), lambda x, y: x.rtag())]) - y1 = UOp.const(1.0, dtypes.float) - y2 = UOp.const(1.0, dtypes.float) + y1 = UOp.const(1.0) + y2 = UOp.const(1.0) c1 = UOp(Ops.ADD, src=(y1, y1)) c2 = UOp(Ops.ADD, src=(y1, y2)) self.assertEqual(matcher.rewrite(c1), c1.rtag()) self.assertEqual(matcher.rewrite(c2), c1.rtag()) def test_dtype(self): - matcher = PatternMatcher([(UPat(Ops.CONST, name="x", dtype=dtypes.float32), lambda x: x.rtag())]) - c1 = UOp.const(1.0, dtypes.float) - c2 = UOp.const(1.0, dtypes.float64) + # a concrete const dtype lives on the pair's CAST + matcher = PatternMatcher([(UPat(Ops.CAST, name="x", dtype=dtypes.float32), lambda x: x.rtag())]) + c1 = UOp.const(1.0).cast(dtypes.float32) + c2 = UOp.const(1.0).cast(dtypes.float64) self.assertEqual(matcher.rewrite(c1), c1.rtag()) self.assertEqual(matcher.rewrite(c2), None) def test_dtype_set(self): - matcher = PatternMatcher([(UPat(Ops.CONST, name="x", dtype={dtypes.float32, dtypes.float64}), lambda x: x.rtag())]) - c1 = UOp.const(1.0, dtypes.float) - c2 = UOp.const(1.0, dtypes.float64) - c3 = UOp.const(1.0, dtypes.float16) - c4 = UOp.const(1, dtypes.int) + matcher = PatternMatcher([(UPat(Ops.CAST, name="x", dtype={dtypes.float32, dtypes.float64}), lambda x: x.rtag())]) + c1 = UOp.const(1.0).cast(dtypes.float32) + c2 = UOp.const(1.0).cast(dtypes.float64) + c3 = UOp.const(1.0).cast(dtypes.float16) + c4 = UOp.const(1).cast(dtypes.int) self.assertEqual(matcher.rewrite(c1), c1.rtag()) self.assertEqual(matcher.rewrite(c2), c2.rtag()) self.assertEqual(matcher.rewrite(c3), None) @@ -141,8 +142,8 @@ class TestPatternMatcher(unittest.TestCase): def test_src_one(self): matcher = PatternMatcher([(UPat(GroupOp.ALU, name="x", src=(UPat(Ops.CONST), UPat(Ops.CONST))), lambda x: x.rtag())]) - c1 = UOp.const(1.0, dtypes.float) - c2 = UOp.const(2.0, dtypes.float) + c1 = UOp.const(1.0) + c2 = UOp.const(2.0) c3 = UOp(Ops.ADD, src=(c1,c2)) self.assertEqual(matcher.rewrite(c3), c3.rtag()) self.assertEqual(matcher.rewrite(c2), None) @@ -158,8 +159,8 @@ class TestPatternMatcher(unittest.TestCase): def test_src_permutations(self): matcher = PatternMatcher([(UPat(GroupOp.ALU, name="x", src=[UPat(Ops.CONST), UPat(GroupOp.ALU)]), lambda x: x.rtag())]) - c1 = UOp.const(1.0, dtypes.float) - c2 = UOp.const(2.0, dtypes.float) + c1 = UOp.const(1.0) + c2 = UOp.const(2.0) c3 = UOp(Ops.ADD, src=(c1,c2)) c4 = UOp(Ops.ADD, src=(c3,c2)) c5 = UOp(Ops.ADD, src=(c2,c3)) @@ -171,8 +172,8 @@ class TestPatternMatcher(unittest.TestCase): def test_src_repeat(self): matcher = PatternMatcher([(UPat(GroupOp.ALU, name="x", src=UPat(Ops.CONST)), lambda x: x.rtag())]) - c1 = UOp.const(1.0, dtypes.float) - c2 = UOp.const(2.0, dtypes.float) + c1 = UOp.const(1.0) + c2 = UOp.const(2.0) c3 = UOp(Ops.ADD, src=(c1,c2)) c4 = UOp(Ops.ADD, src=(c2,c3)) self.assertEqual(matcher.rewrite(c3), c3.rtag()) @@ -180,9 +181,9 @@ class TestPatternMatcher(unittest.TestCase): def test_allow_len(self): matcher = PatternMatcher([(UPat(Ops.MULACC, name="x", src=(UPat(Ops.CONST),), allow_any_len=True), lambda x: x.rtag())]) - c1 = UOp.const(1.0, dtypes.float) - c2 = UOp.const(2.0, dtypes.float) - c3 = UOp.const(3.0, dtypes.float) + c1 = UOp.const(1.0) + c2 = UOp.const(2.0) + c3 = UOp.const(3.0) c4 = UOp(Ops.EXP2, src=(c1,)) c5 = UOp(Ops.ADD, src=(c1,c2)) c6 = UOp(Ops.MULACC, src=(c1,c2,c3)) @@ -191,8 +192,8 @@ class TestPatternMatcher(unittest.TestCase): self.assertEqual(matcher.rewrite(c6), c6.rtag()) def test_deep_src_permutations(self): - c1 = UOp.const(1.0, dtypes.float) - c2 = UOp.const(2.0, dtypes.float) + c1 = UOp.const(1.0) + c2 = UOp.const(2.0) u1 = (c1 + c2) + c1 u2 = (c2 + c1) + c1 matcher = PatternMatcher([ diff --git a/test/null/test_schedule.py b/test/null/test_schedule.py index 7ca7d900ff..2f4cbae50a 100644 --- a/test/null/test_schedule.py +++ b/test/null/test_schedule.py @@ -1640,11 +1640,11 @@ class TestSchedule(unittest.TestCase): self.assertEqual(GlobalCounters.mem_used-base, 0) def test_const_schedule(self): - constv = Tensor.empty(2, 2).uop.const_like(10) + constv = Tensor.empty(2, 2).const_like(10).uop check_schedule(constv, 0) def test_const_schedule_contig(self): - constv = Tensor.empty(2, 2).uop.const_like(10).contiguous() + constv = Tensor.empty(2, 2).const_like(10).uop.contiguous() check_schedule(constv, 0) def test_advanced_simple_indexing_combined(self): diff --git a/test/null/test_simplify_valid_idx.py b/test/null/test_simplify_valid_idx.py index 6760e4402d..20f378000d 100644 --- a/test/null/test_simplify_valid_idx.py +++ b/test/null/test_simplify_valid_idx.py @@ -577,7 +577,7 @@ class TestRangeShrink(unittest.TestCase): # emulates mask.where(x.pad_to(mask.shape), Invalid): range should shrink accordingly from tinygrad.dtype import Invalid r = Range(0, 204) - x = (r < 4).where(UOp.const(1, dtypes.float), Invalid) + x = (r < 4).where(UOp.const(1.0), Invalid) ranges = self.get_ranges(UOp.param(0, dtypes.float, (204,)).index(r).store((r < 4).where(x, Invalid)).sink()) self.assertEqual(len(ranges), 1) self.assertEqual(ranges[0].src[0].arg, 4) @@ -586,7 +586,7 @@ class TestRangeShrink(unittest.TestCase): # above, but flipped from tinygrad.dtype import Invalid r = Range(0, 204) - x = (r < 4).where(UOp.const(1, dtypes.float), Invalid) + x = (r < 4).where(UOp.const(1.0), Invalid) ranges = self.get_ranges(UOp.param(0, dtypes.float, (204,)).index(r).store((r >= 4).where(Invalid, x)).sink()) self.assertEqual(len(ranges), 1) self.assertEqual(ranges[0].src[0].arg, 4) diff --git a/test/null/test_symbolic_failures.py b/test/null/test_symbolic_failures.py index 9fc6b1cf57..efe827833d 100644 --- a/test/null/test_symbolic_failures.py +++ b/test/null/test_symbolic_failures.py @@ -1,5 +1,6 @@ import unittest from tinygrad import Variable +from tinygrad.uop.ops import UOp class TestFuzzFailure(unittest.TestCase): @@ -8,7 +9,7 @@ class TestFuzzFailure(unittest.TestCase): v2=Variable('v2', 0, 2) v3=Variable('v3', 0, 1) expr = (((((((((((((((((((((((0//4)%2)//8)+-2)+-4)+-3)+v1)+-4)+v2)+-2)+v3)+v2)//3)%7)*1)//2)+v2)*-1)+2)+1)+0)+-3)+v3) - v1_val, v2_val, v3_val = v1.const_like(8), v2.const_like(0), v3.const_like(0) + v1_val, v2_val, v3_val = UOp.const(8), UOp.const(0), UOp.const(0) num = expr.simplify().substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify() rn = expr.substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify() self.assertEqual(num, rn) @@ -19,7 +20,7 @@ class TestFuzzFailure(unittest.TestCase): v3=Variable('v3', 0, 3) expr = (((((((((((((((((((((((((0*4)//5)*2)*-1)*-2)+-4)*4)*2)*3)*4)+-4)*4)+v2)+v2)+v3)//3)+v2)+v1)//9)+3)+1)//1)+-4)//4)*2) expr = (((((v1+(v2+(((v3+(v2*2))+1)//3)))+4)//9)+-57)//(9*4)) - v1_val, v2_val, v3_val = v1.const_like(6), v2.const_like(0), v3.const_like(0) + v1_val, v2_val, v3_val = UOp.const(6), UOp.const(0), UOp.const(0) num = expr.simplify().substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify() rn = expr.substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify() self.assertEqual(num, rn) @@ -29,7 +30,7 @@ class TestFuzzFailure(unittest.TestCase): v2=Variable('v2', 0, 1) v3=Variable('v3', 0, 2) expr = (((((((((((((((((((0//2)//3)+v3)+0)+-4)*-2)*-2)+-1)+2)+3)+v3)+0)//8)*-3)+0)*-2)*-4)*-2)//5) - v1_val, v2_val, v3_val = v1.const_like(0), v2.const_like(0), v3.const_like(0) + v1_val, v2_val, v3_val = UOp.const(0), UOp.const(0), UOp.const(0) num = expr.simplify().substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify() rn = expr.substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify() self.assertEqual(num, rn) @@ -39,7 +40,7 @@ class TestFuzzFailure(unittest.TestCase): v2=Variable('v2', 0, 3) v3=Variable('v3', 0, 4) expr = (((((((((((((((((((((((((((((0*-2)+0)*-1)//9)//6)//8)+v1)*-4)+v2)//4)//8)+4)*3)+v1)+v3)//8)//7)+4)+v3)*-4)+1)+v1)*3)+4)*2)//5)//2)//3)*-4) - v1_val, v2_val, v3_val = v1.const_like(2), v2.const_like(0), v3.const_like(2) + v1_val, v2_val, v3_val = UOp.const(2), UOp.const(0), UOp.const(2) num = expr.simplify().substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify() rn = expr.substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify() self.assertEqual(num, rn) @@ -49,7 +50,7 @@ class TestFuzzFailure(unittest.TestCase): v2=Variable('v2', 0, 1) v3=Variable('v3', 0, 3) expr = ((((((((((((((0+v2)+v1)*0)+v2)//1)//7)+-2)+v2)+v1)*4)+-3)//5)+v2)+1) - v1_val, v2_val, v3_val = v1.const_like(0), v2.const_like(0), v3.const_like(0) + v1_val, v2_val, v3_val = UOp.const(0), UOp.const(0), UOp.const(0) num = expr.simplify().substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify() rn = expr.substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify() self.assertEqual(num, rn) @@ -60,7 +61,7 @@ class TestFuzzFailure(unittest.TestCase): v3=Variable('v3', 0, 128) expr = (((((((((((((((((((((((((((((0//3)+4)+v1)//2)+-1)//1)*1)*-1)*4)//5)+v1)//6)+v1)*-1)+-4)+v2)+-2)*-3)+v3)+-4)+-2)*-1)//8)//4)*-4)+3)+v3)* -2)+v2) - v1_val, v2_val, v3_val = v1.const_like(8), v2.const_like(3), v3.const_like(2) + v1_val, v2_val, v3_val = UOp.const(8), UOp.const(3), UOp.const(2) num = expr.simplify().substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify() rn = expr.substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify() self.assertEqual(num, rn) @@ -70,7 +71,7 @@ class TestFuzzFailure(unittest.TestCase): v2=Variable('v2', 0, 5) v3=Variable('v3', 0, 128) expr = (((((((((((((((((((((((((((((0+v2)*-4)+0)//9)+-4)*-2)*3)*4)//9)+v3)+v1)//4)+v1)+v3)+-1)*4)//4)+v2)//7)//3)+v1)+v2)+v3)+1)*2)//4)*3)+-1)*1) - v1_val, v2_val, v3_val = v1.const_like(0), v2.const_like(2), v3.const_like(65) + v1_val, v2_val, v3_val = UOp.const(0), UOp.const(2), UOp.const(65) num = expr.simplify().substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify() rn = expr.substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify() self.assertEqual(num, rn) @@ -80,7 +81,7 @@ class TestFuzzFailure(unittest.TestCase): v2=Variable('v2', 0, 8) v3=Variable('v3', 0, 9) expr = (((((((0+-1)+2)+v1)*-2)//3)+v1)*-4) - v1_val, v2_val, v3_val = v1.const_like(0), v2.const_like(0), v3.const_like(0) + v1_val, v2_val, v3_val = UOp.const(0), UOp.const(0), UOp.const(0) num = expr.simplify().substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify() rn = expr.substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify() self.assertEqual(num, rn) @@ -90,7 +91,7 @@ class TestFuzzFailure(unittest.TestCase): v2=Variable('v2', 0, 1) v3=Variable('v3', 0, 8) expr = (((((((((((((((((((((((((((((0*-2)//1)+3)*-2)+-3)*-4)*1)+v1)+0)%2)%8)%9)+v2)%9)+-4)//4)+-1)*-2)+0)+v1)+v1)+3)+v1)+4)+-4)+0)*2)+-3)%6) - v1_val, v2_val, v3_val = v1.const_like(0), v2.const_like(1), v3.const_like(0) + v1_val, v2_val, v3_val = UOp.const(0), UOp.const(1), UOp.const(0) num = expr.simplify().substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify() rn = expr.substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify() self.assertEqual(num, rn) @@ -101,7 +102,7 @@ class TestFuzzFailure(unittest.TestCase): v3=Variable("v3", 0, 32) x5 = (v1 <= 9).where(v1 * -4 - 4, v1 // 9) // 9 expr = ((x5 >= -4).where(x5, (v2 % 3 + v2) // 5) * -1).maximum(((v1 * -2) % 6 + v3 % 1) * -1) * -1 - v1_val, v2_val, v3_val = v1.const_like(9), v2.const_like(0), v3.const_like(0) + v1_val, v2_val, v3_val = UOp.const(9), UOp.const(0), UOp.const(0) num = expr.simplify().substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify() rn = expr.substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify() self.assertEqual(num, rn) @@ -111,7 +112,7 @@ class TestFuzzFailure(unittest.TestCase): v2=Variable("v2", 0, 128) v3=Variable("v3", 0, 5) expr = (((v2 * 0).maximum(8) - v2 * 2) % 5 + v1 // 6 + v1 + 5) % 5 - v1_val, v2_val, v3_val = v1.const_like(0), v2.const_like(7), v3.const_like(0) + v1_val, v2_val, v3_val = UOp.const(0), UOp.const(7), UOp.const(0) num = expr.simplify().substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify() rn = expr.substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify() self.assertEqual(num, rn) diff --git a/test/null/test_tensor_uop_mixin.py b/test/null/test_tensor_uop_mixin.py index 983b5c6940..fc63e08e1a 100644 --- a/test/null/test_tensor_uop_mixin.py +++ b/test/null/test_tensor_uop_mixin.py @@ -60,7 +60,7 @@ class TestTensorUOpClone(unittest.TestCase): t = _t(3, 4).float() self.assertIs(_strip_unique(t.clone().uop), _strip_unique(t.uop.clone())) def test_clone_deviceless_const(self): - u = UOp.const(2.0, dtypes.float) + u = UOp.const(2.0) self.assertIs(_strip_unique(Tensor(u).clone().uop), _strip_unique(u.clone())) class TestTensorUOpGradient(unittest.TestCase): diff --git a/test/null/test_transcendental_helpers.py b/test/null/test_transcendental_helpers.py index 15a09bbe8a..1653b176c7 100644 --- a/test/null/test_transcendental_helpers.py +++ b/test/null/test_transcendental_helpers.py @@ -10,7 +10,7 @@ class TestTranscendentalFunctions(unittest.TestCase): # TODO: Test constant input when constant folding is fixed (or maybe test both variants) # Load input value from a buffer to prevent constant folding input_buf = UOp.param(1, dtypes.double, (1,)) - loaded_value = input_buf.index(UOp.const(0, dtypes.int)).load() + loaded_value = input_buf.index(UOp.const(0)).load() def eval_payne_hanek_reduction(v:float) -> tuple[float, int]: return tuple(eval_uop(u, [(dtypes.float64, [v])]) for u in payne_hanek_reduction(loaded_value)) @@ -27,48 +27,48 @@ class TestTranscendentalFunctions(unittest.TestCase): np.testing.assert_equal(q, 4) def test_cody_waite_reduction(self): - r, q = (eval_uop(u) for u in cody_waite_reduction(UOp.const(12 * math.pi + 0.1, dtypes.float64))) + r, q = (eval_uop(u) for u in cody_waite_reduction(UOp.const(12 * math.pi + 0.1).cast(dtypes.float64))) np.testing.assert_allclose(r, 0.1) np.testing.assert_equal(q, 12) def test_frexp(self): for x in (1, -1): - mantissa, exponent = (eval_uop(u) for u in frexp(UOp.const(x, dtypes.float64))) + mantissa, exponent = (eval_uop(u) for u in frexp(UOp.const(float(x)).cast(dtypes.float64))) np.testing.assert_equal(mantissa, 0.5) np.testing.assert_equal(exponent, 1) for x in (2, -2): - mantissa, exponent = (eval_uop(u) for u in frexp(UOp.const(2.0, dtypes.float64))) + mantissa, exponent = (eval_uop(u) for u in frexp(UOp.const(2.0).cast(dtypes.float64))) np.testing.assert_equal(mantissa, 0.5) np.testing.assert_equal(exponent, 2) - mantissa, exponent = (eval_uop(u) for u in frexp(UOp.const(5.0, dtypes.float64))) + mantissa, exponent = (eval_uop(u) for u in frexp(UOp.const(5.0).cast(dtypes.float64))) np.testing.assert_equal(mantissa, 0.625) np.testing.assert_equal(exponent, 3) - mantissa, exponent = (eval_uop(u) for u in frexp(UOp.const(1000.0, dtypes.float64))) + mantissa, exponent = (eval_uop(u) for u in frexp(UOp.const(1000.0).cast(dtypes.float64))) np.testing.assert_allclose(mantissa, 0.9765625) np.testing.assert_equal(exponent, 10) def test_rintk(self): - np.testing.assert_allclose(eval_uop(rintk(UOp.const(0.0, dtypes.float))), 0) - np.testing.assert_allclose(eval_uop(rintk(UOp.const(5.0, dtypes.float))), 5) - np.testing.assert_allclose(eval_uop(rintk(UOp.const(5.5, dtypes.float))), 6) - np.testing.assert_allclose(eval_uop(rintk(UOp.const(5.999, dtypes.float))), 6) - np.testing.assert_allclose(eval_uop(rintk(UOp.const(-5.0, dtypes.float))), -5) - np.testing.assert_allclose(eval_uop(rintk(UOp.const(-5.5, dtypes.float))), -6) - np.testing.assert_allclose(eval_uop(rintk(UOp.const(-5.999, dtypes.float))), -6) + np.testing.assert_allclose(eval_uop(rintk(UOp.const(0.0).cast(dtypes.float))), 0) + np.testing.assert_allclose(eval_uop(rintk(UOp.const(5.0).cast(dtypes.float))), 5) + np.testing.assert_allclose(eval_uop(rintk(UOp.const(5.5).cast(dtypes.float))), 6) + np.testing.assert_allclose(eval_uop(rintk(UOp.const(5.999).cast(dtypes.float))), 6) + np.testing.assert_allclose(eval_uop(rintk(UOp.const(-5.0).cast(dtypes.float))), -5) + np.testing.assert_allclose(eval_uop(rintk(UOp.const(-5.5).cast(dtypes.float))), -6) + np.testing.assert_allclose(eval_uop(rintk(UOp.const(-5.999).cast(dtypes.float))), -6) def test_pow2if(self): - np.testing.assert_allclose(eval_uop(pow2if(UOp.const(0, dtypes.int), dtypes.float)), 1.0) - np.testing.assert_allclose(eval_uop(pow2if(UOp.const(1, dtypes.int), dtypes.float)), 2.0) - np.testing.assert_allclose(eval_uop(pow2if(UOp.const(2, dtypes.int), dtypes.float)), 4.0) - np.testing.assert_allclose(eval_uop(pow2if(UOp.const(10, dtypes.int), dtypes.float)), 1024.0) - np.testing.assert_allclose(eval_uop(pow2if(UOp.const(63, dtypes.int), dtypes.float)), 2**63) - np.testing.assert_allclose(eval_uop(pow2if(UOp.const(-1, dtypes.int), dtypes.float)), 0.5) - np.testing.assert_allclose(eval_uop(pow2if(UOp.const(-2, dtypes.int), dtypes.float)), 0.25) - np.testing.assert_allclose(eval_uop(pow2if(UOp.const(-10, dtypes.int), dtypes.float)), 2**-10) - np.testing.assert_allclose(eval_uop(pow2if(UOp.const(-63, dtypes.int), dtypes.float)), 2**-63) + np.testing.assert_allclose(eval_uop(pow2if(UOp.const(0).cast(dtypes.int), dtypes.float)), 1.0) + np.testing.assert_allclose(eval_uop(pow2if(UOp.const(1).cast(dtypes.int), dtypes.float)), 2.0) + np.testing.assert_allclose(eval_uop(pow2if(UOp.const(2).cast(dtypes.int), dtypes.float)), 4.0) + np.testing.assert_allclose(eval_uop(pow2if(UOp.const(10).cast(dtypes.int), dtypes.float)), 1024.0) + np.testing.assert_allclose(eval_uop(pow2if(UOp.const(63).cast(dtypes.int), dtypes.float)), 2**63) + np.testing.assert_allclose(eval_uop(pow2if(UOp.const(-1).cast(dtypes.int), dtypes.float)), 0.5) + np.testing.assert_allclose(eval_uop(pow2if(UOp.const(-2).cast(dtypes.int), dtypes.float)), 0.25) + np.testing.assert_allclose(eval_uop(pow2if(UOp.const(-10).cast(dtypes.int), dtypes.float)), 2**-10) + np.testing.assert_allclose(eval_uop(pow2if(UOp.const(-63).cast(dtypes.int), dtypes.float)), 2**-63) if __name__ == '__main__': unittest.main() diff --git a/test/null/test_uop_graph.py b/test/null/test_uop_graph.py index 9ec5b0ddd3..48471de469 100644 --- a/test/null/test_uop_graph.py +++ b/test/null/test_uop_graph.py @@ -7,9 +7,9 @@ from tinygrad.uop.symbolic import sym from test.helpers import to_uops_list simple_pm = PatternMatcher([ - (UPat.cvar('x', dtypes.int), lambda x: UOp.const(1.0, dtypes.float) + UOp.const(2.0, dtypes.float)), - (UPat.cvar('x') + UPat.cvar('y'), lambda x,y: UOp.const(x.arg+y.arg, dtypes.float)), - (UPat.cvar('x') * UPat.cvar('y') * UPat.cvar('z'), lambda x,y,z: UOp.const(x.arg*y.arg*z.arg, dtypes.float)), + (UPat.cvar('x', dtypes.weakint), lambda x: UOp.const(1.0) + UOp.const(2.0)), + (UPat.cvar('x') + UPat.cvar('y'), lambda x,y: UOp.const(x.arg+y.arg)), + (UPat.cvar('x') * UPat.cvar('y') * UPat.cvar('z'), lambda x,y,z: UOp.const(x.arg*y.arg*z.arg)), ((UPat.var('x') + UPat.cvar('c1')) + UPat.cvar('c2'), lambda x,c1,c2: x + (c1.arg+c2.arg)), ]) @@ -27,15 +27,15 @@ class TestGraphRewriteConst(unittest.TestCase): self.assertEqual(ret.arg, 1) def test_add_const(self): - v1 = UOp.const((0,1,2), dtypes.int) - v2 = UOp.const((5,6,7), dtypes.int) + v1 = UOp.const((0,1,2)) + v2 = UOp.const((5,6,7)) ret = graph_rewrite(v1+v2, sym) self.assertEqual(ret.op, Ops.STACK) self.assertEqual(const_values(ret), (5,7,9)) def test_add_const_lose_v(self): - v1 = UOp.const((0,1,2), dtypes.int) - v2 = UOp.const((2,1,0), dtypes.int) + v1 = UOp.const((0,1,2)) + v2 = UOp.const((2,1,0)) ret = graph_rewrite(v1+v2, sym) self.assertEqual(ret.op, Ops.STACK) self.assertEqual(const_values(ret), (2,2,2)) @@ -103,62 +103,62 @@ class TestGraphRewrite(unittest.TestCase): # NOTE: this shows why we can't have a UOp in arg @unittest.expectedFailure def test_no_dedup_args(self): - a1 = UOp.variable("a1", UOp.const(0, dtypes.int), UOp.const(11, dtypes.int), dtypes.int) - a2 = UOp.variable("a2", UOp.const(0, dtypes.int), UOp.const(11, dtypes.int), dtypes.int) + a1 = UOp.variable("a1", UOp.const(0), UOp.const(11), dtypes.int) + a2 = UOp.variable("a2", UOp.const(0), UOp.const(11), dtypes.int) sink = a1.sink(a2) variables = [x for x in graph_rewrite(sink, PatternMatcher([])).toposort() if x.op is Ops.PARAM and x.addrspace is AddrSpace.ALU] self.assertEqual(len(variables), 1) def test_simple(self): - c1 = UOp.const(1.0, dtypes.float) - c2 = UOp.const(2.0, dtypes.float) + c1 = UOp.const(1.0) + c2 = UOp.const(2.0) nout = graph_rewrite(c1+c2, simple_pm) self.assertEqual(nout.op, Ops.CONST) self.assertEqual(nout.arg, 3.0) def test_depth_2_late(self): - c1 = UOp.const(1.0, dtypes.float) - c2 = UOp.const(2.0, dtypes.float) - c3 = UOp.const(3.0, dtypes.float) + c1 = UOp.const(1.0) + c2 = UOp.const(2.0) + c3 = UOp.const(3.0) nout = graph_rewrite(c1*c2*(c3+c3), simple_pm) self.assertEqual(nout.op, Ops.CONST) self.assertEqual(nout.arg, 12.0) def test_double(self): - c1 = UOp.const(1.0, dtypes.float) - c2 = UOp.const(2.0, dtypes.float) - c3 = UOp.const(3.0, dtypes.float) + c1 = UOp.const(1.0) + c2 = UOp.const(2.0) + c3 = UOp.const(3.0) nout = graph_rewrite(c1+c2+c3, simple_pm) self.assertEqual(nout.op, Ops.CONST) self.assertEqual(nout.arg, 6.0) def test_triple(self): - c1 = UOp.const(1.0, dtypes.float) - c2 = UOp.const(2.0, dtypes.float) - c3 = UOp.const(3.0, dtypes.float) - c4 = UOp.const(4.0, dtypes.float) + c1 = UOp.const(1.0) + c2 = UOp.const(2.0) + c3 = UOp.const(3.0) + c4 = UOp.const(4.0) nout = graph_rewrite(c1+c2+c3+c4, simple_pm) self.assertEqual(nout.op, Ops.CONST) self.assertEqual(nout.arg, 10.0) def test_diamond(self): - c1 = UOp.const(1.0, dtypes.float) - c2 = UOp.const(2.0, dtypes.float) - c3 = UOp.const(3.0, dtypes.float) + c1 = UOp.const(1.0) + c2 = UOp.const(2.0) + c3 = UOp.const(3.0) nout = graph_rewrite((c1+c2)+(c1+c3), simple_pm) self.assertEqual(nout.op, Ops.CONST) self.assertEqual(nout.arg, 7.0) def test_magic_4(self): - c1 = UOp.const(4.0, dtypes.int) + c1 = UOp.const(4) nout = graph_rewrite(c1, simple_pm) self.assertEqual(nout.op, Ops.CONST) self.assertEqual(nout.arg, 3.0) def test_depth_2_fold(self): v = UOp.variable("v", 0, 1, dtypes.float) - c1 = UOp.const(1.0, dtypes.float) - c2 = UOp.const(2.0, dtypes.float) + c1 = UOp.const(1.0) + c2 = UOp.const(2.0) nout = graph_rewrite(v+c1+c2, simple_pm) self.assertEqual(nout.op, Ops.ADD) self.assertEqual(nout.src[0].op, Ops.PARAM) @@ -174,14 +174,14 @@ class TestGraphRewrite(unittest.TestCase): a = UOp.variable('a', 0, 1) tst = (2+a).simplify() self.assertIs(tst.src[0], a) - self.assertIs(tst.src[1], a.const_like(2)) + self.assertIs(tst.src[1], UOp.const(2)) def test_consts_go_last(self): a = UOp.variable('a', 0, 1) b = UOp.variable('b', 0, 1) c = UOp.variable('c', 0, 1) d = UOp.variable('d', 0, 1) - outs = [2+a, 2+a+d+3+b+c+4, a.const_like(2)+a, (4+d)+c+(2+a)+b] + outs = [2+a, 2+a+d+3+b+c+4, UOp.const(2)+a, (4+d)+c+(2+a)+b] for out in outs: sink = graph_rewrite(out, sym) print(sink.render()) @@ -243,7 +243,7 @@ class TestUOpGraph(unittest.TestCase): @unittest.expectedFailure def test_const_shape_change_bitcast(self): - bf = UOp.const(0x3F, dtypes.uint8) + bf = UOp.const(0x3F).cast(dtypes.uint8) out = bf.bitcast(dtypes.half) uops = to_uops_list([out]) self.assertEqual(len(uops), 2) # +1 for SINK @@ -259,7 +259,7 @@ class TestUOpGraph(unittest.TestCase): @unittest.skip("this test isn't valid uops") def test_noop_vectorize_fold(self): d0 = UOp.param(0, dtypes.float, (1,)) - idx = UOp.const(0, dtypes.int) + idx = UOp.const(0) ld = d0.load(idx, dtype=dtypes.float) vec = UOp(Ops.STACK, dtypes.float, (ld,)) x = vec.index(0) @@ -273,7 +273,7 @@ class TestUOpGraph(unittest.TestCase): d0 = UOp.param(0, dtypes.float, (1,)) d1 = UOp.param(1, dtypes.float, (1,)) d2 = UOp.param(2, dtypes.float, (1,)) - idx = UOp.const(0, dtypes.int) + idx = UOp.const(0) def _test_vec(geps, count=4): vec = UOp(Ops.STACK, dtypes.float, geps) out = d0.index(idx).store(vec) @@ -320,7 +320,7 @@ class TestUOpGraph(unittest.TestCase): def test_cast_alu_fold(self): d0 = UOp.param(0, dtypes.bool, (1,)) d1 = UOp.param(1, dtypes.int, (1,)) - idx = UOp.const(0, dtypes.int) + idx = UOp.const(0) ld = d1.index(idx) alu = (ld<1).cast(dtypes.bool) out = d0.index(idx).store(alu) @@ -353,7 +353,7 @@ class TestUOpGraph(unittest.TestCase): def test_bitcast_to_same_dtype_fold(self): for dt in dtypes.ints + dtypes.floats + (dtypes.bool,): d0 = UOp.param(0, dt, (1,)) - v = d0.index(UOp.const(0, dtypes.int)) + v = d0.index(UOp.const(0)) uops = to_uops_list([v.bitcast(dt)]) self.assertEqual(len([x for x in uops if x.op is Ops.BITCAST and x.dtype is dt]), 0, f"dtype = {dt}") @@ -400,7 +400,7 @@ class TestUOpGraph(unittest.TestCase): ridx0 = UOp.range(100, 0) d0 = UOp.param(0, dtypes.float, (100,)) ld = d0.index(ridx0.valid(ridx0<50)) - w = ((ridx0<50) & (ridx0>30)).where(ld, UOp.const(0, dtypes.float)).cast(dtypes.half) + w = ((ridx0<50) & (ridx0>30)).where(ld, UOp.const(0.0)).cast(dtypes.half) out = UOp.param(1, dtypes.half, (100,)) uops = to_uops_list([out.index(ridx0).store(w)]) for u in uops: @@ -410,7 +410,7 @@ class TestUOpGraph(unittest.TestCase): ridx0 = UOp.range(100, 0) d0 = UOp.param(0, dtypes.float, (100,)) ld = d0.index(ridx0.valid(ridx0<50)) - w = ((ridx0<50) & (ridx0>30)).where(UOp.const(0, dtypes.float), ld).cast(dtypes.half) + w = ((ridx0<50) & (ridx0>30)).where(UOp.const(0.0), ld).cast(dtypes.half) out = UOp.param(1, dtypes.half, (100,)) uops = to_uops_list([out.index(ridx0).store(w)]) for u in uops: @@ -470,7 +470,7 @@ class TestUOpGraph(unittest.TestCase): glbl0 = UOp.param(0, dtypes.int, (1,)) glbl1 = UOp.param(1, dtypes.int, (1,)) glbl2 = UOp.param(2, dtypes.int, (1,)) - idx = UOp.const(0, dtypes.int) + idx = UOp.const(0) ld0 = glbl1.index(UOp.invalid()) ld1 = glbl2.index(idx.valid(UOp.const(True))) uops = to_uops_list([glbl0.index(idx).store(ld1+ld0)]) @@ -492,8 +492,8 @@ class TestUOpGraph(unittest.TestCase): def test_fold_gated_store(self): glbl = UOp.param(0, dtypes.int, (1,)) - idx0 = UOp.const(0, dtypes.int) - val = UOp.const(42, dtypes.int) + idx0 = UOp.const(0) + val = UOp.const(42) st0 = glbl.index(UOp.invalid()).store(val) st1 = glbl.index(idx0.valid(UOp.const(True))).store(val) uops = to_uops_list([st0, st1]) @@ -503,9 +503,9 @@ class TestUOpGraph(unittest.TestCase): @unittest.skip("this is a uop type error") def test_asserts_bad_gate(self): glbl0 = UOp.param(0, dtypes.int, (1,)) - idx = UOp.const(0, dtypes.int) - bad_gate = UOp.const(1, dtypes.int) - with self.assertRaises(AssertionError): to_uops_list([UOp(Ops.STORE, src=(glbl0, idx, UOp.const(42, dtypes.int), bad_gate))]) + idx = UOp.const(0) + bad_gate = UOp.const(1) + with self.assertRaises(AssertionError): to_uops_list([UOp(Ops.STORE, src=(glbl0, idx, UOp.const(42), bad_gate))]) def test_after_end(self): r = UOp.range(10, 0) @@ -513,7 +513,7 @@ class TestUOpGraph(unittest.TestCase): c = r + 1 self.assertIn(r, c.ranges) - e = UOp.const(1, dtypes.int).end(r) + e = UOp.const(1).end(r) self.assertNotIn(r, e.ranges) a = c.after(e) @@ -563,7 +563,7 @@ class TestConstBufferize(unittest.TestCase): CONST doesn't depend on ranges (constant is same value everywhere). """ from tinygrad.schedule.rangeify import pm_const_buffer_folding, BufferizeOpts - c = UOp.const(42.0, dtypes.float) + c = UOp.const(42.0) r1 = UOp.range(3, 0) bufferize_with_range = UOp(Ops.STAGE, src=(c, r1), arg=BufferizeOpts(device="CPU")) self.assertEqual(len(bufferize_with_range.src), 2) # const + 1 range @@ -571,13 +571,13 @@ class TestConstBufferize(unittest.TestCase): result = graph_rewrite(bufferize_with_range, pm_const_buffer_folding, name='test') # BUFFERIZE should be removed, result is const broadcast to shape self.assertNotEqual(result.op, Ops.STAGE) - const_vals = [u.arg for u in result.toposort() if u.op is Ops.CONST and u.dtype == dtypes.float] + const_vals = [u.arg for u in result.toposort() if u.op is Ops.CONST and u.dtype is dtypes.weakfloat] self.assertIn(42.0, const_vals) def test_const_bufferize_with_multiple_ranges(self): """Test CONST.BUFFERIZE with multiple ranges is also folded.""" from tinygrad.schedule.rangeify import pm_const_buffer_folding, BufferizeOpts - c = UOp.const(3.14, dtypes.float) + c = UOp.const(3.14) r1 = UOp.range(3, 0) r2 = UOp.range(4, 1) bufferize_with_ranges = UOp(Ops.STAGE, src=(c, r1, r2), arg=BufferizeOpts(device="CPU")) @@ -586,12 +586,12 @@ class TestConstBufferize(unittest.TestCase): result = graph_rewrite(bufferize_with_ranges, pm_const_buffer_folding, name='test') # BUFFERIZE should be removed self.assertNotEqual(result.op, Ops.STAGE) - const_vals = [u.arg for u in result.toposort() if u.op is Ops.CONST and u.dtype == dtypes.float] + const_vals = [u.arg for u in result.toposort() if u.op is Ops.CONST and u.dtype is dtypes.weakfloat] self.assertIn(3.14, const_vals) class TestUOpTags(unittest.TestCase): def test_inc_by_one(self): - g = UOp.const(1, dtypes.int) + UOp.const(1, dtypes.int) + g = UOp.const(1) + UOp.const(1) assert g.ssimplify() == 2 pm_plus_1 = PatternMatcher([(UPat(Ops.CONST, name="x"), lambda x: x.replace(arg=x.arg+1, tag=1) if x.tag is None else None)]) pm_strip_tags = PatternMatcher([(UPat(GroupOp.All, name="x"), lambda x: x.replace(tag=None) if x.tag is not None else None)]) @@ -682,36 +682,36 @@ class TestUOpGetItem(unittest.TestCase): class TestUOpBroadcast(unittest.TestCase): def test_broadcast_row(self): - a = UOp.const(1, dtypes.float).expand((4, 8)) - b = UOp.const(2, dtypes.float).expand((4, 1)) + a = UOp.const(1.0).expand((4, 8)) + b = UOp.const(2.0).expand((4, 1)) c = a + b self.assertEqual(c.shape, (4, 8)) self.assertEqual(c.op, Ops.ADD) def test_broadcast_col(self): - a = UOp.const(1, dtypes.float).expand((4, 8)) - b = UOp.const(2, dtypes.float).expand((1, 8)) + a = UOp.const(1.0).expand((4, 8)) + b = UOp.const(2.0).expand((1, 8)) c = a + b self.assertEqual(c.shape, (4, 8)) self.assertEqual(c.op, Ops.ADD) def test_broadcast_lower_dim(self): - a = UOp.const(1, dtypes.float).expand((4, 8)) - b = UOp.const(2, dtypes.float).expand((8,)) + a = UOp.const(1.0).expand((4, 8)) + b = UOp.const(2.0).expand((8,)) c = a * b self.assertEqual(c.shape, (4, 8)) self.assertEqual(c.op, Ops.MUL) def test_broadcast_scalar(self): - a = UOp.const(1, dtypes.float).expand((4, 8)) + a = UOp.const(1.0).expand((4, 8)) c = a * 2 self.assertEqual(c.shape, (4, 8)) self.assertEqual(c.op, Ops.MUL) def test_broadcast_symbolic_same_shape(self): t = Variable("t", 1, 10) - a = UOp.const(1, dtypes.float).expand((1, 1, t)) - b = UOp.const(2, dtypes.float).expand((1, 1, t)) + a = UOp.const(1.0).expand((1, 1, t)) + b = UOp.const(2.0).expand((1, 1, t)) c = a + b self.assertEqual(c.op, Ops.ADD) diff --git a/test/null/test_uop_repr.py b/test/null/test_uop_repr.py index ef4f292e59..1c2b1bd252 100644 --- a/test/null/test_uop_repr.py +++ b/test/null/test_uop_repr.py @@ -1,34 +1,34 @@ import unittest -from tinygrad import UOp, dtypes +from tinygrad import UOp class TestUOpRepr(unittest.TestCase): def test_simple_const(self): - a = UOp.const(42, dtypes.int) - self.assertEqual(repr(a), "UOp(Ops.CONST, dtypes.int, arg=42, src=())") + a = UOp.const(42) + self.assertEqual(repr(a), "UOp(Ops.CONST, dtypes.weakint, arg=42, src=())") def test_different_consts(self): - a, b = UOp.const(42, dtypes.int), UOp.const(3, dtypes.int) + a, b = UOp.const(42), UOp.const(3) expected = ( - "UOp(Ops.ADD, dtypes.int, arg=None, src=(\n" + - " UOp(Ops.CONST, dtypes.int, arg=42, src=()),\n" + - " UOp(Ops.CONST, dtypes.int, arg=3, src=()),))" + "UOp(Ops.ADD, dtypes.weakint, arg=None, src=(\n" + + " UOp(Ops.CONST, dtypes.weakint, arg=42, src=()),\n" + + " UOp(Ops.CONST, dtypes.weakint, arg=3, src=()),))" ) self.assertEqual(repr(a+b), expected) def test_walrus_operator_indentation(self): # The reference should have the same indentation as the definition - a = UOp.const(42, dtypes.int) + a = UOp.const(42) expected = ( - "UOp(Ops.ADD, dtypes.int, arg=None, src=(\n" + - " x0:=UOp(Ops.CONST, dtypes.int, arg=42, src=()),\n" + + "UOp(Ops.ADD, dtypes.weakint, arg=None, src=(\n" + + " x0:=UOp(Ops.CONST, dtypes.weakint, arg=42, src=()),\n" + " x0,))" ) self.assertEqual(repr(a+a), expected) def test_nested_walrus_indentation(self): # Ensure indentation is consistent at multiple levels - b = (a:=UOp.const(1, dtypes.int)) + a + b = (a:=UOp.const(1)) + a expected = ( - "UOp(Ops.MUL, dtypes.int, arg=None, src=(\n" + - " x0:=UOp(Ops.ADD, dtypes.int, arg=None, src=(\n" + - " x1:=UOp(Ops.CONST, dtypes.int, arg=1, src=()),\n" + + "UOp(Ops.MUL, dtypes.weakint, arg=None, src=(\n" + + " x0:=UOp(Ops.ADD, dtypes.weakint, arg=None, src=(\n" + + " x1:=UOp(Ops.CONST, dtypes.weakint, arg=1, src=()),\n" + " x1,)),\n" + " x0,))" ) diff --git a/test/null/test_uop_resolve.py b/test/null/test_uop_resolve.py index 490244068b..a0b1d5809d 100644 --- a/test/null/test_uop_resolve.py +++ b/test/null/test_uop_resolve.py @@ -12,7 +12,7 @@ class TestUOpResolve(unittest.TestCase): self.assertEqual(int(u), 11) def test_lt(self): - u = UOp.const(4, dtypes.int) < 7 + u = UOp.const(4) < 7 self.assertTrue(u) def test_rfloordiv(self): @@ -24,24 +24,24 @@ class TestUOpResolve(unittest.TestCase): self.assertEqual(float(u), 2.25) def test_leq(self): - u = UOp.const(4, dtypes.int) <= 4 + u = UOp.const(4) <= 4 self.assertTrue(u) def test_ne(self): - u = UOp.const(4, dtypes.int) != 7 + u = UOp.const(4) != 7 self.assertTrue(u) def test_ne_f(self): - u = UOp.const(4, dtypes.int) != 4 + u = UOp.const(4) != 4 self.assertFalse(u) def test_ngt(self): - u = UOp.const(4, dtypes.int) > 7 + u = UOp.const(4) > 7 self.assertFalse(u) def test_ssimplify(self): - self.assertEqual((8 % UOp.const(4, dtypes.int)).ssimplify(), 0) - self.assertEqual((8 * UOp.const(4, dtypes.int)).ssimplify(), 32) + self.assertEqual((8 % UOp.const(4)).ssimplify(), 0) + self.assertEqual((8 * UOp.const(4)).ssimplify(), 32) def test_ambiguous_less_than(self): u = UOp.variable("i", 1, 10) diff --git a/test/null/test_uop_symbolic.py b/test/null/test_uop_symbolic.py index fa77ae4806..8678d762a0 100644 --- a/test/null/test_uop_symbolic.py +++ b/test/null/test_uop_symbolic.py @@ -12,8 +12,7 @@ from tinygrad.uop.validate import uops_to_z3 def check_uop_against_string(self, v:UOp, s:str): sym_vars = {v.render():v for v in v.toposort() if v.op in (Ops.RANGE, Ops.SPECIAL, Ops.PARAM)} s_eval = eval(s, sym_vars) - if isinstance(s_eval, int) and v.dtype==dtypes.weakint: s_eval = UOp.const(s_eval, dtypes.weakint) - elif isinstance(s_eval, (bool, int, float)): s_eval = UOp.const(s_eval) + if isinstance(s_eval, (bool, int, float)): s_eval = UOp.const(s_eval) s_eval = graph_rewrite(s_eval, commutative, name="cannonicalize eval") self.assertIs(s_eval, v, f"eval did not match simplified: {s_eval} != {v.render()} for {s}") @@ -104,16 +103,16 @@ class TestSymbolic(unittest.TestCase): self.assertEqual(UOp.gcd(a, a*b, a*3).simplify(), a) self.assertEqual(UOp.gcd(a*a*a, a*b*a, a*3*a).simplify(), a*a) self.assertEqual(UOp.gcd(a*a*10, b*a*5, a*a*5).simplify(), a*5) - self.assertEqual(UOp.gcd(a*10, b*5, a*5).simplify(), a.const_like(5)) - self.assertEqual(UOp.gcd(a, b*5, a*5).simplify(), a.const_like(1)) + self.assertEqual(UOp.gcd(a*10, b*5, a*5).simplify(), uconst(5)) + self.assertEqual(UOp.gcd(a, b*5, a*5).simplify(), uconst(1)) def test_divides_exact(self): a = Variable("a", 1, 8) b = Variable("b", 1, 8) self.assertEqual((a*a*3).divide_exact(a).simplify(), a*3) - self.assertEqual((a*a*3).divide_exact(a*a*3).simplify(), a.const_like(1)) - self.assertEqual((a*a*6).divide_exact(a*a*3).simplify(), a.const_like(2)) - self.assertEqual((a*b*3).divide_exact(a.const_like(3)).simplify(), a*b) + self.assertEqual((a*a*3).divide_exact(a*a*3).simplify(), uconst(1)) + self.assertEqual((a*a*6).divide_exact(a*a*3).simplify(), uconst(2)) + self.assertEqual((a*b*3).divide_exact(uconst(3)).simplify(), a*b) self.assertEqual((a*a*3).divide_exact(a*(-3)).simplify(), a*-1) self.assertEqual((a*a*b*3).divide_exact(a*b).simplify(), a*3) self.assertEqual((a*3+a*b).divide_exact(a).simplify(), b+3) @@ -883,7 +882,7 @@ class TestSymbolic(unittest.TestCase): idx = Variable("idx", 0, 24) self.helper_test_variable(idx//4, 0, 6, "(idx//4)") # TODO: simplify the true branch - self.helper_test_variable((idx<4).where(idx//4, idx.const_like(-1)), -1, 6, "(idx<4).where((idx//4), -1)") + self.helper_test_variable((idx<4).where(idx//4, uconst(-1)), -1, 6, "(idx<4).where((idx//4), -1)") def test_floordiv_lt(self): # x//d x0, and <=> c*d NOOP rule. This rule matches patterns that EMERGE during simplification.""" @@ -1375,11 +1374,11 @@ class TestStoreLoadFolding(unittest.TestCase): # Direct: store(idx, load(idx)) -> NOOP self.assertEqual(graph_rewrite(index.store(index.load()), sym).op, Ops.NOOP) # Emergent: store(idx, load(idx) + 0) -> store(idx, load(idx)) -> NOOP - self.assertEqual(graph_rewrite(index.store(index.load() + UOp.const(0, dtypes.int)), sym).op, Ops.NOOP) + self.assertEqual(graph_rewrite(index.store(index.load() + UOp.const(0)), sym).op, Ops.NOOP) # Emergent: store(idx, load(idx) * 1) -> store(idx, load(idx)) -> NOOP - self.assertEqual(graph_rewrite(index.store(index.load() * UOp.const(1, dtypes.int)), sym).op, Ops.NOOP) + self.assertEqual(graph_rewrite(index.store(index.load() * UOp.const(1)), sym).op, Ops.NOOP) # Negative: store(idx, load(idx) + 1) should NOT fold - self.assertEqual(graph_rewrite(index.store(index.load() + UOp.const(1, dtypes.int)), sym).op, Ops.STORE) + self.assertEqual(graph_rewrite(index.store(index.load() + UOp.const(1)), sym).op, Ops.STORE) class TestMoveWhereOnLoad(unittest.TestCase): def test_bool_index_preserves_dtype(self): @@ -1390,7 +1389,7 @@ class TestMoveWhereOnLoad(unittest.TestCase): 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)) - expr = cond.where(idx, idx.const_like(0)) + expr = cond.where(idx, UOp.const(0)) out = graph_rewrite(expr, pm_move_where_on_load) type_verify(out, spec_shared) # Invalid matches any dtype @@ -1475,7 +1474,7 @@ class TestFuzzFailure(unittest.TestCase): v2=Variable('v2', 0, 2) v3=Variable('v3', 0, 1) expr = (((((((((((((((((((((((0//4)%2)//8)+-2)+-4)+-3)+v1)+-4)+v2)+-2)+v3)+v2)//3)%7)*1)//2)+v2)*-1)+2)+1)+0)+-3)+v3) - v1_val, v2_val, v3_val = v1.const_like(8), v2.const_like(0), v3.const_like(0) + v1_val, v2_val, v3_val = UOp.const(8), UOp.const(0), UOp.const(0) num = expr.simplify().substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify() rn = expr.substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify() assert num==rn, f"{num} != {rn}" diff --git a/test/null/test_uop_vmin_vmax.py b/test/null/test_uop_vmin_vmax.py index b426cea615..fb7a6bee4d 100644 --- a/test/null/test_uop_vmin_vmax.py +++ b/test/null/test_uop_vmin_vmax.py @@ -5,12 +5,12 @@ from tinygrad.dtype import dtypes, Invalid class TestVminVmaxProperties(unittest.TestCase): def test_vmin_vmax_constant(self): # vmin and vmax for a constant - uop = UOp.const(42, dtypes.int32) + uop = UOp.const(42) self.assertEqual(uop.vmin, 42) self.assertEqual(uop.vmax, 42) def test_vmin_vmax_cmpne(self): - uop = UOp.const(42, dtypes.int32) + uop = UOp.const(42) def test_bool(u, x): self.assertEqual(u.vmin, x) self.assertEqual(u.vmax, x) @@ -81,8 +81,8 @@ class TestVminVmaxProperties(unittest.TestCase): def test_vmin_vmax_multiplication_0_inf(self): # vmin and vmax for multiplication with a variable - x = UOp.const(0.0, dtypes.float) - y = UOp.load(UOp.param(0, dtypes.float, (1,)), UOp.const(0, dtypes.int), dtype=dtypes.float) + x = UOp.const(0.0) + y = UOp.load(UOp.param(0, dtypes.float, (1,)), UOp.const(0), dtype=dtypes.float) uop = x * y # TODO: these should be 0, but definitely should not be nan self.assertEqual(uop.vmin, -math.inf) @@ -167,7 +167,7 @@ class TestVminVmaxProperties(unittest.TestCase): self.assertNotEqual(i.vmin, i.vmax) def test_vmin_vmax_invalid_vconst(self): - x = UOp.const((0, 4, Invalid, Invalid), dtypes.weakint) + x = UOp.const((0, 4, Invalid, Invalid)) self.assertEqual((x.vmin, x.vmax), (0, 4)) class TestVminVmaxDivMod(unittest.TestCase): @@ -198,14 +198,14 @@ class TestVminVmaxDivMod(unittest.TestCase): def test_vmin_vmax_floordiv_floormod(self): x = UOp.variable('x', -7, 7) - floordiv = x.alu(Ops.FLOORDIV, x.const_like(3)) + floordiv = x.alu(Ops.FLOORDIV, UOp.const(3)) self.assertEqual(floordiv.vmin, -3) self.assertEqual(floordiv.vmax, 2) - floormod = x.alu(Ops.FLOORMOD, x.const_like(3)) + floormod = x.alu(Ops.FLOORMOD, UOp.const(3)) self.assertEqual(floormod.vmin, 0) self.assertEqual(floormod.vmax, 2) # negative const divisor: floormod range is [c+1, 0] - floormod_neg = x.alu(Ops.FLOORMOD, x.const_like(-3)) + floormod_neg = x.alu(Ops.FLOORMOD, UOp.const(-3)) self.assertEqual(floormod_neg.vmin, -2) self.assertEqual(floormod_neg.vmax, 0) @@ -286,31 +286,31 @@ class TestVminVmaxDivMod(unittest.TestCase): class TestVminVmaxVConst(unittest.TestCase): def test_vmin_vmax_vconst_single_element(self): # vmin and vmax for a single-element vector constant - uop = UOp.const((42,), dtypes.int32) + uop = UOp.const((42,)) self.assertEqual(uop.vmin, 42) self.assertEqual(uop.vmax, 42) def test_vmin_vmax_vconst_multiple_elements(self): # vmin and vmax for a multi-element vector constant - uop = UOp.const((10, 20, -5, 7), dtypes.int32) + uop = UOp.const((10, 20, -5, 7)) self.assertEqual(uop.vmin, -5) self.assertEqual(uop.vmax, 20) def test_vmin_vmax_vconst_all_equal(self): # vmin and vmax for a vector where all elements are equal - uop = UOp.const((7, 7, 7), dtypes.int32) + uop = UOp.const((7, 7, 7)) self.assertEqual(uop.vmin, 7) self.assertEqual(uop.vmax, 7) def test_vmin_vmax_vconst_with_negative_values(self): # vmin and vmax for a vector constant containing negative values - uop = UOp.const((-10, -20, -5, -15), dtypes.int32) + uop = UOp.const((-10, -20, -5, -15)) self.assertEqual(uop.vmin, -20) self.assertEqual(uop.vmax, -5) def test_vmin_vmax_vconst_with_floats(self): # vmin and vmax for a vector constant of float values - uop = UOp.const((1.5, -3.2, 0.0), dtypes.float32) + uop = UOp.const((1.5, -3.2, 0.0)) self.assertEqual(uop.vmin, -3.2) self.assertEqual(uop.vmax, 1.5) @@ -323,7 +323,7 @@ class TestVminVmaxVConst(unittest.TestCase): def test_vmin_vmax_vector_with_gep(self): # vmin and vmax for a vector constant of bool values d1 = UOp.param(1, dtypes.int, (1,)) - idx = UOp.const(0, dtypes.int) + idx = UOp.const(0) val = UOp(Ops.LOAD, src=(d1.index(idx),)) uop = (val // 32) self.assertEqual(uop.vmin, -67108864) @@ -332,17 +332,17 @@ class TestVminVmaxVConst(unittest.TestCase): class TestConstFactor(unittest.TestCase): def test_const_factor_constant(self): # const_factor for a constant - uop = UOp.const(42, dtypes.int32) + uop = UOp.const(42) self.assertEqual(uop.const_factor(), 42) def test_const_factor_addition(self): # const_factor for an addition of constants - uop = UOp.const(30, dtypes.int32) + UOp.const(12, dtypes.int32) + uop = UOp.const(30) + UOp.const(12) self.assertEqual(uop.const_factor(), 6) # GCD(30, 12) = 6 def test_const_factor_multiplication(self): # const_factor for a multiplication of constants - uop = UOp.const(5, dtypes.int32) * UOp.const(7, dtypes.int32) + uop = UOp.const(5) * UOp.const(7) self.assertEqual(uop.const_factor(), 5) # For multiplication, it's one of the factors def test_const_factor_with_variable(self): @@ -377,14 +377,14 @@ class TestConstFactor(unittest.TestCase): class TestDivides(unittest.TestCase): def test_divides_constant_exact(self): # Divides a constant by an exact divisor - uop = UOp.const(42, dtypes.int32) + uop = UOp.const(42) result = uop.divides(7) self.assertIsNotNone(result) self.assertEqual(result.const_factor(), 6) # 42 / 7 = 6 def test_divides_constant_inexact(self): # Try to divide a constant by a non-exact divisor - uop = UOp.const(42, dtypes.int32) + uop = UOp.const(42) result = uop.divides(5) self.assertIsNone(result) # 42 is not divisible by 5 diff --git a/test/null/test_uops.py b/test/null/test_uops.py index 1e85d26201..95362b767d 100644 --- a/test/null/test_uops.py +++ b/test/null/test_uops.py @@ -12,8 +12,8 @@ from test.helpers import eval_uop, to_uops_list class TestDTypeFromUOp(unittest.TestCase): def test_broadcastable_promotion(self): - self.assertEqual(dtype_from_uop(Ops.ADD, (UOp.const(1.0, dtypes.float32), UOp.const(1.0, dtypes.float16)), None), dtypes.float32) - self.assertEqual(dtype_from_uop(Ops.MUL, (UOp.const(1, dtypes.int8), UOp.const(1, dtypes.int32)), None), dtypes.int32) + self.assertEqual(dtype_from_uop(Ops.ADD, (UOp.const(1.0).cast(dtypes.float32), UOp.const(1.0).cast(dtypes.float16)), None), dtypes.float32) + self.assertEqual(dtype_from_uop(Ops.MUL, (UOp.const(1).cast(dtypes.int8), UOp.const(1).cast(dtypes.int32)), None), dtypes.int32) def test_same_dtype_fast_path(self): src = (UOp.const(1), UOp.const(2)) @@ -21,7 +21,8 @@ class TestDTypeFromUOp(unittest.TestCase): def test_where_promotion(self): cond = UOp.const(True) - self.assertEqual(dtype_from_uop(Ops.WHERE, (cond, UOp.const(1.0, dtypes.float32), UOp.const(1.0, dtypes.float16)), None), dtypes.float32) + srcs = (cond, UOp.const(1.0).cast(dtypes.float32), UOp.const(1.0).cast(dtypes.float16)) + self.assertEqual(dtype_from_uop(Ops.WHERE, srcs, None), dtypes.float32) idx = UOp.range(4, 0) self.assertEqual(idx.valid(idx < 4).dtype, dtypes.weakint) @@ -200,7 +201,7 @@ class TestGatedStoreRewrite(unittest.TestCase): gidx0 = UOp.special(4, 'gidx0') gate = gidx0EXPAND should be folded into the ALU node, not shown as separate EXPAND nodes - c = UOp.const(1.0, dtypes.float).expand((3,4)) # creates CONST->EXPAND chain + c = UOp.const(1.0).expand((3,4)) # creates CONST->EXPAND chain a = UOp.variable("a", 0.0, 10.0, dtypes.float) alu = a + c with save_viz() as viz: @@ -267,13 +267,13 @@ class TestViz(unittest.TestCase): def test_stack_movement_not_folded_unless_all_const(self): a = UOp.variable("a", 0, 10, dtype=dtypes.int) - c = UOp.const(1, dtypes.int) + c = UOp.const(1) stack = a.stack(c) reshaped = stack.reshape((1, 2)) graph = uop_to_json(VizData(), reshaped) self.assertFalse(graph[id(stack)]["exclude"]) - const_stack = c.stack(UOp.const(2, dtypes.int)) + const_stack = c.stack(UOp.const(2)) const_reshaped = const_stack.reshape((1, 2)) const_graph = uop_to_json(VizData(), const_reshaped) self.assertTrue(const_graph[id(const_stack)]["exclude"]) @@ -401,7 +401,7 @@ class TestVizIntegration(unittest.TestCase): with save_viz() as viz: def test(root): return graph_rewrite(root, sym) - test(c:=UOp.const(1, dtypes.int)) + test(c:=UOp.const(1)) test(c+1) ls = viz.list_items() self.assertEqual(len(ls), 1) @@ -414,7 +414,7 @@ class TestVizIntegration(unittest.TestCase): @track_rewrites() def test(root): return graph_rewrite(root, sym) - test(c:=UOp.const(1, dtypes.int)) + test(c:=UOp.const(1)) test(c+1) ls = viz.list_items() self.assertEqual(len(ls), 2) @@ -425,14 +425,15 @@ class TestVizIntegration(unittest.TestCase): with save_viz() as viz: def default_test(root): return graph_rewrite(root, sym) tracked_test = track_rewrites()(default_test) - c = UOp.const(1, dtypes.int) + c = UOp.const(1) default_test(c+1) # goes to the default group tracked_test(c) # all rewrites after this go inside the second group. default_test(c+2) ls = viz.list_items() self.assertEqual(len(ls), 2) graph = next(viz.get_details(0, 0))["graph"] - self.assertEqual(list(graph), [id(c), id((c+1).src[1]), id(c+1)]) + # both operands of c+1 are the same bare weak CONST, so the graph has two nodes + self.assertEqual(list(graph), [id(c), id(c+1)]) self.assertTrue(graph[id(c)]["exclude"]) self.assertFalse(graph[id(c+1)]["exclude"]) self.assertEqual(list(next(viz.get_details(1, 0))["graph"]), [id(c)]) diff --git a/test/unit/test_assign.py b/test/unit/test_assign.py index f7f3e72bbe..ee93250480 100644 --- a/test/unit/test_assign.py +++ b/test/unit/test_assign.py @@ -661,7 +661,7 @@ class TestAssign(unittest.TestCase): def test_assign_deviceless_const(self): s = Tensor.empty(4, device="CPU:1", dtype=dtypes.float) - s.assign(Tensor(UOp.const(2.0, dtypes.float))) + s.assign(Tensor(UOp.const(2.0).cast(dtypes.float))) np.testing.assert_equal(s.numpy(), [2, 2, 2, 2]) def test_nested_after_contiguous_store(self): diff --git a/test/unit/test_jit.py b/test/unit/test_jit.py index 8aaee9eec2..428bdb05ec 100644 --- a/test/unit/test_jit.py +++ b/test/unit/test_jit.py @@ -351,13 +351,13 @@ class TestJit(unittest.TestCase): @TinyJit def f(x:Tensor) -> Tensor: return (x + 1).realize() with self.assertRaises(JitError): - f(Tensor(UOp.const(2.0, dtypes.float))).item() + f(Tensor(UOp.const(2.0).cast(dtypes.float))).item() def test_jit_deviceless_compute_input(self): @TinyJit def f(x:Tensor) -> Tensor: return (x + 1).realize() with self.assertRaises(JitError): - f(Tensor(UOp.const(2.0, dtypes.float) + UOp.const(1.0, dtypes.float))).item() + f(Tensor(UOp.const(2.0).cast(dtypes.float) + UOp.const(1.0).cast(dtypes.float))).item() def test_jit_init_empty_alt(self): @TinyJit diff --git a/test/unit/test_multitensor.py b/test/unit/test_multitensor.py index 3a2cccac9a..aa1612df6f 100644 --- a/test/unit/test_multitensor.py +++ b/test/unit/test_multitensor.py @@ -60,8 +60,8 @@ class TestMultiTensor(unittest.TestCase): def test_shard_elementwise(self): self._test_shard_op(lambda t:(t+t).reshape(2, 2), [[2.,2.],[2.,2.]]) def test_alu_deviceless_const(self): s = Tensor([1.0, 2, 3, 4]).shard((f"{Device.DEFAULT}:0", f"{Device.DEFAULT}:1"), axis=0) - np.testing.assert_equal((s + Tensor(UOp.const(1.0, dtypes.float))).numpy(), [2, 3, 4, 5]) - np.testing.assert_equal((s + Tensor(UOp.const(1.0, dtypes.float)).reshape((1,)).expand((4,))).numpy(), [2, 3, 4, 5]) + np.testing.assert_equal((s + Tensor(UOp.const(1.0).cast(dtypes.float))).numpy(), [2, 3, 4, 5]) + np.testing.assert_equal((s + Tensor(UOp.const(1.0).cast(dtypes.float)).reshape((1,)).expand((4,))).numpy(), [2, 3, 4, 5]) def test_add_rank_expand_shard(self): # a sharded src keeps its own rank under implicit broadcast, its shard axis right-aligns into the output diff --git a/test/unit/test_tensor_data.py b/test/unit/test_tensor_data.py index 8874e93efe..6b2d1c10cd 100644 --- a/test/unit/test_tensor_data.py +++ b/test/unit/test_tensor_data.py @@ -66,7 +66,7 @@ class TestTensorData(unittest.TestCase): assert dat.shape == () def test_const_dtype_for_uop(self): - self.assertEqual(Tensor.const(UOp.const(1.0, dtypes.float32), dtypes.int8).dtype, dtypes.int8) + self.assertEqual(Tensor.const(UOp.const(1.0).cast(dtypes.float32), dtypes.int8).dtype, dtypes.int8) self.assertEqual(Tensor.const(UOp.variable("x", 1, 10).bind(5), dtypes.int32).item(), 5) def test_data_float32(self): From 8dc225e28e3dbac7c23fda9e1eb02a5e025807a8 Mon Sep 17 00:00:00 2001 From: chenyu Date: Fri, 31 Jul 2026 14:33:15 -0400 Subject: [PATCH 27/44] dtype_from_uop(INS) is None [PR] (#17337) --- tinygrad/uop/ops.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/tinygrad/uop/ops.py b/tinygrad/uop/ops.py index b182e88345..32a0d9ac5b 100644 --- a/tinygrad/uop/ops.py +++ b/tinygrad/uop/ops.py @@ -124,8 +124,10 @@ def dtype_from_uop(op:Ops, src:tuple[UOp,...], arg:Any) -> DType|None: case Ops.CALL: # a CALL of an opaque body is void, a CALL of an address can return a value return dtypes.void if src[0].dtype is dtypes.void else None - case Ops.CUSTOM | Ops.CUSTOMI | Ops.INS | Ops.PYLITERAL: + case Ops.CUSTOM | Ops.CUSTOMI | Ops.PYLITERAL: return dtypes.void + case Ops.INS: + return None case Ops.NOOP: # NOOP can be void or carry any dtype (e.g. x.f(Ops.NOOP) or substitute base with NOOP) return None @@ -1741,7 +1743,7 @@ def graph_rewrite(sink:UOp, pm:PatternMatcher, ctx=None, bottom_up=False, name=N def _rebuild_dtype(n:UOp, new_src:tuple[UOp,...]) -> DType: # TODO: delete this once the dtype field is removed, every rebuild will re-derive # TODO: these ops keep their stored dtype until dtype_from_uop works - if n.op in {Ops.INS, Ops.INDEX, Ops.CUSTOM, Ops.CUSTOMI, Ops.PYLITERAL} or \ + if n.op in {Ops.INDEX, Ops.CUSTOM, Ops.CUSTOMI, Ops.PYLITERAL} or \ all(a.dtype is b.dtype or b.base.arg is Invalid for a,b in zip(n.src, new_src)): return n.dtype return dtype_from_uop(n.op, new_src, n.arg) or n.dtype From 277433259eb71b5fc3d6d5cc33c5a1be1458e9fa Mon Sep 17 00:00:00 2001 From: chenyu Date: Fri, 31 Jul 2026 14:38:38 -0400 Subject: [PATCH 28/44] fix sym_infer for CAST (#17338) --- test/null/test_uop_symbolic.py | 4 ++++ tinygrad/uop/render.py | 2 ++ 2 files changed, 6 insertions(+) diff --git a/test/null/test_uop_symbolic.py b/test/null/test_uop_symbolic.py index 8678d762a0..183175ef09 100644 --- a/test/null/test_uop_symbolic.py +++ b/test/null/test_uop_symbolic.py @@ -1202,6 +1202,10 @@ class TestSymInfer(unittest.TestCase): # floor: 1 % -1000 = -999, 1 // -1000 = -1 assert sym_infer(a%b, var_vals) == -999 assert sym_infer(a//b, var_vals) == -1 + def test_sym_infer_with_cast(self): + a = Variable("a", 0, 100, dtypes.int) + assert sym_infer(a.cast(dtypes.long) + 1, {a.expr: 5}) == 6 + assert sym_infer(a.cast(dtypes.float) * 0.5, {a.expr: 5}) == 2.5 def test_sym_infer_with_bitcast(self): a = Variable("a", 1, 10, dtypes.int) expr = ((a.bitcast(dtypes.uint) << UOp.const(1)).bitcast(dtypes.int) + 2) diff --git a/tinygrad/uop/render.py b/tinygrad/uop/render.py index b88c97f13b..97ebe79db4 100644 --- a/tinygrad/uop/render.py +++ b/tinygrad/uop/render.py @@ -58,6 +58,8 @@ renderer_infer = PatternMatcher([ (UPat(Ops.CDIV, name="x"), lambda ctx,x: f"cdiv({ctx[x.src[0]]}, {ctx[x.src[1]]})"), (UPat(Ops.FLOORMOD, name="x"), lambda ctx,x: f"floormod({ctx[x.src[0]]}, {ctx[x.src[1]]})"), (UPat(Ops.FLOORDIV, name="x"), lambda ctx,x: f"floordiv({ctx[x.src[0]]}, {ctx[x.src[1]]})"), + (UPat(Ops.CAST, name="x"), + lambda ctx,x: f"{'float' if dtypes.is_float(x.dtype) else 'bool' if x.dtype is dtypes.bool else 'int'}({ctx[x.src[0]]})"), (UPat(Ops.BITCAST, name="x"), lambda ctx,x: f"bitcast({ctx[x.src[0]]}, {x.src[0].dtype!r}, {x.dtype!r})"), ]) + renderer From 85ced44db66c7b25a71d727f88333275cb57ac8d Mon Sep 17 00:00:00 2001 From: Christopher Milan Date: Fri, 31 Jul 2026 16:48:51 -0400 Subject: [PATCH 29/44] tc: don't allow reduce over output dims (#17340) --- test/opt/test_tensor_cores.py | 7 +++++++ tinygrad/codegen/opt/postrange.py | 2 ++ 2 files changed, 9 insertions(+) diff --git a/test/opt/test_tensor_cores.py b/test/opt/test_tensor_cores.py index 930177b147..81eb7c6a42 100644 --- a/test/opt/test_tensor_cores.py +++ b/test/opt/test_tensor_cores.py @@ -81,6 +81,13 @@ class TestTensorCores(unittest.TestCase): for tc in Device[Device.DEFAULT].renderer.tensor_cores: helper_tc_allclose(tc.dims[0], tc.dims[1], tc.dims[2], tc.dtype_in, tc.dtype_out, axis=0, tc_opt=0) + @unittest.skipUnless(Device[Device.DEFAULT].renderer.tensor_cores, "test requires tensor cores") + def test_tensor_cores_nested_reduce(self): + tc = Device[Device.DEFAULT].renderer.tensor_cores[0] + a, b = Tensor.empty(tc.dims[1]*2, tc.dims[2], dtype=tc.dtype_in), Tensor.empty(tc.dims[2], tc.dims[0], dtype=tc.dtype_in) + ast = replace_opts(a.matmul(b, dtype=tc.dtype_out).sum(0).schedule_linear().src[-1].src[0], [Opt(OptOps.TC, 0, (-1, 0, 1))]) + with self.assertRaises(KernelOptError): to_program(ast, Device[Device.DEFAULT].renderer) + @Context(ALLOW_TF32=1) @unittest.skipIf(Device.DEFAULT == "PYTHON", "not generated on EMULATED device") @unittest.skipUnless(Device[Device.DEFAULT].renderer.tensor_cores, "test requires tensor cores") diff --git a/tinygrad/codegen/opt/postrange.py b/tinygrad/codegen/opt/postrange.py index 1cf597b0f7..987aea76ba 100644 --- a/tinygrad/codegen/opt/postrange.py +++ b/tinygrad/codegen/opt/postrange.py @@ -245,6 +245,8 @@ class Scheduler: if not (axis < len(axis_choices)): continue axes = list(axis_choices[axis]) + if any(a.arg[-1] is AxisType.REDUCE for a in axes[:2]): raise KernelOptError("tensor core X/Y axes can't be REDUCE") + # tag the reduceop self.ast = self.ast.substitute({reduceop: reduceop.replace(tag="TC")}) From 15d515299ed9b9e1cb1cbe195c0f3d4d63d1f149 Mon Sep 17 00:00:00 2001 From: Christopher Milan Date: Fri, 31 Jul 2026 19:25:18 -0400 Subject: [PATCH 30/44] heuristics: try multiple TC axes (#17341) --- .github/workflows/benchmark.yml | 2 +- tinygrad/codegen/opt/heuristic.py | 28 ++++++++++++---------------- 2 files changed, 13 insertions(+), 17 deletions(-) diff --git a/.github/workflows/benchmark.yml b/.github/workflows/benchmark.yml index b4427d9d05..6027ed90ea 100644 --- a/.github/workflows/benchmark.yml +++ b/.github/workflows/benchmark.yml @@ -516,7 +516,7 @@ jobs: echo "CACHEDB=/tmp/staging.db" >> $GITHUB_ENV rm -f /tmp/staging.db /tmp/staging.db-shm /tmp/staging.db-wal - name: openpilot compile3 big_driving_supercombo - run: BENCHMARK_LOG=usbgpu_openpilot_big_driving_supercombo PICKLE_OOB=1 PYTHONPATH="." GMMU=0 DEV=USB+AMD:LLVM ASSERT_MIN_STEP_TIME=50 python3 examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/10926f2c0911821ca0e72439c1c3bf3ec11f0a08789aa14b7ee8f25379b2afa4 openpilot.pkl + run: BENCHMARK_LOG=usbgpu_openpilot_big_driving_supercombo PICKLE_OOB=1 PYTHONPATH="." TC_OPT=2 GMMU=0 DEV=USB+AMD:LLVM ASSERT_MIN_STEP_TIME=50 python3 examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/10926f2c0911821ca0e72439c1c3bf3ec11f0a08789aa14b7ee8f25379b2afa4 openpilot.pkl - name: openpilot load_pickle big_driving_supercombo run: BENCHMARK_LOG=usbgpu_openpilot_big_driving_supercombo_load_pickle PICKLE_OOB=1 PYTHONPATH="." GMMU=0 DEV=USB+AMD ASSERT_MIN_LOAD_TIME=25 python3 examples/openpilot/load_pickle.py openpilot.pkl - name: openpilot run_pickle big_driving_supercombo diff --git a/tinygrad/codegen/opt/heuristic.py b/tinygrad/codegen/opt/heuristic.py index c9b133fa5b..5b687ae04d 100644 --- a/tinygrad/codegen/opt/heuristic.py +++ b/tinygrad/codegen/opt/heuristic.py @@ -26,22 +26,18 @@ def hand_coded_optimizations(k:Scheduler) -> Scheduler: """ # NOTE: unless TC_OPT is > 0, we only trigger tensor cores if there's only one reduce axis if USE_TC > 0 and (len(k.axes_of(AxisType.GROUP_REDUCE, AxisType.REDUCE)) == 1 or (TC_OPT.value >= 1)): - good_tc_opt = False - tk = k.copy() - try: # check TC first and apply hand-coded opts if successful - rngs = tk.apply_opt(Opt(OptOps.TC, 0, (TC_SELECT.value, TC_OPT.value, USE_TC.value))) - good_tc_opt = True - except KernelOptError: - pass - if good_tc_opt: - if rngs is not None: - for tc_dim in [1,0]: # attempt to upcast M and N - szs = [sz for sz in [5,4,3,2] if rngs[tc_dim].src[0].divides(sz) is not None] - if szs: - # set it to the replaced range - rngs[tc_dim] = tk.apply_opt(Opt(OptOps.UPCAST, tk.rngs.index(rngs[tc_dim]), szs[0]))[0] - if (szs := [sz for sz in [4,2] if rngs[0].src[0].divides(sz) is not None]): # attempt to local N - tk.apply_opt(Opt(OptOps.LOCAL, tk.rngs.index(rngs[0]), szs[0])) + for axis in range(3): + tk = k.copy() + # check TC first and apply hand-coded opts if successful + try: rngs = tk.apply_opt(Opt(OptOps.TC, axis, (TC_SELECT.value, TC_OPT.value, USE_TC.value))) + except KernelOptError: continue + for tc_dim in [1,0]: # attempt to upcast M and N + szs = [sz for sz in [5,4,3,2] if rngs[tc_dim].src[0].divides(sz) is not None] + if szs: + # set it to the replaced range + rngs[tc_dim] = tk.apply_opt(Opt(OptOps.UPCAST, tk.rngs.index(rngs[tc_dim]), szs[0]))[0] + if (szs := [sz for sz in [4,2] if rngs[0].src[0].divides(sz) is not None]): # attempt to local N + tk.apply_opt(Opt(OptOps.LOCAL, tk.rngs.index(rngs[0]), szs[0])) return tk # make a copy so it does not mutate the input From 850989115d9c2622c7acd564bc9fbaf4f8acdef2 Mon Sep 17 00:00:00 2001 From: chenyu Date: Fri, 31 Jul 2026 19:56:37 -0400 Subject: [PATCH 31/44] __int__ and __float__ work for weak (#17342) --- test/null/test_uop_resolve.py | 4 ++++ tinygrad/uop/ops.py | 4 ++-- 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/test/null/test_uop_resolve.py b/test/null/test_uop_resolve.py index a0b1d5809d..5520d534ff 100644 --- a/test/null/test_uop_resolve.py +++ b/test/null/test_uop_resolve.py @@ -7,6 +7,10 @@ class TestUOpResolve(unittest.TestCase): u = UOp.const(4, dtypes.int) self.assertEqual(int(u), 4) + def test_weak_const(self): + self.assertEqual(int(UOp.const(5)), 5) + self.assertEqual(float(UOp.const(1.5)), 1.5) + def test_int_add(self): u = UOp.const(4, dtypes.int) + 7 self.assertEqual(int(u), 11) diff --git a/tinygrad/uop/ops.py b/tinygrad/uop/ops.py index 32a0d9ac5b..17809ff3c1 100644 --- a/tinygrad/uop/ops.py +++ b/tinygrad/uop/ops.py @@ -529,8 +529,8 @@ class UOp(RandMixin, metaclass=UOpMetaClass): assert isinstance(vmin, expected_type), f"vmin is wrong dtype {type(vmin)} != {expected_type}" return vmin def __bool__(self): return self._eval((dtypes.bool,), bool) - def __int__(self): return self._eval(dtypes.ints, int) - def __float__(self): return float(self._eval(dtypes.floats, float)) + def __int__(self): return self._eval(dtypes.ints+(dtypes.weakint,), int) + def __float__(self): return float(self._eval(dtypes.floats+(dtypes.weakfloat,), float)) def substitute(self, dvars:dict[UOp, UOp], name:str|None=None, extra_pm:PatternMatcher|None=None, walk:bool=False, enter_calls:bool=False): dvars = {k:v for k,v in dvars.items() if k is not v} if len(dvars) == 0: return self From a88f832f0ce040c00f2d7e703da584f95f3b6d6d Mon Sep 17 00:00:00 2001 From: George Hotz <72895+geohot@users.noreply.github.com> Date: Fri, 31 Jul 2026 18:58:38 -0700 Subject: [PATCH 32/44] remove UOp.val (#17345) --- examples/gpt2.py | 19 +++++-------------- test/unit/test_llm_server.py | 14 ++++++-------- tinygrad/llm/model.py | 5 +++-- tinygrad/uop/ops.py | 2 -- 4 files changed, 14 insertions(+), 26 deletions(-) diff --git a/examples/gpt2.py b/examples/gpt2.py index 604840fc55..7e12258f54 100644 --- a/examples/gpt2.py +++ b/examples/gpt2.py @@ -22,10 +22,6 @@ class Attention: self.head_dim = dim // n_heads def __call__(self, x:Tensor, start_pos:Variable, mask:Optional[Tensor]) -> Tensor: - if mask is not None or start_pos.val == 0: - # no symbolic shape qkv when consuming prompts - start_pos = start_pos.val - if HALF: x = x.half() xqkv = self.c_attn(x).reshape(None, None, 3, self.n_heads, self.head_dim) xq, xk, xv = [xqkv[:, :, i, :, :] for i in range(3)] @@ -38,12 +34,8 @@ class Attention: # update the cache self.cache_kv[:, :, start_pos:start_pos+seqlen, :, :].assign(Tensor.stack(xk, xv)).realize() - if start_pos > 0: - keys = self.cache_kv[0][:, :start_pos+seqlen, :, :] - values = self.cache_kv[1][:, :start_pos+seqlen, :, :] - else: - keys = xk - values = xv + keys = self.cache_kv[0][:, :start_pos+seqlen, :, :] + values = self.cache_kv[1][:, :start_pos+seqlen, :, :] xq, keys, values = xq.transpose(1, 2), keys.transpose(1, 2), values.transpose(1, 2) return self.c_proj(xq.scaled_dot_product_attention(keys, values, mask).transpose(1, 2).reshape(bsz, seqlen, self.dim)) @@ -86,15 +78,14 @@ class Transformer: seqlen = tokens.shape[1] tok_emb = self.wte(tokens) - # not symbolic when consuming the prompt - selected_pos = (0, seqlen) if start_pos.val == 0 else (start_pos, start_pos+1) - pos_emb = self.wpe(self.allpos.shrink((None, selected_pos))) + # start_pos is a bound Variable, so everything below it stays symbolic + pos_emb = self.wpe(self.allpos.shrink((None, (start_pos, start_pos+seqlen)))) h = tok_emb + pos_emb if HALF: h = h.half() - mask = Tensor.full((1, 1, seqlen, start_pos.val+seqlen), float("-inf"), dtype=h.dtype).triu(start_pos.val+1) if seqlen > 1 else None + mask = Tensor.full((1, 1, seqlen, start_pos+seqlen), float("-inf"), dtype=h.dtype).triu(start_pos+1) if seqlen > 1 else None for hi in self.h: h = hi(h, start_pos, mask) diff --git a/test/unit/test_llm_server.py b/test/unit/test_llm_server.py index 9a7d0e9908..bbc8c90907 100644 --- a/test/unit/test_llm_server.py +++ b/test/unit/test_llm_server.py @@ -6,6 +6,8 @@ from tinygrad.llm.model import Transformer, TransformerConfig TEST_CONFIG = TransformerConfig(num_blocks=1, dim=64, hidden_dim=128, n_heads=2, n_kv_heads=2, norm_eps=1e-5, vocab_size=100, head_dim=32, rope_theta=10000.0, rope_dim=32, v_head_dim=32, max_context=32) +V_START_POS = UOp.variable("start_pos", 0, TEST_CONFIG.max_context-1) +V_TOKS = UOp.variable("toks", 1, 32) # 32 is the default chunk_size in generate class TestTransformerGenerate(unittest.TestCase): def test_kv_cache_reuse(self): @@ -14,7 +16,7 @@ class TestTransformerGenerate(unittest.TestCase): captured_inputs = [] def mock_call(self, tokens, start_pos, temperature): - captured_inputs.append((tokens.shape, start_pos if isinstance(start_pos, int) else start_pos.val)) + captured_inputs.append((tokens.shape, start_pos)) return Tensor([[42]]) with patch.object(Transformer, '__call__', mock_call): @@ -31,9 +33,7 @@ class TestTransformerGenerate(unittest.TestCase): next(gen) # should process tokens[6:] = [42, 10, 11, 12] since first 6 have cached k/v - toks_shape = captured_inputs[0][0][-1] - self.assertEqual(toks_shape.val if isinstance(toks_shape, UOp) else toks_shape, 4) - self.assertEqual(captured_inputs[0][1], 6) + self.assertEqual(captured_inputs, [((1, V_TOKS.bind(4)), V_START_POS.bind(6))]) def test_kv_cache_invalidation(self): """Test that generate invalidates the KV cache when tokens diverge from the cached prefix.""" @@ -41,7 +41,7 @@ class TestTransformerGenerate(unittest.TestCase): captured_inputs = [] def mock_call(self, tokens, start_pos, temperature): - captured_inputs.append((tokens.shape, start_pos if isinstance(start_pos, int) else start_pos.val)) + captured_inputs.append((tokens.shape, start_pos)) return Tensor([[42]]) with patch.object(Transformer, '__call__', mock_call): @@ -55,9 +55,7 @@ class TestTransformerGenerate(unittest.TestCase): next(gen) # should process all 3 tokens from start - toks_shape = captured_inputs[0][0][-1] - self.assertEqual(toks_shape.val if isinstance(toks_shape, UOp) else toks_shape, 3) - self.assertEqual(captured_inputs[0][1], 0) + self.assertEqual(captured_inputs, [((1, V_TOKS.bind(3)), V_START_POS.bind(0))]) def test_two_prompts_schedule_cache(self): """Third prompt should hit the schedule cache, not miss (first two warm up both jits: prefill + decode).""" diff --git a/tinygrad/llm/model.py b/tinygrad/llm/model.py index 421005342a..e802c7b394 100644 --- a/tinygrad/llm/model.py +++ b/tinygrad/llm/model.py @@ -432,9 +432,10 @@ class Transformer: if start_pos < len(self._cached_tokens) and (resets := [r for b in self.blk for r in b._state_reset_ops()]): Tensor.realize(*resets) out, prompt_len = None, len(tokens) while len(tokens) < self.max_context: - sp, nt = v_start_pos.bind(start_pos), v_toks.bind(min(chunk_size, len(tokens) - start_pos)) + n_toks = min(chunk_size, len(tokens) - start_pos) + sp, nt = v_start_pos.bind(start_pos), v_toks.bind(n_toks) out = self(t[:, sp:sp+nt] if start_pos < prompt_len or out is None else out, sp, temp).realize() - start_pos += nt.val + start_pos += n_toks # chunked prefill: keep processing until all prompt tokens are consumed if start_pos < len(tokens): continue tokens.append(int(out.item())) diff --git a/tinygrad/uop/ops.py b/tinygrad/uop/ops.py index 17809ff3c1..1ad814049a 100644 --- a/tinygrad/uop/ops.py +++ b/tinygrad/uop/ops.py @@ -988,8 +988,6 @@ class UOp(RandMixin, metaclass=UOpMetaClass): def unbind_all(self) -> tuple[UOp, dict[Variable, int]]: ret:dict[Variable, int] = {} return graph_rewrite(self, pm_unbind, ctx=ret), ret - @property - def val(self) -> int: return self.unbind()[1] def variables(self) -> list[Variable]: return sorted({x for x in self.backward_slice_with_self if x.op is Ops.PARAM and x.arg.addrspace is AddrSpace.ALU}, key=lambda v: v.expr) From 099d69ff7d5773cd2324c36f014fa6e65285458d Mon Sep 17 00:00:00 2001 From: George Hotz <72895+geohot@users.noreply.github.com> Date: Fri, 31 Jul 2026 19:41:18 -0700 Subject: [PATCH 33/44] ci: split macos unit test into metal and mock runners (#17346) --- .github/workflows/test.yml | 23 +++++++++++++++++++---- 1 file changed, 19 insertions(+), 4 deletions(-) diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index e4b54e6c21..cb1cddf171 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -643,12 +643,8 @@ jobs: with: key: unittest-macos deps: testing_unit - amd: 'true' - ocelot: 'true' - name: Run unit tests run: DEV=METAL python -m pytest -n=auto test/unit/ --durations=20 - - name: Run NULL backend tests - run: SPEC=2 DEV=NULL python -m pytest -n=auto test/null/ --durations=20 - name: Test tensor core ops (fake) run: DEV=METAL DEBUG=3 TC=2 python test/backend/test_ops.py TestOps.test_gemm - name: Test tensor core ops (real) @@ -659,6 +655,25 @@ jobs: run: DEV=METAL python3 -m pytest test/device/test_metal.py #- name: Fuzz Test linearizer # run: DEV=METAL DEPTH=4 FUZZ_N=50 FUZZ_MAX_SIZE=1000000 python test/external/fuzz_linearizer.py + - name: Run process replay tests + uses: ./.github/actions/process-replay + + unittestmacosmock: + name: MacOS (unit, mock) + runs-on: macos-26 + timeout-minutes: 20 + steps: + - name: Checkout Code + uses: actions/checkout@v6 + - name: Setup Environment + uses: ./.github/actions/setup-tinygrad + with: + key: unittest-macos-mock + deps: testing_unit + amd: 'true' + ocelot: 'true' + - name: Run NULL backend tests + run: SPEC=2 DEV=NULL python -m pytest -n=auto test/null/ --durations=20 - name: Run pytest (amd) env: DEV: MOCKKFD+AMD From 9082ecef5dbdbaaf6a8b5e50bc305f078996367b Mon Sep 17 00:00:00 2001 From: George Hotz <72895+geohot@users.noreply.github.com> Date: Fri, 31 Jul 2026 19:51:42 -0700 Subject: [PATCH 34/44] use .val to access the value of Ops.CONST (#17347) --- tinygrad/codegen/decomp/dtype.py | 2 +- tinygrad/codegen/decomp/op.py | 14 +++++------ tinygrad/codegen/decomp/transcendental.py | 4 +-- tinygrad/codegen/late/coalesce.py | 10 ++++---- tinygrad/codegen/opt/heuristic.py | 4 +-- tinygrad/codegen/simplify.py | 2 +- tinygrad/engine/realize.py | 2 +- tinygrad/mixin/elementwise.py | 2 +- tinygrad/renderer/cstyle.py | 30 +++++++++++------------ tinygrad/renderer/isa/x86.py | 26 ++++++++++---------- tinygrad/renderer/llvmir.py | 4 +-- tinygrad/renderer/nir.py | 4 +-- tinygrad/renderer/ptx.py | 6 ++--- tinygrad/renderer/wgsl.py | 6 ++--- tinygrad/runtime/ops_dsp.py | 2 +- tinygrad/runtime/ops_python.py | 2 +- tinygrad/runtime/support/hcq2.py | 2 +- tinygrad/schedule/multi.py | 2 +- tinygrad/schedule/rangeify.py | 4 +-- tinygrad/uop/divandmod.py | 10 ++++---- tinygrad/uop/ops.py | 26 +++++++++++--------- tinygrad/uop/symbolic.py | 24 +++++++++--------- tinygrad/uop/upat.py | 2 +- 23 files changed, 97 insertions(+), 93 deletions(-) diff --git a/tinygrad/codegen/decomp/dtype.py b/tinygrad/codegen/decomp/dtype.py index df9a24e47b..d682aad8f5 100644 --- a/tinygrad/codegen/decomp/dtype.py +++ b/tinygrad/codegen/decomp/dtype.py @@ -158,7 +158,7 @@ pm_long_decomp = PatternMatcher([ (UPat(Ops.LOAD, tuple(l2i_dt.keys()), src=(UPat.var('idx'),), name='x'), lambda x,idx: x.replace(dtype=l2i_dt[x.dtype], src=(reindex(idx, x.tag[0]).replace(dtype=l2i_dt[x.dtype], tag=None),), tag=None) if x.tag is not None else None), (UPat(Ops.CONST, tag={(w, dt) for w in (0, 1) for dt in l2i_dt.values()}, name='x'), lambda x: - UOp.const(truncate[x.tag[1]]((x.arg >> 32) if x.tag[0] == 1 else (x.arg & 0xFFFFFFFF)), x.tag[1])) + UOp.const(truncate[x.tag[1]]((x.val >> 32) if x.tag[0] == 1 else (x.val & 0xFFFFFFFF)), x.tag[1])) ]) # float decomposition patterns - ctx is (fr, to) tuple diff --git a/tinygrad/codegen/decomp/op.py b/tinygrad/codegen/decomp/op.py index 6efe4147d1..6b0526b498 100644 --- a/tinygrad/codegen/decomp/op.py +++ b/tinygrad/codegen/decomp/op.py @@ -77,7 +77,7 @@ def get_simplifying_rewrite_patterns(ops:tuple[Ops, ...]) -> PatternMatcher: # these are rewrites that make things simpler pat: list[tuple[UPat, Callable]] = [(UPat.var("a")//UPat.var("b"), floordiv_to_idiv)] # FLOORMOD by 2**y -> x & (2**y-1) (correct floor mod for any sign in two's complement); fires before floormod_to_mod - if Ops.AND in ops: pat.append((UPat.var("x", dtypes.ints)%UPat.cvar("c"), lambda x,c: x & (c.arg-1) if c.arg in powers_of_two else None)) + if Ops.AND in ops: pat.append((UPat.var("x", dtypes.ints)%UPat.cvar("c"), lambda x,c: x & (c.val-1) if c.val in powers_of_two else None)) pat.append((UPat.var("a")%UPat.var("b"), floormod_to_mod)) # no real hardware supports THREEFRY, but NullRenderer does if Ops.THREEFRY not in ops: pat.append((UPat(Ops.THREEFRY, dtype=dtypes.uint64, src=(UPat.var("x"), UPat.var("key"))), threefry2x32)) @@ -91,19 +91,19 @@ def get_late_rewrite_patterns(ops:tuple[Ops, ...], disable_fast_idiv:bool) -> Pa if Ops.OR in ops: pat += [(UPat.var("x", dtypes.bool).logical_not()&UPat.var("y", dtypes.bool).logical_not(), lambda x,y: (x | y).logical_not())] # rewrite MUL/CDIV to SHL+SHR: x*(2**y) -> shl(x,y) and x//(2**y) -> shr(x,y) - if Ops.SHL in ops: pat += [(UPat.var("x", dtypes.ints)*UPat.cvar("c"), lambda c,x: x << v if (v:=powers_of_two.get(c.arg, 0)) else None)] + if Ops.SHL in ops: pat += [(UPat.var("x", dtypes.ints)*UPat.cvar("c"), lambda c,x: x << v if (v:=powers_of_two.get(c.val, 0)) else None)] if Ops.SHR in ops: # uint CDIV by 2**v -> x >> v (FLOORDIV is lowered to CDIV by the rule above before reaching here) pat += [(UPat(Ops.CDIV, src=(UPat.var("x", dtypes.uints), UPat.cvar("c"))), - lambda x,c: x >> v if (v:=powers_of_two.get(c.arg, 0)) else None)] + lambda x,c: x >> v if (v:=powers_of_two.get(c.val, 0)) else None)] # signed CDIV (trunc) by 2**v -> (x + (x<0 ? c-1 : 0)) >> v pat += [(UPat(Ops.CDIV, src=(UPat.var("x", dtypes.ints), UPat.cvar("c"))), lambda x,c: (x+(l.const_like(l.vmin) if (l:=(x<0)).vmin==l.vmax else l).where(c-1, 0)) >> v - if (v:=powers_of_two.get(c.arg, 0)) else None)] + if (v:=powers_of_two.get(c.val, 0)) else None)] if not disable_fast_idiv: # fast_idiv handles non-pow2: only fire on non-negative inputs (signed magic-mul is unreliable for x<0) pat += [(UPat(Ops.CDIV, src=(UPat.var("x", dtypes.ints), UPat.cvar("d"))), - lambda ctx, x, d: fast_idiv(ctx, x, d.arg) if x.vmin >= 0 or x.dtype in dtypes.uints else None)] + lambda ctx, x, d: fast_idiv(ctx, x, d.val) if x.vmin >= 0 or x.dtype in dtypes.uints else None)] # rewrite raw CMOD -> x - d*CDIV(x,d) so fast_idiv can pick up the CDIV. only on non-negative inputs; # avoids disturbing floormod_to_mod's general-path output (which uses a trunc Ops.CMOD as an implementation detail) pat += [(UPat(Ops.CMOD, src=(UPat.var("x", dtypes.ints), UPat.var("d"))), @@ -119,13 +119,13 @@ def get_late_rewrite_patterns(ops:tuple[Ops, ...], disable_fast_idiv:bool) -> Pa (UPat.var("x", dtypes.sints)*-1 < UPat.var("y", dtypes.sints)*UPat.cvar("c"), lambda x,y,c: y*(-c) x==c + lambda x,c1,c2: x.eq(c1+1) if c1.val+1==c2.val-1 else None), # (c-1) x==c ] if Ops.CMPEQ in ops: pat += [(UPat.var('x').ne(UPat.var('y')).logical_not(), lambda x,y: x.alu(Ops.CMPEQ, y))] if Ops.MULACC in ops: pat += [(UPat.var('a')*UPat.var('b')+UPat.var('c'), lambda a,b,c: a.alu(Ops.MULACC, b, c))] # also fuse (x << n) + c → MULACC(x, 2^n, c) since MUL→SHL may run first - if Ops.SHL in ops: pat += [(UPat.var('x').alu(Ops.SHL, UPat.cvar('n'))+UPat.var('c'), lambda x,n,c: x.alu(Ops.MULACC, x.const_like(1< a/b if Ops.FDIV in ops: pat += [(UPat.var("x").reciprocal(), lambda x: x.const_like(1).alu(Ops.FDIV, x))] diff --git a/tinygrad/codegen/decomp/transcendental.py b/tinygrad/codegen/decomp/transcendental.py index 23b1723503..b2771acb37 100644 --- a/tinygrad/codegen/decomp/transcendental.py +++ b/tinygrad/codegen/decomp/transcendental.py @@ -16,8 +16,8 @@ def exponent_bias(d:DType) -> int: return (1 << (dtypes.finfo(d)[0] - 1)) - (0 i def exponent_mask(d:DType) -> int: return (1 << dtypes.finfo(d)[0]) - 1 # **** utils **** -def shr(x:UOp|int, y:UOp|int) -> UOp: return x // (2**(y.simplify().arg) if isinstance(y, UOp) else 2**y) -def shl(x:UOp|int, y:UOp|int) -> UOp: return x * (2**(y.simplify().arg) if isinstance(y, UOp) else 2**y) +def shr(x:UOp|int, y:UOp|int) -> UOp: return x // (2**(y.simplify().val) if isinstance(y, UOp) else 2**y) +def shl(x:UOp|int, y:UOp|int) -> UOp: return x * (2**(y.simplify().val) if isinstance(y, UOp) else 2**y) def rintk(d:UOp) -> UOp: """round d:float to int away from 0""" diff --git a/tinygrad/codegen/late/coalesce.py b/tinygrad/codegen/late/coalesce.py index 436032e35f..47ec328017 100644 --- a/tinygrad/codegen/late/coalesce.py +++ b/tinygrad/codegen/late/coalesce.py @@ -84,7 +84,7 @@ def transform_to_image(ctx, buf:UOp, x:UOp) -> UOp|None: h, w, cidx = cands[0] if len(cands) == 1 else min(cands, key=lambda cand: len(cand[2].index(1).simplify().backward_slice)) buf = buf.replace(src=(shape_to_shape_arg((h, w, 4)),)) shapes[buf.arg.slot] = (h, w) - if valid.op is not Ops.CONST or valid.arg is not True: + if valid.op is not Ops.CONST or valid.val is not True: return buf.index(cidx.src[1].valid(valid), cidx.src[0].valid(valid), dtype=dtypes.float) else: return buf.index(cidx.src[1], cidx.src[0], dtype=dtypes.float) @@ -111,10 +111,10 @@ def memory_coalescing(sink:UOp, ctx:Renderer) -> UOp: if buf.addrspace == AddrSpace.REG: continue idx, valid = idx_u.get_idx(), idx_u.get_valid() root_src: UOp|str - if idx.op is Ops.ADD and idx.src[1].op is Ops.CONST: root_src, arg = idx.src[0], idx.src[1].arg - elif idx.op is Ops.ADD and idx.src[0].op is Ops.CONST: root_src, arg = idx.src[1], idx.src[0].arg - elif idx.op is Ops.CONST and idx.arg is Invalid: root_src, arg = "INVALID", 0 - elif idx.op is Ops.CONST: root_src, arg = "CONST", idx.arg + if idx.op is Ops.ADD and idx.src[1].op is Ops.CONST: root_src, arg = idx.src[0], idx.src[1].val + elif idx.op is Ops.ADD and idx.src[0].op is Ops.CONST: root_src, arg = idx.src[1], idx.src[0].val + elif idx.op is Ops.CONST and idx.val is Invalid: root_src, arg = "INVALID", 0 + elif idx.op is Ops.CONST: root_src, arg = "CONST", idx.val else: root_src, arg = idx, 0 memory[(u.op, buf, root_src, valid)].setdefault(arg, []).append(u) diff --git a/tinygrad/codegen/opt/heuristic.py b/tinygrad/codegen/opt/heuristic.py index 5b687ae04d..f78dcb2b78 100644 --- a/tinygrad/codegen/opt/heuristic.py +++ b/tinygrad/codegen/opt/heuristic.py @@ -126,8 +126,8 @@ def hand_coded_optimizations(k:Scheduler) -> Scheduler: if rng in idx.backward_slice: num_strides += 1 for c in idx.split_uop(Ops.ADD): if c is rng: sum_strides += 1 - if c.op is Ops.MUL and c.src[0] is rng and c.src[1].op is Ops.CONST: sum_strides += c.src[1].arg - if c.op is Ops.MUL and c.src[1] is rng and c.src[0].op is Ops.CONST: sum_strides += c.src[0].arg + if c.op is Ops.MUL and c.src[0] is rng and c.src[1].op is Ops.CONST: sum_strides += c.src[1].val + if c.op is Ops.MUL and c.src[1] is rng and c.src[0].op is Ops.CONST: sum_strides += c.src[0].val xb_choices.append((num_strides, sum_strides, axis, upcast_amount)) if xb_choices: xb_choices = sorted(xb_choices) diff --git a/tinygrad/codegen/simplify.py b/tinygrad/codegen/simplify.py index 647a6c1f11..608a656d25 100644 --- a/tinygrad/codegen/simplify.py +++ b/tinygrad/codegen/simplify.py @@ -62,7 +62,7 @@ pm_simplify_ranges = PatternMatcher([ def mark_range_mod(ctx:dict[UOp, UOp|None], r:UOp, c:UOp) -> None: # ranges that aren't looped over can't be split if r not in ctx and r.arg[-1] not in {AxisType.WARP, AxisType.DEVICE} \ - and r.src[0].op is Ops.CONST and r.src[0].divides(c.arg) is not None: ctx[r] = c + and r.src[0].op is Ops.CONST and r.src[0].divides(c.val) is not None: ctx[r] = c def do_substitute(ctx:dict, x: UOp, sub_fxn:Callable[[UOp, UOp], UOp]) -> UOp|None: ret = x.substitute({k:sub_fxn(k,v) for k,v in ctx.items() if v is not None}) diff --git a/tinygrad/engine/realize.py b/tinygrad/engine/realize.py index f11c9688d4..a5f8155162 100644 --- a/tinygrad/engine/realize.py +++ b/tinygrad/engine/realize.py @@ -198,7 +198,7 @@ def exec_validate(ctx:ExecContext, call:UOp, ast:UOp) -> float|None: def exec_encdec(ctx:ExecContext, call:UOp, ast:UOp) -> float|None: bufs = [cast(Buffer, b.buffer).ensure_allocated() for b in resolve_params(call, ctx.input_uops)] - shape, pos_var = tuple(s.arg for s in ast.src if s.op is Ops.CONST), ast.variables()[0].expr + shape, pos_var = tuple(s.val for s in ast.src if s.op is Ops.CONST), ast.variables()[0].expr with track_stats(ctx, call, bufs[0].device, bufs, ctx.var_vals): bufs[0].allocator._encode_decode(bufs[0]._buf, bufs[1]._buf, bufs[2]._buf, [x._buf for x in bufs[3:]], shape, ctx.var_vals[pos_var]) return None diff --git a/tinygrad/mixin/elementwise.py b/tinygrad/mixin/elementwise.py index f73e0a0bea..c74baa8b39 100644 --- a/tinygrad/mixin/elementwise.py +++ b/tinygrad/mixin/elementwise.py @@ -24,7 +24,7 @@ class ElementwiseMixin(CreationMixin): out_dtype = least_upper_dtype(x.dtype, y.dtype) # keep weak CONST weak, might lift weakint -> weakfloat def promote(t): - if t.dtype in dtypes.weaks and t._uop.base.op is Ops.CONST: return t._wrap_uop(t._uop.const_like(t._uop.base.arg, weak_dtype(out_dtype))) + if t.dtype in dtypes.weaks and t._uop.base.op is Ops.CONST: return t._wrap_uop(t._uop.const_like(t._uop.base.val, weak_dtype(out_dtype))) return t.cast(out_dtype) return promote(x), promote(y) diff --git a/tinygrad/renderer/cstyle.py b/tinygrad/renderer/cstyle.py index dfb18ff382..6dbf090bb4 100644 --- a/tinygrad/renderer/cstyle.py +++ b/tinygrad/renderer/cstyle.py @@ -34,18 +34,18 @@ base_rewrite = PatternMatcher([ # const (UPat(Ops.CONST, arg=math.inf, name="x"), lambda ctx, x: f"({ctx.render_cast(x, ctx.infinity)})"), (UPat(Ops.CONST, arg=-math.inf, name="x"), lambda ctx, x: f"({ctx.render_cast(x, f'-{ctx.infinity}')})"), - (UPat(Ops.CONST, dtype=dtypes.floats, name="x"), lambda ctx,x: f"({ctx.render_cast(x, ctx.nan)})" if math.isnan(x.arg) else None), - (UPat(Ops.CONST, dtype=dtypes.float, name="x"), lambda ctx,x: f"{x.arg}f"), - (UPat(Ops.CONST, dtype=dtypes.int64, name="x"), lambda ctx,x: f"{x.arg}l"), - (UPat(Ops.CONST, dtype=dtypes.uint64, name="x"), lambda ctx,x: f"{truncate[x.dtype](x.arg)}ul"), - (UPat(Ops.CONST, dtype=dtypes.uint32, name="x"), lambda ctx,x: f"{truncate[x.dtype](x.arg)}u"), - (UPat(Ops.CONST, dtype=dtypes.bool, name="x"), lambda ctx,x: "1" if x.arg else "0"), + (UPat(Ops.CONST, dtype=dtypes.floats, name="x"), lambda ctx,x: f"({ctx.render_cast(x, ctx.nan)})" if math.isnan(x.val) else None), + (UPat(Ops.CONST, dtype=dtypes.float, name="x"), lambda ctx,x: f"{x.val}f"), + (UPat(Ops.CONST, dtype=dtypes.int64, name="x"), lambda ctx,x: f"{x.val}l"), + (UPat(Ops.CONST, dtype=dtypes.uint64, name="x"), lambda ctx,x: f"{truncate[x.dtype](x.val)}ul"), + (UPat(Ops.CONST, dtype=dtypes.uint32, name="x"), lambda ctx,x: f"{truncate[x.dtype](x.val)}u"), + (UPat(Ops.CONST, dtype=dtypes.bool, name="x"), lambda ctx,x: "1" if x.val else "0"), # consts are rendered to larger type and casted - (UPat(Ops.CONST, (*dtypes.fp8s, dtypes.bfloat16, dtypes.half), name="x"), lambda ctx,x: f"({ctx.render_cast(x, f'{x.arg}f')})"), - (UPat(Ops.CONST, (dtypes.uint8, dtypes.uint16), name="x"), lambda ctx,x: f"({ctx.render_cast(x, f'{x.arg}u')})"), - (UPat(Ops.CONST, (dtypes.int8, dtypes.int16), name="x"), lambda ctx,x: f"({ctx.render_cast(x, str(x.arg))})"), + (UPat(Ops.CONST, (*dtypes.fp8s, dtypes.bfloat16, dtypes.half), name="x"), lambda ctx,x: f"({ctx.render_cast(x, f'{x.val}f')})"), + (UPat(Ops.CONST, (dtypes.uint8, dtypes.uint16), name="x"), lambda ctx,x: f"({ctx.render_cast(x, f'{x.val}u')})"), + (UPat(Ops.CONST, (dtypes.int8, dtypes.int16), name="x"), lambda ctx,x: f"({ctx.render_cast(x, str(x.val))})"), # default const render - (UPat(Ops.CONST, name="x"), lambda ctx,x: str(x.arg)), + (UPat(Ops.CONST, name="x"), lambda ctx,x: str(x.val)), # SHRINK/INDEX (UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var('idx')), name="x"), lambda ctx,**kwargs: ctx.render_index(**kwargs)), @@ -164,7 +164,7 @@ class CStyleLanguage(Renderer): if buf.addrspace == AddrSpace.ALU: # this is lane access in C if idx.op is not Ops.CONST: return f"({self[buf]})[{self[idx]}]" - return self[buf]+(f"[{idx.arg}]" if buf.max_numel() > self.gep_arr_threshold else f".{'xyzwabcd'[idx.arg]}") + return self[buf]+(f"[{idx.val}]" if buf.max_numel() > self.gep_arr_threshold else f".{'xyzwabcd'[idx.val]}") return f"({self[buf]}+{strip_parens(self[idx]) if idx.arg == Ops.ADD else self[idx]})" def render_buffer(self, x:UOp): @@ -494,10 +494,10 @@ class HIPRenderer(CStyleLanguage): (UPat(Ops.WMMA, name="x"), lambda ctx,x: f"__{_wmma_name(x)}({ctx[x.src[0]]}, {ctx[x.src[1]]}, {ctx[x.src[2]]}," f" {fp8_index(x.src[0].dtype)}, {fp8_index(x.src[0].dtype)}, 0, 0, 0, 0)" if x.arg[0][2] == 128 else None), (UPat(Ops.WMMA, name="x"), lambda ctx,x: f"__{_wmma_name(x)}({ctx[x.src[0]]}, {ctx[x.src[1]]}, {ctx[x.src[2]]}, 0, 0, 0)"), - (UPat(Ops.CONST, dtypes.fp8s, name="x"), lambda ctx,x: f"f32_to_fp8({ctx.nan}, {fp8_index(x.dtype)})" if math.isnan(x.arg) else None), + (UPat(Ops.CONST, dtypes.fp8s, name="x"), lambda ctx,x: f"f32_to_fp8({ctx.nan}, {fp8_index(x.dtype)})" if math.isnan(x.val) else None), (UPat(Ops.CONST, dtypes.fp8s, arg=math.inf, name="x"), lambda ctx,x: f"f32_to_fp8({ctx.infinity}, {fp8_index(x.dtype)})"), (UPat(Ops.CONST, dtypes.fp8s, arg=-math.inf, name="x"), lambda ctx,x: f"f32_to_fp8(-{ctx.infinity}, {fp8_index(x.dtype)})"), - (UPat(Ops.CONST, dtypes.fp8s, name="x"), lambda ctx,x: f"f32_to_fp8({x.arg}f, {fp8_index(x.dtype)})"), + (UPat(Ops.CONST, dtypes.fp8s, name="x"), lambda ctx,x: f"f32_to_fp8({x.val}f, {fp8_index(x.dtype)})"), (UPat(Ops.CAST, dtypes.fp8s, (UPat(dtype=dtypes.float),), name="x",), lambda ctx,x: f"f32_to_fp8({ctx[x.src[0]]}, {fp8_index(x.dtype)})"), (UPat(Ops.CAST, dtypes.float, (UPat.var("y", dtypes.fp8s),), name="x",), @@ -522,7 +522,7 @@ class HIPRenderer(CStyleLanguage): lambda x: x.replace(src=(x.src[0].bitcast(dtypes.uint64), x.src[1].bitcast(dtypes.uint64), x.src[2])) if x.src[0].max_numel() == 8 and x.src[0].dtype in dtypes.fp8_ocp else None), # bfloat16 constant casting - (UPat.cvar('x', dtypes.bfloat16), lambda x: cast_float_to_bf16(UOp.const(x.arg, dtypes.float))), + (UPat.cvar('x', dtypes.bfloat16), lambda x: cast_float_to_bf16(UOp.const(x.val, dtypes.float))), ]) def asm(self, prg:UOp, lin:UOp) -> bytes: @@ -538,7 +538,7 @@ class HIPRenderer(CStyleLanguage): prefix, ockl = [], [] type_map = { dtypes.bfloat16: "bf16", dtypes.float: "f32", dtypes.half: "f16", dtypes.fp8e4m3: "_fp8_fp8", dtypes.fp8e5m2: "_bf8_bf8" } used_dtypes = uops_to_dtypes(uops) - if any(u.op is Ops.CONST and not math.isfinite(u.arg) for u in uops): + if any(u.op is Ops.CONST and not math.isfinite(u.val) for u in uops): prefix += ["#define INFINITY (__builtin_inff())", "#define NAN (__builtin_nanf(\"\"))"] if any(u.op is Ops.SPECIAL for u in uops): prefix.append("typedef long unsigned int size_t;") diff --git a/tinygrad/renderer/isa/x86.py b/tinygrad/renderer/isa/x86.py index 2436d8bcf4..4ea41f643c 100644 --- a/tinygrad/renderer/isa/x86.py +++ b/tinygrad/renderer/isa/x86.py @@ -289,8 +289,8 @@ def fold_address(x:UOp) -> tuple[UOp, UOp, UOp, UOp]: # buffers are indexed by element, everything else (the stack pointer) by byte scale = base.dtype.itemsize if base.op in {Ops.PARAM, Ops.BUFFER, Ops.AFTER} else 1 sz = imm(dtypes.uint8, base.dtype.itemsize) - if idx.op is Ops.ADD and idx.src[1].op is Ops.CONST: return (base, _cast(idx.src[0]), _disp(idx.src[1].arg * scale), sz) - if idx.op is Ops.CONST: return (base, UOp(Ops.NOOP), _disp(idx.arg * scale), sz) + if idx.op is Ops.ADD and idx.src[1].op is Ops.CONST: return (base, _cast(idx.src[0]), _disp(idx.src[1].val * scale), sz) + if idx.op is Ops.CONST: return (base, UOp(Ops.NOOP), _disp(idx.val * scale), sz) return (base, _cast(idx), _disp(0), sz) def abi(ctx:IselContext, x:UOp) -> UOp|None: @@ -353,7 +353,7 @@ isel_matcher = PatternMatcher([ # cast of void is a noop (UPat.var("y").cast(name="x"), lambda y,x: y if y.dtype == dtypes.void else None), # range is lowered to acc, cmp, jmp after regalloc - (UPat(Ops.RANGE, src=(UPat.cvar("c"),), allow_any_len=True, name="x"), lambda c,x: x.replace(src=(imm(c.dtype, c.arg),) + x.src[1:])), + (UPat(Ops.RANGE, src=(UPat.cvar("c"),), allow_any_len=True, name="x"), lambda c,x: x.replace(src=(imm(c.dtype, c.val),) + x.src[1:])), (UPat(Ops.RANGE, name="x"), lambda ctx,x: x.replace(tag=(ctx.vreg(WGPR),)) if not isinstance(x.tag, tuple) else None), # really all a backedge END is is an IF with a tag referencing the RANGE start label (UPat(Ops.END, src=(UPat(), UPat(), UPat(GroupOp.Comparison, name="cond")), name="x"), @@ -367,10 +367,10 @@ isel_matcher = PatternMatcher([ # function abi constraints (UPat((Ops.PARAM, Ops.SPECIAL), name="x"), abi), # constants that can't be immediates, move them to registers - (UPat.cvar("x", dtypes.int64s), lambda x: x.ins(X86Ops.MOVABS, src=(imm(x.dtype, x.arg),)) if not x.tag else None), - (UPat.cvar("x", dtypes.ints+(dtypes.bool,)), lambda x: x.ins(X86Ops.MOVi, src=(imm(x.dtype, x.arg),)) if not x.tag else None), + (UPat.cvar("x", dtypes.int64s), lambda x: x.ins(X86Ops.MOVABS, src=(imm(x.dtype, x.val),)) if not x.tag else None), + (UPat.cvar("x", dtypes.ints+(dtypes.bool,)), lambda x: x.ins(X86Ops.MOVi, src=(imm(x.dtype, x.val),)) if not x.tag else None), (UPat.cvar("x", dtypes.floats), lambda x: - UOp.const(struct.unpack((dt:=to_int(x.dtype)).fmt, struct.pack(x.dtype.fmt, x.arg))[0], dt).bitcast(x.dtype) if not x.tag else None), + UOp.const(struct.unpack((dt:=to_int(x.dtype)).fmt, struct.pack(x.dtype.fmt, x.val))[0], dt).bitcast(x.dtype) if not x.tag else None), # conditional moves that use masks NOTE: these currently assume a mask producing cmp exists (UPat.var("m").where(UPat.var("a", dtypes.int8s+dtypes.int16s+dtypes.int32s+(dtypes.int64,)), UPat.var("b")), lambda m,a,b: a.ins(X86Ops.VPBLENDVB, src=(b, a, m.replace(dtype=m.src[0].dtype))) if a.max_numel() > 1 else None), @@ -453,9 +453,9 @@ isel_matcher = PatternMatcher([ # scalar int binary ((UPat(dtype=dtypes.ints).alu(Ops.CDIV, UPat())).named("x"), idiv), # scalar int binary with immediate - (UPat.var("a", dtypes.ints) << UPat.cvar("c"), lambda a,c: a.ins(X86Ops.SHLi, src=(a, imm(dtypes.uint8, c.arg)))), - (UPat.var("a", dtypes.uints) >> UPat.cvar("c"), lambda a,c: a.ins(X86Ops.SHRi, src=(a, imm(dtypes.uint8, c.arg)))), - (UPat.var("a", dtypes.sints) >> UPat.cvar("c"), lambda a,c: a.ins(X86Ops.SARi, src=(a, imm(dtypes.uint8, c.arg)))), + (UPat.var("a", dtypes.ints) << UPat.cvar("c"), lambda a,c: a.ins(X86Ops.SHLi, src=(a, imm(dtypes.uint8, c.val)))), + (UPat.var("a", dtypes.uints) >> UPat.cvar("c"), lambda a,c: a.ins(X86Ops.SHRi, src=(a, imm(dtypes.uint8, c.val)))), + (UPat.var("a", dtypes.sints) >> UPat.cvar("c"), lambda a,c: a.ins(X86Ops.SARi, src=(a, imm(dtypes.uint8, c.val)))), (UPat.var("a", dtypes.ints) + UPat.cvar("c"), lambda a,c: a.ins(X86Ops.ADDi, src=(a, i)) if (i:=to_imm(c)) is not None else None), (UPat.var("a", dtypes.ints) * UPat.cvar("c"), lambda a,c: a.ins(X86Ops.IMULi, src=(a, i)) if (i:=to_imm(c)) is not None else None), (UPat.var("a", dtypes.ints+(dtypes.bool,)) & UPat.cvar("c"), lambda a,c: a.ins(X86Ops.ANDi, src=(a, i)) if (i:=to_imm(c)) is not None else None), @@ -664,10 +664,10 @@ def encode(x:UOp, opc:int, reg:int|None=None, pp:int=0, sel:int=0, we:int=0) -> # DISP byte if mod == 0b01 or mod == 0b10: assert disp_uop is not None - inst += struct.pack(unwrap(disp_uop.dtype.fmt), disp_uop.arg) + inst += struct.pack(unwrap(disp_uop.dtype.fmt), disp_uop.val) # IMM byte if imm_uop is not None: - if imm_uop.op is Ops.CONST: inst += struct.pack(unwrap(imm_uop.dtype.fmt), imm_uop.arg) + if imm_uop.op is Ops.CONST: inst += struct.pack(unwrap(imm_uop.dtype.fmt), imm_uop.val) elif isinstance(greg(imm_uop), Register): inst += bytes([(greg(imm_uop).index & 0b1111) << 4 | 0b0000]) return inst @@ -840,10 +840,10 @@ class X86Renderer(ISARenderer): def _format_op(x:UOp) -> str: return f" {(o[7:-1] if (o:=str(x.arg))[-1] in ('i', 'm') else o[7:]).lower():7s}" def _format_operands(x:UOp) -> str: def _format(src:tuple[UOp, ...]) -> list[str]: - return [str(s.arg) if s.op is Ops.CONST else reg_strs[o].get(s.dtype.itemsize, o) if \ + return [str(s.val) if s.op is Ops.CONST else reg_strs[o].get(s.dtype.itemsize, o) if \ (o:=str(greg(s))) in reg_strs else o for s in src if greg(s) is not None] def _mem_adress(base:UOp, idx:UOp, disp:UOp, sz:UOp) -> list[str]: - return [f"[{greg(base)}" + (f" + {greg(idx)}*{sz.arg}" if greg(idx) else "") + (f" + {disp.arg}" if disp.arg else "") + "]"] + return [f"[{greg(base)}" + (f" + {greg(idx)}*{sz.val}" if greg(idx) else "") + (f" + {disp.val}" if disp.val else "") + "]"] if len(x.src) > 4 and x.arg in X86GroupOp.WriteMem: ret = _mem_adress(*x.src[:4]) + _format(x.src[4:]) elif len(x.src) > 3 and x.arg in X86GroupOp.Rm1st: ret = _format((x,)) + _mem_adress(*x.src[:4]) + _format(x.src[4:]) diff --git a/tinygrad/renderer/llvmir.py b/tinygrad/renderer/llvmir.py index 6ae37820d8..56b3b38027 100644 --- a/tinygrad/renderer/llvmir.py +++ b/tinygrad/renderer/llvmir.py @@ -67,7 +67,7 @@ base_rewrite = PatternMatcher([ f" {ctx[x]} = getelementptr inbounds {ldt(x.dtype)}, {ldt(x.dtype, ptr=True)} {ctx[x.src[0]]}, {ldt(x.src[1].dtype)} {ctx[x.src[1]]}"), # register index (UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.cvar("idx")), name="x"), lambda ctx,buf,idx,x: - f" {ctx[x]} = extractelement {ldt(buf.dtype, buf.max_numel())} {ctx[buf]}, i32 {idx.arg}" if buf.addrspace == AddrSpace.ALU else None), + f" {ctx[x]} = extractelement {ldt(buf.dtype, buf.max_numel())} {ctx[buf]}, i32 {idx.val}" if buf.addrspace == AddrSpace.ALU else None), # load/store (UPat(Ops.LOAD, src=(UPat.var("idx"), UPat.var("alt"), UPat.var("mask")), name="x"), @@ -170,7 +170,7 @@ class LLVMRenderer(Renderer): kernel.append(f" {r[u]} = addrspacecast [{size} x {ldt(u.dtype)}] addrspace(3)* @{r[u][1:]} to [{size} x {ldt(u.dtype)}]*") else: kernel.append(f" {r[u]} = alloca [{size} x {ldt(u.dtype)}], align 16") - elif u.op is Ops.CONST: r[u] = lconst(u.arg, u.dtype) + elif u.op is Ops.CONST: r[u] = lconst(u.val, u.dtype) elif u.op is Ops.CAST and ldt(u.dtype) == ldt(u.src[0].dtype): r[u] = r[u.src[0]] # cast from signed to unsigned of the same size is a noop, or pointer cast else: diff --git a/tinygrad/renderer/nir.py b/tinygrad/renderer/nir.py index fbac922a10..a531007ff7 100644 --- a/tinygrad/renderer/nir.py +++ b/tinygrad/renderer/nir.py @@ -122,7 +122,7 @@ class NIRRenderer(Renderer): extra_matcher = PatternMatcher([ # handle negative unsigned CONST - (UPat.cvar("x", dtypes.uints), lambda x: UOp.const(x.dtype.max+x.arg+1, x.dtype) if x.arg < 0 else None), + (UPat.cvar("x", dtypes.uints), lambda x: UOp.const(x.dtype.max+x.val+1, x.dtype) if x.val < 0 else None), # from ptx (UPat.var('x', dtype=dtypes.bool) uint8 @@ -144,7 +144,7 @@ class NIRRenderer(Renderer): ]) def_rewrite = PatternMatcher([ - (UPat(Ops.CONST, name="x"), lambda ctx,x: nimm(ctx.b, x.arg, x.dtype)), + (UPat(Ops.CONST, name="x"), lambda ctx,x: nimm(ctx.b, x.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"))), diff --git a/tinygrad/renderer/ptx.py b/tinygrad/renderer/ptx.py index 36a46a8d1d..1b6b77859a 100644 --- a/tinygrad/renderer/ptx.py +++ b/tinygrad/renderer/ptx.py @@ -79,8 +79,8 @@ def modifier(a: DType, b: DType): return '.rzi' if dtypes.is_int(a) and dtypes.i (a.itemsize < b.itemsize or dtypes.is_int(b) or b == dtypes.bool) else '' string_rewrite = PatternMatcher([ - (UPat.cvar("x", dtypes.bool), lambda ctx, x: f"setp.ne.s16 {ctx.r[x]}, {render_val(x.arg, x.dtype)}, 0;"), - (UPat.cvar("x"), lambda ctx, x: f"mov.b{ctx.types[x.dtype][1:]} {ctx.r[x]}, {render_val(x.arg, x.dtype)};"), + (UPat.cvar("x", dtypes.bool), lambda ctx, x: f"setp.ne.s16 {ctx.r[x]}, {render_val(x.val, x.dtype)}, 0;"), + (UPat.cvar("x"), lambda ctx, x: f"mov.b{ctx.types[x.dtype][1:]} {ctx.r[x]}, {render_val(x.val, x.dtype)};"), (UPat(Ops.SPECIAL, name="x"), lambda ctx,x: f"mov.u32 %{x.arg}, %{'ctaid' if x.arg[0] == 'g' else 'tid'}.{chr(120+int(x.arg[-1]))};"), (UPat(Ops.PARAM, name="x"), lambda ctx, x: f"ld.param.{ctx.types[dtypes.ulong] if x.addrspace is AddrSpace.GLOBAL else ctx.mem_types[x.dtype]} {ctx.r[x]}, [data{x.arg.slot}+0];"), @@ -203,7 +203,7 @@ class PTXRenderer(Renderer): # on REG, INDEX/SHRINK pick the register (must be CONST) and LOAD is a noop if u.op is not Ops.LOAD and u.src[1].op is not Ops.CONST: raise RuntimeError(f"PTX does not support dynamic register indexing: {u}") - r[u] = r[u.src[0]] if u.op is Ops.LOAD else r[u.src[0]][u.src[1].arg] + r[u] = r[u.src[0]] if u.op is Ops.LOAD else r[u.src[0]][u.src[1].val] continue if u.op is Ops.SPECIAL: r[u] = "%" + u.arg elif u.op is Ops.LOAD: diff --git a/tinygrad/renderer/wgsl.py b/tinygrad/renderer/wgsl.py index 71c969629a..e9e2bfdcc3 100644 --- a/tinygrad/renderer/wgsl.py +++ b/tinygrad/renderer/wgsl.py @@ -68,10 +68,10 @@ class WGSLRenderer(CStyleLanguage): string_rewrite = PatternMatcher([ (UPat(Ops.NEG, dtypes.uints, src=(UPat.var('x'))), lambda ctx,x: f"(0-{ctx[x]})"), - (UPat.cvar("x", dtype=dtypes.bool), lambda x: "true" if x.arg else "false"), + (UPat.cvar("x", dtype=dtypes.bool), lambda x: "true" if x.val else "false"), (UPat(Ops.CONST, dtype=(dtypes.uchar, dtypes.ushort, dtypes.uint32), name="x"), - lambda x: f"bitcast({x.arg})" if x.arg < 0 else f"{x.arg&0xFFFFFFFF}u"), - (UPat(Ops.CONST, dtype=dtypes.int32, name="x"), lambda ctx,x: f"{truncate[x.dtype](x.arg)}"), + lambda x: f"bitcast({x.val})" if x.val < 0 else f"{x.val&0xFFFFFFFF}u"), + (UPat(Ops.CONST, dtype=dtypes.int32, name="x"), lambda ctx,x: f"{truncate[x.dtype](x.val)}"), (UPat(Ops.BUFFER, name="x"), lambda ctx,x: f"var{'' if x.addrspace == AddrSpace.LOCAL else ''} {ctx[x]}: array<{ctx.buf_map(x)},{_packed_size(x)}>;"), (UPat(Ops.BITCAST, dtype=dtypes.half, name="x", src=(UPat(dtype=(dtypes.short, dtypes.ushort, dtypes.uint32),),)), diff --git a/tinygrad/runtime/ops_dsp.py b/tinygrad/runtime/ops_dsp.py index 9419ea962e..e8ffa2ef83 100644 --- a/tinygrad/runtime/ops_dsp.py +++ b/tinygrad/runtime/ops_dsp.py @@ -13,7 +13,7 @@ from tinygrad.uop.ops import PatternMatcher, UPat # NOTE: this just increases readability of the generated code dsp_string = PatternMatcher([ - (UPat(Ops.CONST, (dtypes.int8, dtypes.uint8), name="x"), lambda ctx,x: str(x.arg)), + (UPat(Ops.CONST, (dtypes.int8, dtypes.uint8), name="x"), lambda ctx,x: str(x.val)), ]) class DSPRenderer(ClangRenderer): diff --git a/tinygrad/runtime/ops_python.py b/tinygrad/runtime/ops_python.py index 5b2f14483b..5112df0040 100644 --- a/tinygrad/runtime/ops_python.py +++ b/tinygrad/runtime/ops_python.py @@ -102,7 +102,7 @@ class PythonProgram(Program['PythonDevice']): elif u.op is Ops.SPECIAL: if u.arg[0] == 'g': values[u] = [idxs[2-int(u.arg[-1])]] * warp_size elif u.arg[0] == 'l': values[u] = [x[2-int(u.arg[-1])] for x in warp] - elif u.op is Ops.CONST: values[u] = [u.arg] * warp_size + elif u.op is Ops.CONST: values[u] = [u.val] * warp_size elif u.op in {Ops.INDEX, Ops.SHRINK}: ret:list = [] if u.src[0].addrspace == AddrSpace.ALU: diff --git a/tinygrad/runtime/support/hcq2.py b/tinygrad/runtime/support/hcq2.py index 3952f30e9e..498703bc3a 100644 --- a/tinygrad/runtime/support/hcq2.py +++ b/tinygrad/runtime/support/hcq2.py @@ -55,7 +55,7 @@ def make_cmdbuf(lin, devs, buf:UOp|None=None, dep:UOp|None=None): blob, patches = bytearray(), [] for s in (s for ins in lin.src for s in ins.src): if s.op is not Ops.CONST: patches.append((len(blob), s)) - blob.extend(struct.pack(f'<{s.dtype.fmt}', s.arg if s.op is Ops.CONST else 0x0)) + blob.extend(struct.pack(f'<{s.dtype.fmt}', s.val if s.op is Ops.CONST else 0x0)) cmdbuf = buf if buf is not None else UOp.placeholder((len(blob) // 4,), dtypes.uint32, next(UOp.unique_num), device=devs).rtag("cmdbuf") writable = cmdbuf.after(dep) if dep is not None else cmdbuf return cmdbuf.after(make_binary_patch(writable, bytes(blob)), *((make_patches(writable, patches),) if patches else ())) diff --git a/tinygrad/schedule/multi.py b/tinygrad/schedule/multi.py index 440b1565f9..17110647f4 100644 --- a/tinygrad/schedule/multi.py +++ b/tinygrad/schedule/multi.py @@ -209,7 +209,7 @@ def index_multi(root:UOp, multi:UOp): continue # strided ownership: idx ≡ rng (mod shard_sz), intra-shard position is (idx - rng) // shard_sz diff = (idxs[ax] - rng).simplify() - if (mod:=(diff % shard_sz).simplify()).op is Ops.CONST and mod.arg == 0: + if (mod:=(diff % shard_sz).simplify()).op is Ops.CONST and mod.val == 0: local = (diff // shard_sz).simplify() if local.vmin >= 0 and local.vmax < shard_sz: idxs[ax] = local diff --git a/tinygrad/schedule/rangeify.py b/tinygrad/schedule/rangeify.py index ff15b77b24..0e6b44b376 100644 --- a/tinygrad/schedule/rangeify.py +++ b/tinygrad/schedule/rangeify.py @@ -283,7 +283,7 @@ def remove_bufferize(src:UOp, buf:UOp, idx:UOp): # if it makes it here, the bufferize is removed # this is the ranges replaced # NOTE: if buf src is a const, we don't replace it. if idx is Invalid (dead load), don't replace it either - replaced = {k:v for k,v in zip(buf.src[1:], idx.src[1:]) if k.op is not Ops.CONST and not (v.op is Ops.CONST and v.arg is Invalid)} + replaced = {k:v for k,v in zip(buf.src[1:], idx.src[1:]) if k.op is not Ops.CONST and not (v.op is Ops.CONST and v.val is Invalid)} return src.substitute(replaced, extra_pm=pm_gate_substitute) def remove_noop_bufferize(idx,b2): @@ -304,7 +304,7 @@ pm_const_buffer_folding = pm_mops+PatternMatcher([ (UPat(Ops.INDEX, src=(UPat(Ops.STAGE),), allow_any_len=True, name="idx").f(Ops.NOOP).f(Ops.STAGE, allow_any_len=True, name="b2"), remove_noop_bufferize), # no buffers for const (ranges don't matter for const - it's the same value everywhere) - (UPat(Ops.CONST, name='c').f(Ops.STAGE, allow_any_len=True, name="b"), lambda c,b: b.const_like(c.arg)), + (UPat(Ops.CONST, name='c').f(Ops.STAGE, allow_any_len=True, name="b"), lambda c,b: b.const_like(c.val)), # indexing a const is a const (UPat(Ops.INDEX, src=(UPat(Ops.CONST, name="c"),),), lambda c: c), # indexing an after with all fully invalid stores is invalid diff --git a/tinygrad/uop/divandmod.py b/tinygrad/uop/divandmod.py index 8ad865baca..bb082bc3d3 100644 --- a/tinygrad/uop/divandmod.py +++ b/tinygrad/uop/divandmod.py @@ -12,7 +12,7 @@ def fold_divmod_general(d: UOp) -> UOp|None: # x//y is constant if (xdiv:=x//y).vmin == xdiv.vmax: return x - xdiv.vmin*y if d.op is Ops.FLOORMOD else xdiv.const_like(xdiv.vmin) # PARAM // c is irreducible - if x.op is Ops.PARAM and y.op is Ops.CONST and x.arg.multiple_of % y.arg == 0: return d.const_like(0) if d.op is Ops.FLOORMOD else None + if x.op is Ops.PARAM and y.op is Ops.CONST and x.arg.multiple_of % y.val == 0: return d.const_like(0) if d.op is Ops.FLOORMOD else None # split uops for the rest of the processing x_peeled, const = x.pop_const() @@ -20,7 +20,7 @@ def fold_divmod_general(d: UOp) -> UOp|None: # ** Constant Denominator Rules ** # these rules strictly require y to be a scalar constant > 0 - if y.op is Ops.CONST and (c := y.arg) > 0: + if y.op is Ops.CONST and (c := y.val) > 0: # nested_div: (x%(k*c))//c -> (x//c)%k (requires k>0); the mod case is handled by remove_nested_mod below if d.op is Ops.FLOORDIV and x.op is Ops.FLOORMOD and (k := x.src[1].divides(c)) is not None and k > 0: return x.src[0] // y % k @@ -76,7 +76,7 @@ def fold_divmod_general(d: UOp) -> UOp|None: # divide_by_gcd: x//y -> (x//gcd)//(y//gcd) gcd = UOp.gcd(*all_uops, y).simplify() - if not (gcd.op is Ops.CONST and gcd.arg==1): + if not (gcd.op is Ops.CONST and gcd.val==1): ret = unwrap(x.divide_exact(gcd)).alu(d.op, unwrap(y.divide_exact(gcd))) return ret*gcd if d.op is Ops.FLOORMOD else ret @@ -85,8 +85,8 @@ def fold_divmod_general(d: UOp) -> UOp|None: quo, rem = [], [] for u in all_uops: if (q:=u.divide_exact(y)) is not None: quo.append(q) - elif y.op is Ops.CONST and (c:=u.const_factor())%y.arg!=c: - rem.append(u.divides(c)*(c%y.arg)) + elif y.op is Ops.CONST and (c:=u.const_factor())%y.val!=c: + rem.append(u.divides(c)*(c%y.val)) quo.append(u.divides(c)*(c//y.arg) if d.op is Ops.FLOORDIV else u.const_like(0)) else: rem.append(u) diff --git a/tinygrad/uop/ops.py b/tinygrad/uop/ops.py index 1ad814049a..17899ce79f 100644 --- a/tinygrad/uop/ops.py +++ b/tinygrad/uop/ops.py @@ -259,6 +259,10 @@ class UOp(RandMixin, metaclass=UOpMetaClass): if (self.op, self.dtype, self.src, self.arg, self.tag) == new_args: return self return UOp(*new_args) def rtag(self, tag=True): return self.replace(tag=tag) + @property + def val(self): + assert self.op is Ops.CONST, f"val is only for CONST, got {self.op}" + return self.arg @recursive_property def key(self) -> bytes: return hashlib.sha256(str((self.op, self.dtype, self.arg)).encode() + b"".join([s.key for s in self.src])).digest() @@ -521,7 +525,7 @@ class UOp(RandMixin, metaclass=UOpMetaClass): from tinygrad.uop.symbolic import symbolic with Context(TRACK_MATCH_STATS=0 if not tracked else TRACK_MATCH_STATS.value): return graph_rewrite(self, symbolic, name="simplify") - def ssimplify(self) -> UOp|ConstType: return ret.arg if (ret:=self.simplify()).op is Ops.CONST else ret + def ssimplify(self) -> UOp|ConstType: return ret.val if (ret:=self.simplify()).op is Ops.CONST else ret def _eval(self, dtype, expected_type:Type[T]) -> T: assert self.dtype in dtype, f"eval with wrong dtype {self}" vmin, vmax = (simple_self:=self.simplify())._min_max @@ -765,9 +769,9 @@ class UOp(RandMixin, metaclass=UOpMetaClass): # cached property here makes external_uop_gc fail, why? @property def as_shape(self) -> tuple[sint, ...]: - if self.op is Ops.CONST: return (self.arg,) + if self.op is Ops.CONST: return (self.val,) if self.op is not Ops.STACK: return (ssimplify(self),) - return tuple(s.arg if s.op is Ops.CONST else ssimplify(s) for s in self.src) + return tuple(s.val if s.op is Ops.CONST else ssimplify(s) for s in self.src) @functools.cached_property def marg(self): @@ -895,7 +899,7 @@ class UOp(RandMixin, metaclass=UOpMetaClass): idx = self.flatten().index(UOp.range(self.numel(), 0)) out = graph_rewrite(idx, pm_mops+symbolic+pm_contiguous_view_offset, ctx=self, name="contiguous_view_offset") - return out.arg if out.op is Ops.CONST and isinstance(out.arg, int) else None + return out.val if out.op is Ops.CONST and isinstance(out.val, int) else None def has_buffer_identity(self, after_ok=False): """Check if this UOp has a concrete buffer identity in the graph (RESHAPE/UNSHARD -> BUFFER chain).""" @@ -984,7 +988,7 @@ class UOp(RandMixin, metaclass=UOpMetaClass): return UOp(Ops.BIND, src=(self, uval)) def unbind(self) -> tuple[Variable, int]: assert self.op is Ops.BIND and self.src[0].op is Ops.PARAM and self.src[1].op is Ops.CONST, f"can't unbind {self}" - return self.src[0], self.src[1].arg + return self.src[0], self.src[1].val def unbind_all(self) -> tuple[UOp, dict[Variable, int]]: ret:dict[Variable, int] = {} return graph_rewrite(self, pm_unbind, ctx=ret), ret @@ -997,15 +1001,15 @@ class UOp(RandMixin, metaclass=UOpMetaClass): def const_factor(self) -> int: """largest known int that divides self""" # TODO: for negatives it's not the largest - if self.op is Ops.CONST: return self.arg + if self.op is Ops.CONST: return self.val if self.op is Ops.STACK: return math.gcd(*[x.const_factor() for x in self.src]) if self.op is Ops.ADD: return math.gcd(self.src[0].const_factor(), self.src[1].const_factor()) - if self.op is Ops.MUL: return self.src[0].arg if self.src[0].op is Ops.CONST else self.src[1].arg if self.src[1].op is Ops.CONST else 1 + if self.op is Ops.MUL: return self.src[0].val if self.src[0].op is Ops.CONST else self.src[1].val if self.src[1].op is Ops.CONST else 1 if self.op is Ops.PARAM and self.arg.multiple_of is not None: return self.arg.multiple_of return 1 def divides(self, v:int) -> UOp|None: if v==1: return self - if self.op is Ops.CONST: return self.const_like(self.arg//v) if self.arg%v == 0 else None + if self.op is Ops.CONST: return self.const_like(self.val//v) if self.val%v == 0 else None if self.op is Ops.STACK: srcs = tuple(s.divides(v) for s in self.src) return None if any(s is None for s in srcs) else UOp(Ops.STACK, src=cast(tuple[UOp, ...], srcs)) @@ -1016,7 +1020,7 @@ class UOp(RandMixin, metaclass=UOpMetaClass): if self.op is Ops.PARAM and self.arg.multiple_of is not None: return self // v if self.arg.multiple_of%v == 0 else None return None # generic None if we aren't sure def pop_const(self, op=Ops.ADD) -> tuple[UOp, PyConst]: # NOTE: assume Invalid ALU is resolved - return (self.src[0], self.src[1].arg) if self.op is op and self.src[1].op is Ops.CONST else (self, identity_element(op, self.dtype)) + return (self.src[0], self.src[1].val) if self.op is op and self.src[1].op is Ops.CONST else (self, identity_element(op, self.dtype)) @staticmethod def gcd(*uops: UOp) -> UOp: terms, factors = zip(*[(u.divides(f:=u.const_factor()),f) for u in uops]) @@ -1024,7 +1028,7 @@ class UOp(RandMixin, metaclass=UOpMetaClass): return math.prod([*count.elements(), terms[0].const_like(math.gcd(*factors))]) # put the const at the top def divide_exact(self, v:UOp) -> UOp|None: if self is v: return self.const_like(1) - if v.op is Ops.CONST: return self.divides(v.arg) + if v.op is Ops.CONST: return self.divides(v.val) if self.op is Ops.ADD: return None if (s0:=self.src[0].divide_exact(v)) is None or (s1:=self.src[1].divide_exact(v)) is None else s0+s1 if self.op is Ops.MUL: (fac, const), (div_fac, div_const) = self.pop_const(Ops.MUL), v.pop_const(Ops.MUL) @@ -1082,7 +1086,7 @@ class UOp(RandMixin, metaclass=UOpMetaClass): if self.op in (Ops.RANGE, Ops.SPECIAL) and self.dtype is not dtypes.void: return 0, (self.src[0]-1).vmax if self.op is Ops.BIND: return self.src[0]._min_max # ignore the bound value if self.op is Ops.STACK: return min(x.vmin for x in self.src), max(x.vmax for x in self.src) - if self.op is Ops.CONST and self.arg is not Invalid: return self.arg, self.arg + if self.op is Ops.CONST and self.val is not Invalid: return self.val, self.val if self.op is Ops.INDEX: return self.src[0]._min_max if self.op is Ops.CAST: # a cast to unsigned keeps exact bounds when the source fits diff --git a/tinygrad/uop/symbolic.py b/tinygrad/uop/symbolic.py index 0c39015d09..3e21f6eb0b 100644 --- a/tinygrad/uop/symbolic.py +++ b/tinygrad/uop/symbolic.py @@ -13,10 +13,10 @@ from tinygrad.codegen.decomp.transcendental import xpow # ******** phase 1 of symbolic used to live in ops, it's the most generic folding rules ******** def simplify_pow(x:UOp, c:UOp) -> UOp|None: - if c.arg < 0: return x.reciprocal().pow(-c) - if c.arg == 0: return x.const_like(1) - if int(c.arg-0.5)+0.5 == c.arg: return x.pow(c.const_like(c.arg-0.5)) * x.sqrt() - if int(c.arg) == c.arg: return (y := x.pow(c.const_like(c.arg//2))) * y * (x if c.arg%2 == 1 else 1) + if c.val < 0: return x.reciprocal().pow(-c) + if c.val == 0: return x.const_like(1) + if int(c.val-0.5)+0.5 == c.val: return x.pow(c.const_like(c.val-0.5)) * x.sqrt() + if int(c.val) == c.val: return (y := x.pow(c.const_like(c.val//2))) * y * (x if c.val%2 == 1 else 1) return None def fold_bitcast(root:UOp, c:UOp) -> UOp|None: @@ -110,8 +110,8 @@ symbolic_simple = pm_data_invalid + PatternMatcher([ ((UPat.var() % UPat.var("y")).named("base") % UPat.var("y"), lambda base,y: base), # (x%y)%y = -> x%y (rewritten with base for speed) # variations of (x%c)+(x//c)*c = x (UPat(Ops.ADD, dtype=dtypes.weakint, name="x"), fold_add_divmod_recombine), - (UPat.var("x", dtype=dtypes.bool) & UPat.cvar("c"), lambda x,c: x if c.arg else c), - (UPat.var("x", dtype=dtypes.bool) | UPat.cvar("c"), lambda x,c: c if c.arg else x), + (UPat.var("x", dtype=dtypes.bool) & UPat.cvar("c"), lambda x,c: x if c.val else c), + (UPat.var("x", dtype=dtypes.bool) | UPat.cvar("c"), lambda x,c: c if c.val else x), (UPat.var("x", dtype=dtypes.bool) != UPat.const(False, dtypes.bool), lambda x: x), # x != False -> x (UPat(GroupOp.Idempotent, src=(UPat.var("x"), UPat.var("x"))), lambda x: x), (UPat.var("x", dtype=dtypes.bool).logical_not().logical_not(), lambda x: x), @@ -153,10 +153,10 @@ symbolic_simple = pm_data_invalid + PatternMatcher([ # if x is nan or inf it should render the nan value. # NOTE: this can be wrong for loaded NaN (UPat.var("x") * 0, lambda x: x.const_like(float("nan") if x.op is Ops.CONST - and isinstance(x.arg, float) and (math.isnan(x.arg) or math.isinf(x.arg)) else 0)), + and isinstance(x.val, float) and (math.isnan(x.val) or math.isinf(x.val)) else 0)), # *** cast/bitcast *** # TODO: delete this once CONST has no dtype - (UPat(Ops.CAST, name="root", src=(UPat.cvar("c"),)), lambda root, c: root.const_like(c.arg)), + (UPat(Ops.CAST, name="root", src=(UPat.cvar("c"),)), lambda root, c: root.const_like(c.val)), (UPat((Ops.CAST, Ops.BITCAST), name="root"), lambda root: root.src[0] if root.dtype == root.src[0].dtype else None), (UPat(Ops.BITCAST, name="root", src=(UPat.cvar("c"),)), fold_bitcast), # b.cast(a).cast(b) -> b if a preserves all values in b @@ -167,7 +167,7 @@ symbolic_simple = pm_data_invalid + PatternMatcher([ # ** pow ** (UPat.var("x").alu(Ops.POW, UPat.cvar("c")), simplify_pow), # positive const ** x - (UPat.cvar("c").alu(Ops.POW, UPat.var("x")), lambda c,x: c if c.arg == 1 else (x*math.log2(c.arg)).exp2() if c.arg > 0 else None), + (UPat.cvar("c").alu(Ops.POW, UPat.var("x")), lambda c,x: c if c.val == 1 else (x*math.log2(c.val)).exp2() if c.val > 0 else None), # rules for threefry ((UPat.var('x', dtypes.uint64)&0xFFFFFFFF).cast(dtypes.uint32), lambda x: x.cast(dtypes.uint32)), (((UPat.var(None, dtypes.uint64)*(1<<32)) | UPat.var('y', dtypes.uint32).cast(dtypes.uint64)).cast(dtypes.uint32), lambda y: y), @@ -178,7 +178,7 @@ symbolic_simple = pm_data_invalid + PatternMatcher([ # ** simple where folding ** # a conditional with the same results either way is a noop, also fold const conditionals (UPat.var().where(UPat.var("val"), UPat.var("val")), lambda val: val), - (UPat.cvar("gate").where(UPat.var("c0"), UPat.var("c1")), lambda gate, c0, c1: c0 if gate.arg else c1), + (UPat.cvar("gate").where(UPat.var("c0"), UPat.var("c1")), lambda gate, c0, c1: c0 if gate.val else c1), # a.where(b.where(c, d), d) -> (a & b).where(c, d) (UPat.var("a").where(UPat.var("b").where(UPat.var("c"), UPat.var("d")), UPat.var("d")), lambda a,b,c,d: (a&b).where(c,d)), # a.where(c, b.where(c, d)) -> (a | b).where(c, d) @@ -274,7 +274,7 @@ symbolic = symbolic_simple+commutative+PatternMatcher([ ((UPat.var("x") * UPat.cvar("c1")) * UPat.var("y"), lambda x,c1,y: (x*y)*c1 if y.op is not Ops.CONST else None), # *** rules from symbolic *** # generic lt folding - (UPat.var("x", dtypes.weakint) 0. NOTE: not x < 1 means x > 0 ((UPat.var("x", dtypes.weakint)<1).ne(True), lambda x: (newx<1).ne(True) if (newx:=canonicalize_simplex(x)) is not None else None), @@ -453,7 +453,7 @@ sym = symbolic+pm_simplify_valid+PatternMatcher([ (UPat.var("x") * ((1+UPat.var("x")).reciprocal().named("d")*UPat.var("y")), lambda x,y,d: y*(1-d)), (UPat.var("x") * ((1+UPat.var("x")).reciprocal().named("d")+UPat.var("y")), lambda x,y,d: (1-d)+x*y), # move const multiply after REDUCE (NOTE: the mul chain can do this, but only if it's a same dtype reduce) - ((UPat.var("x")*UPat.cvar("c")).reduce(arg=Ops.ADD, name="r", allow_any_len=True), lambda x,c,r: r.replace(src=(x,)+r.src[1:])*c.arg), + ((UPat.var("x")*UPat.cvar("c")).reduce(arg=Ops.ADD, name="r", allow_any_len=True), lambda x,c,r: r.replace(src=(x,)+r.src[1:])*c.val), # reduce mul chain, move muls after the reduce (UPat(Ops.MUL).reduce(name="r", allow_any_len=True), reduce_mul_chain), # ** combine terms (opinionated) ** diff --git a/tinygrad/uop/upat.py b/tinygrad/uop/upat.py index 2eb5cf18ca..9bbcab452f 100644 --- a/tinygrad/uop/upat.py +++ b/tinygrad/uop/upat.py @@ -118,7 +118,7 @@ pm_renderer = PatternMatcher([ lambda r,x: r.replace(src=(UOp(Ops.CUSTOMI, arg="(" + ' and '.join(y.arg for y in x.src) + ")"),)+r.src[1:])), (UPat(Ops.CUSTOM, src=UPat(Ops.CUSTOMI), name="x"), lambda x: UOp(Ops.CUSTOMI, arg=x.arg.format(*[y.arg for y in x.src]))), - (UPat(Ops.INDEX, src=(UPat(Ops.CUSTOMI, name="x"), UPat(Ops.CONST, name="c")), name="g"), lambda x,c,g: x.replace(arg=x.arg+f".src[{c.arg}]")) + (UPat(Ops.INDEX, src=(UPat(Ops.CUSTOMI, name="x"), UPat(Ops.CONST, name="c")), name="g"), lambda x,c,g: x.replace(arg=x.arg+f".src[{c.val}]")) ], compiled=False) def _final_render(x:UOp, has_ctx:bool, depth=1) -> list[str]: From 20b8ecff50ddfb2b64f3bfe82ed87f396667e5c3 Mon Sep 17 00:00:00 2001 From: Adeeb Shihadeh Date: Fri, 31 Jul 2026 21:45:06 -0700 Subject: [PATCH 35/44] support new USB vendor ID (#17348) --- tinygrad/runtime/ops_amd.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tinygrad/runtime/ops_amd.py b/tinygrad/runtime/ops_amd.py index a9a1d610f3..7404a1fca3 100644 --- a/tinygrad/runtime/ops_amd.py +++ b/tinygrad/runtime/ops_amd.py @@ -909,7 +909,7 @@ class PCIIface(PCIIfaceBase): class USBIface(PCIIface): def __init__(self, dev, dev_id): # pylint: disable=super-init-not-called - if dev_id >= len(visible:=hcq_filter_visible_devices(USB3.list_devices(0xADD1, 0x0001), "AMD")): + if dev_id >= len(visible:=hcq_filter_visible_devices(USB3.list_devices(0xADD1, 0x0001) + USB3.list_devices(0x3801, 0x0001), "AMD")): raise RuntimeError(f"AMD:{dev_id} does not exist ({pluralize('device', len(visible))} available)") self.dev, self.pci_dev, self.vram_bar, self.count = dev, USBPCIDevice("AM", *visible[dev_id]), 0, len(visible) self.dev_impl = AMDev(self.pci_dev) From 5a1c641f791cc987ddd471d42472fcf1b2dfa633 Mon Sep 17 00:00:00 2001 From: George Hotz <72895+geohot@users.noreply.github.com> Date: Fri, 31 Jul 2026 22:44:24 -0700 Subject: [PATCH 36/44] more arg -> val (#17349) * more arg -> val * kimi * more --- tinygrad/callify.py | 2 +- tinygrad/codegen/simplify.py | 2 +- tinygrad/engine/realize.py | 4 ++-- tinygrad/function.py | 2 +- tinygrad/renderer/cstyle.py | 2 +- tinygrad/renderer/isa/x86.py | 28 ++++++++++++++-------------- tinygrad/renderer/nir.py | 4 ++-- tinygrad/runtime/graph/metal.py | 2 +- tinygrad/runtime/support/hcq2.py | 2 +- tinygrad/schedule/__init__.py | 2 +- tinygrad/uop/divandmod.py | 6 +++--- tinygrad/uop/movement.py | 4 ++-- tinygrad/uop/ops.py | 6 +++--- tinygrad/uop/render.py | 10 +++++----- tinygrad/uop/spec.py | 4 ++-- tinygrad/uop/symbolic.py | 30 +++++++++++++++--------------- tinygrad/uop/validate.py | 4 ++-- tinygrad/viz/serve.py | 4 ++-- 18 files changed, 59 insertions(+), 59 deletions(-) diff --git a/tinygrad/callify.py b/tinygrad/callify.py index 5ff272a4ec..b7bb8d4972 100644 --- a/tinygrad/callify.py +++ b/tinygrad/callify.py @@ -57,7 +57,7 @@ def _make_buffer_view(src:UOp) -> UOp|None: if (offset := src.contiguous_view_offset()) is None: return None buf = src.base if buf.op is Ops.SLICE: - byte_offset = buf.src[1].arg * buf.src[0].dtype.itemsize + offset * src.dtype.itemsize + byte_offset = buf.src[1].val * buf.src[0].dtype.itemsize + offset * src.dtype.itemsize buf = buf.src[0] if byte_offset % buf.dtype.itemsize != 0: return None offset = byte_offset // buf.dtype.itemsize diff --git a/tinygrad/codegen/simplify.py b/tinygrad/codegen/simplify.py index 608a656d25..4aadeb6d6f 100644 --- a/tinygrad/codegen/simplify.py +++ b/tinygrad/codegen/simplify.py @@ -47,7 +47,7 @@ def mark_gated(ctx, idx): guards = {r:c for v in cond.split_uop(Ops.AND) if v.op is Ops.CMPLT and (r:=v.src[0]).op is Ops.RANGE and (c:=v.src[1]).op is Ops.CONST} else: x, guards = idx, {} # ensure that we choose max(c_i) for all i where r < c_i - ctx |= {r:c for r,c in guards.items() if (r not in ctx or ctx[r].arg < c.arg)} + ctx |= {r:c for r,c in guards.items() if (r not in ctx or ctx[r].val < c.val)} # but if a range is ever ungated, we cannot shrink it ctx |= {r:r.src[0] for r in x.ranges if r not in guards} diff --git a/tinygrad/engine/realize.py b/tinygrad/engine/realize.py index a5f8155162..91b9e331a6 100644 --- a/tinygrad/engine/realize.py +++ b/tinygrad/engine/realize.py @@ -28,7 +28,7 @@ def get_call_name(call:UOp, bufs:Sequence[Buffer|UOp], var_vals:dict[str, int]|N ast, arg_uops = call.src[0], get_call_arg_uops(call) if ast.op is Ops.PROGRAM: return ast.arg.name if ast.op is Ops.SLICE: - offset = ast.src[1].arg * arg_uops[1].dtype.itemsize + offset = ast.src[1].val * arg_uops[1].dtype.itemsize return colored(f"view {_uop_sz_to_str(arg_uops[0]):>10} @ {offset:<10d}", "yellow") if ast.op is Ops.COPY: return colored(f"copy {_uop_sz_to_str(arg_uops[0]):>10}, {_dev_str(bufs[0]):>7s} <- {_dev_str(bufs[1]):7s}", "yellow") if ast.op is Ops.CUSTOM_FUNCTION and ast.arg == "encdec": return colored(f"enc/dec {_uop_sz_to_str(arg_uops[0])}", "yellow") @@ -156,7 +156,7 @@ def unwrap_multi(call:UOp, resolved:list[UOp]) -> Iterator[tuple[list[Buffer], d def exec_view(ctx:ExecContext, call:UOp, ast:UOp) -> float|None: resolved = resolve_params(call, ctx.input_uops) bufs = [cast(Buffer, b.buffer) for b in resolved] - bv = bufs[1].view(resolved[0].max_numel(), ast.dtype, ast.src[1].arg*bufs[1].dtype.itemsize) + bv = bufs[1].view(resolved[0].max_numel(), ast.dtype, ast.src[1].val*bufs[1].dtype.itemsize) with track_stats(ctx, call, bv.device, [bv, bufs[1]], ctx.var_vals): buffers[resolved[0]] = bv return None diff --git a/tinygrad/function.py b/tinygrad/function.py index e63adaad70..8b540e3381 100644 --- a/tinygrad/function.py +++ b/tinygrad/function.py @@ -22,7 +22,7 @@ def invalid_outputs(uret:UOp) -> set[UOp]: # invalids() returns fresh write-only scratch: a clone storing CONST(Invalid) # don't capture it as an input; only skip fresh buffers, not realized ones return {u.src[0].buf_uop for u in uret.backward_slice_with_self - if u.op is Ops.STORE and u.src[1].base.op is Ops.CONST and u.src[1].base.arg is Invalid + if u.op is Ops.STORE and u.src[1].base.op is Ops.CONST and u.src[1].base.val is Invalid and not u.src[0].buf_uop.is_realized} ReturnType = TypeVar('ReturnType') diff --git a/tinygrad/renderer/cstyle.py b/tinygrad/renderer/cstyle.py index 6dbf090bb4..e25ab0c605 100644 --- a/tinygrad/renderer/cstyle.py +++ b/tinygrad/renderer/cstyle.py @@ -320,7 +320,7 @@ class OpenCLRenderer(CStyleLanguage): (UPat(Ops.BITCAST, name="x"), lambda ctx,x: f"as_{ctx.render_dtype(x.dtype)}(({ctx.render_dtype(x.src[0].dtype)})({ctx[x.src[0]]}))"), # bfloat16 constants need to be rendered as their bit pattern since bf16 is stored as ushort (UPat(Ops.CONST, dtypes.bfloat16, name="x"), - lambda ctx,x: f"{(struct.unpack('I', struct.pack('f', float_to_bf16(x.arg)))[0] >> 16)}u"), + lambda ctx,x: f"{(struct.unpack('I', struct.pack('f', float_to_bf16(x.val)))[0] >> 16)}u"), # load/store image (OpenCL) (UPat.var('buf').index(UPat.var('idx_y'), UPat.var('idx_x')), lambda ctx,buf,idx_y,idx_x: f"IMAGE<{ctx[buf]}, {ctx[idx_y]}, {ctx[idx_x]}>"), (UPat(Ops.LOAD, dtype=dtypes.float, src=(UPat.var('buf').index(UPat.var('idx_y'), UPat.var('idx_x')), UPat.var("var"), UPat.var("gate"))), diff --git a/tinygrad/renderer/isa/x86.py b/tinygrad/renderer/isa/x86.py index 4ea41f643c..1346a02678 100644 --- a/tinygrad/renderer/isa/x86.py +++ b/tinygrad/renderer/isa/x86.py @@ -221,15 +221,15 @@ reg_strs = {"rax": {4:"eax", 2:"ax", 1:"al"}, "rcx": {4:"ecx", 2:"cx", 1:"cl"}, # ***** X86 instruction selection ***** def base(x:UOp, i:int) -> UOp: return s.src[0] if (s:=x.src[i]).op is Ops.INDEX else s -def lane(x:UOp, i:int) -> int: return s.src[1].arg if (s:=x.src[i]).op is Ops.INDEX else 0 +def lane(x:UOp, i:int) -> int: return s.src[1].val if (s:=x.src[i]).op is Ops.INDEX else 0 def to_int(dt:DType): return {dtypes.float16: dtypes.int16, dtypes.float32: dtypes.int32, dtypes.float64: dtypes.int64}[dt] def def_reg(dt:DType, reg:Register|None=None) -> UOp: return UOp(Ops.INS, dt, arg=X86Ops.DEFINE, tag=None if reg is None else (reg,)) def imm(dt:DType, v:int) -> UOp: return UOp.const(truncate[dt](v), dt).rtag() def to_imm(c:UOp) -> UOp|None: if c.op is not Ops.CONST: return None - if c.dtype is dtypes.int64: return imm(dtypes.int32, c.arg) if not c.overflows(dtypes.int32) else None - if c.dtype is dtypes.uint64: return imm(dtypes.uint32, c.arg) if not c.overflows(dtypes.uint32) else None - if c.dtype in dtypes.ints+(dtypes.bool,): return imm(c.dtype, c.arg) + if c.dtype is dtypes.int64: return imm(dtypes.int32, c.val) if not c.overflows(dtypes.int32) else None + if c.dtype is dtypes.uint64: return imm(dtypes.uint32, c.val) if not c.overflows(dtypes.uint32) else None + if c.dtype in dtypes.ints+(dtypes.bool,): return imm(c.dtype, c.val) return None def cmp(x:UOp) -> UOp: if x.src[0].dtype is dtypes.float32: return x.ins(X86Ops.VUCOMISS, dtype=dtypes.void) @@ -421,15 +421,15 @@ isel_matcher = PatternMatcher([ (UPat(Ops.STACK, dtypes.ints+(dtypes.bool,), name="x"), vpins), # INDEX on a vector register value extracts a single element (UPat.var("y", dtypes.int8s+(dtypes.bool,)).index(UPat.cvar("c"), name="x"), - lambda y,c,x: x.ins(X86Ops.VPEXTRB, src=(y, imm(dtypes.uint8, c.arg))) if _is_vec_xmm(y) else None), + lambda y,c,x: x.ins(X86Ops.VPEXTRB, src=(y, imm(dtypes.uint8, c.val))) if _is_vec_xmm(y) else None), (UPat.var("y", dtypes.int16s).index(UPat.cvar("c"), name="x"), - lambda y,c,x: x.ins(X86Ops.VPEXTRW, src=(y, imm(dtypes.uint8, c.arg))) if _is_vec_xmm(y) else None), + lambda y,c,x: x.ins(X86Ops.VPEXTRW, src=(y, imm(dtypes.uint8, c.val))) if _is_vec_xmm(y) else None), (UPat.var("y", dtypes.int32s).index(UPat.cvar("c"), name="x"), - lambda y,c,x: x.ins(X86Ops.VPEXTRD, src=(y, imm(dtypes.uint8, c.arg))) if _is_vec_xmm(y) else None), + lambda y,c,x: x.ins(X86Ops.VPEXTRD, src=(y, imm(dtypes.uint8, c.val))) if _is_vec_xmm(y) else None), (UPat.var("y", dtypes.int64s).index(UPat.cvar("c"), name="x"), - lambda y,c,x: x.ins(X86Ops.VPEXTRQ, src=(y, imm(dtypes.uint8, c.arg))) if _is_vec_xmm(y) else None), + lambda y,c,x: x.ins(X86Ops.VPEXTRQ, src=(y, imm(dtypes.uint8, c.val))) if _is_vec_xmm(y) else None), (UPat.var("y", dtypes.floats).index(UPat.cvar("c"), name="x"), - lambda y,c,x: x.ins(X86Ops.VPSRLDQ, src=(y, imm(dtypes.uint8, c.arg * x.dtype.itemsize))) if _is_vec_xmm(y) else None), + lambda y,c,x: x.ins(X86Ops.VPSRLDQ, src=(y, imm(dtypes.uint8, c.val * x.dtype.itemsize))) if _is_vec_xmm(y) else None), # packed bitwise ((UPat() & UPat()).named("x"), lambda x: x.ins(X86Ops.VPAND) if x.max_numel() > 1 else None), ((UPat() | UPat()).named("x"), lambda x: x.ins(X86Ops.VPOR) if x.max_numel() > 1 else None), @@ -614,7 +614,7 @@ def encode(x:UOp, opc:int, reg:int|None=None, pp:int=0, sel:int=0, we:int=0) -> rm = cast(Register, greg(rm_uop)).index idx = cast(Register, greg(idx_uop)).index if idx_uop is not None and greg(idx_uop) is not None else 4 # for a memory operand the rm size is the element size from the address, otherwise it's the size of the value in the register - rm_sz = sz_uop.arg if sz_uop is not None else rm_uop.dtype.itemsize + rm_sz = sz_uop.val if sz_uop is not None else rm_uop.dtype.itemsize reg_sz = reg_uop.dtype.itemsize if reg_uop is not None else 0 sz = reg_sz or rm_sz @@ -650,7 +650,7 @@ def encode(x:UOp, opc:int, reg:int|None=None, pp:int=0, sel:int=0, we:int=0) -> assert disp_uop.op is Ops.CONST, "displacement must be a constant" assert disp_uop.dtype in (dtypes.int8, dtypes.int32), "displacement can only be 1 or 4 byte signed int" # rbp/r13 always require a displacement - if disp_uop.arg != 0 or rm == 0b101: mod = 0b01 if disp_uop.dtype.itemsize == 1 else 0b10 + if disp_uop.val != 0 or rm == 0b101: mod = 0b01 if disp_uop.dtype.itemsize == 1 else 0b10 else: mod = 0b00 else: mod = 0b11 # x 0b0 and idx 0b100 means rsp which means no index exists @@ -701,7 +701,7 @@ def encode(x:UOp, opc:int, reg:int|None=None, pp:int=0, sel:int=0, we:int=0) -> encodings = { # moves X86Ops.MOVABS: lambda x: - bytes([0b0100 << 4 | 0b1 << 3 | 0b00 << 2 | greg(x).index >> 3, 0xB8 + (greg(x).index & 0b111)]) + struct.pack(x.dtype.fmt, x.src[0].arg), + bytes([0b0100 << 4 | 0b1 << 3 | 0b00 << 2 | greg(x).index >> 3, 0xB8 + (greg(x).index & 0b111)]) + struct.pack(x.dtype.fmt, x.src[0].val), X86Ops.MOV: lambda x: encode(x, 0x8B), X86Ops.MOVi: lambda x: encode(x, 0xC7, reg=0), X86Ops.MOVm: lambda x: encode(x, 0x89), X86Ops.LEA: lambda x: encode(x, 0x8D), X86Ops.VMOVSS: lambda x: encode(x, 0x10, pp=2, sel=1), X86Ops.VMOVSSm: lambda x: encode(x, 0x11, pp=2, sel=1), @@ -724,8 +724,8 @@ encodings = { X86Ops.VCVTPS2PD: lambda x: encode(x, 0x5A, pp=0, sel=1), X86Ops.VCVTPD2PS: lambda x: encode(x, 0x5A, pp=1, sel=1), X86Ops.VCVTTPS2DQ: lambda x: encode(x, 0x5B, pp=2, sel=1), X86Ops.VCVTTPD2DQ: lambda x: encode(x, 0xE6, pp=1, sel=1), # the int src is the 2nd src (the rm field), if it was folded into a memory operand its width is the element size of the address - X86Ops.VCVTSI2SS: lambda x: encode(x, 0x2A, pp=2, sel=1, we=(x.src[4].arg if len(x.src) > 4 else x.src[1].dtype.itemsize) == 8), - X86Ops.VCVTSI2SD: lambda x: encode(x, 0x2A, pp=3, sel=1, we=(x.src[4].arg if len(x.src) > 4 else x.src[1].dtype.itemsize) == 8), + X86Ops.VCVTSI2SS: lambda x: encode(x, 0x2A, pp=2, sel=1, we=(x.src[4].val if len(x.src) > 4 else x.src[1].dtype.itemsize) == 8), + X86Ops.VCVTSI2SD: lambda x: encode(x, 0x2A, pp=3, sel=1, we=(x.src[4].val if len(x.src) > 4 else x.src[1].dtype.itemsize) == 8), X86Ops.VCVTTSS2SI: lambda x: encode(x, 0x2C, pp=2, sel=1, we=x.dtype.itemsize == 8), X86Ops.VCVTTSD2SI: lambda x: encode(x, 0x2C, pp=3, sel=1, we=x.dtype.itemsize == 8), # int division diff --git a/tinygrad/renderer/nir.py b/tinygrad/renderer/nir.py index a531007ff7..b0d3a38f2c 100644 --- a/tinygrad/renderer/nir.py +++ b/tinygrad/renderer/nir.py @@ -185,7 +185,7 @@ 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].arg + 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 self.r: dict[UOp, Any] = {} self.param_idx = 0 ranges: list[mesa.nir_def|None] = [] @@ -194,7 +194,7 @@ class NIRRenderer(Renderer): if u.op in {Ops.NOOP, Ops.GROUP} 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].arg) + 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) elif u.op is Ops.AFTER: self.r[u] = self.r[u.src[0]] elif u.op == Ops.SINK: diff --git a/tinygrad/runtime/graph/metal.py b/tinygrad/runtime/graph/metal.py index 9094f36692..409cfb973d 100644 --- a/tinygrad/runtime/graph/metal.py +++ b/tinygrad/runtime/graph/metal.py @@ -113,5 +113,5 @@ class MetalGraph(GraphRunner): @staticmethod def supports_uop(batch_devs, new_call:UOp) -> bool: # Metal ICB replay encodes offsets as uint32; reject if any Metal buffer offset exceeds 32-bit range. - if any(b.op is Ops.SLICE and b.src[1].arg * b.src[0].dtype.itemsize > 0xFFFFFFFF for b in new_call.src[1:]): return False + if any(b.op is Ops.SLICE and b.src[1].val * b.src[0].dtype.itemsize > 0xFFFFFFFF for b in new_call.src[1:]): return False return GraphRunner.supports_uop(batch_devs, new_call) diff --git a/tinygrad/runtime/support/hcq2.py b/tinygrad/runtime/support/hcq2.py index 498703bc3a..0f36f516de 100644 --- a/tinygrad/runtime/support/hcq2.py +++ b/tinygrad/runtime/support/hcq2.py @@ -334,7 +334,7 @@ pm_replace_params = PatternMatcher([ def resolve_getaddr_slice(bv:UOp, g:UOp) -> UOp: base = bv.src[0].after(*g.src[0].src[1:] if g.src[0].op is Ops.AFTER else ()) itemsize = bv.src[0].dtype.itemsize if bv.src[0].without_after.op in (Ops.BUFFER, Ops.SLICE, Ops.MSTACK, Ops.MSELECT) else bv.dtype.itemsize - return UOp(Ops.GETADDR, src=(base,), arg=g.arg) + UOp.const(bv.src[1].arg * itemsize, dtypes.uint64) + return UOp(Ops.GETADDR, src=(base,), arg=g.arg) + UOp.const(bv.src[1].val * itemsize, dtypes.uint64) pm_early_simplify = PatternMatcher([ (UPat(Ops.GETADDR, src=(UPat.any(sl:=UPat(Ops.SLICE, name="bv"), sl.after(allow_any_len=True)),), name="g"), resolve_getaddr_slice), diff --git a/tinygrad/schedule/__init__.py b/tinygrad/schedule/__init__.py index 0845e2e354..99d228945f 100644 --- a/tinygrad/schedule/__init__.py +++ b/tinygrad/schedule/__init__.py @@ -186,7 +186,7 @@ def create_linear_with_vars(big_sink:UOp) -> tuple[UOp, dict[str, int]]: if b.op is Ops.BIND: nm = b.src[0].expr if nm not in used_vars: continue - val = b.src[1].arg + val = b.src[1].val if var_vals.get(nm, val) != val: raise RuntimeError(f"bind mismatch on {nm}, {var_vals[nm]} != {val}") var_vals[nm] = val diff --git a/tinygrad/uop/divandmod.py b/tinygrad/uop/divandmod.py index bb082bc3d3..03dbd1e138 100644 --- a/tinygrad/uop/divandmod.py +++ b/tinygrad/uop/divandmod.py @@ -87,7 +87,7 @@ def fold_divmod_general(d: UOp) -> UOp|None: if (q:=u.divide_exact(y)) is not None: quo.append(q) elif y.op is Ops.CONST and (c:=u.const_factor())%y.val!=c: rem.append(u.divides(c)*(c%y.val)) - quo.append(u.divides(c)*(c//y.arg) if d.op is Ops.FLOORDIV else u.const_like(0)) + quo.append(u.divides(c)*(c//y.val) if d.op is Ops.FLOORDIV else u.const_like(0)) else: rem.append(u) if not quo: return None @@ -101,8 +101,8 @@ div_and_mod_symbolic = PatternMatcher([ ((UPat.var("x")//UPat.cvar("c") + UPat.cvar("a"))//UPat.cvar("d"), lambda x,c,a,d: (x+a*c)//(c*d) if d.vmin>0 else None), # (x+c)//d -> (x+c%d)//d + c//d ; (x+c)%d -> (x+c%d)%d (split the multiple of d out of the const, holds for any d!=0) (UPat((Ops.FLOORDIV, Ops.FLOORMOD), src=(UPat.var("x", dtypes.weakint)+UPat.cvar("c"), UPat.cvar("d")), name="n"), - lambda n,x,c,d: None if d.arg==0 or c.arg%d.arg==c.arg else - (x+c.arg%d.arg)//d + c.arg//d.arg if n.op is Ops.FLOORDIV else (x+c.arg%d.arg)%d), + lambda n,x,c,d: None if d.val==0 or c.val%d.val==c.val else + (x+c.val%d.val)//d + c.val//d.val if n.op is Ops.FLOORDIV else (x+c.val%d.val)%d), # ** 2. Slow Rules ** (UPat((Ops.FLOORDIV, Ops.FLOORMOD), dtypes.weakint, name="d"), lambda d: fold_divmod_general(d)), diff --git a/tinygrad/uop/movement.py b/tinygrad/uop/movement.py index 686f7a7ae7..295efced07 100644 --- a/tinygrad/uop/movement.py +++ b/tinygrad/uop/movement.py @@ -13,10 +13,10 @@ mop_cleanup = PatternMatcher([ (UPat(Ops.PERMUTE, name="x"), lambda x: x.src[0] if list(x.arg) == list(range(len(x.arg))) else None), # STACK on INDEX CONST (UPat(Ops.STACK, src=UPat(Ops.INDEX, src=(UPat.var("src"), UPat(Ops.CONST))), name="stk"), - lambda src,stk: src if stk.shape == src.shape and list(range(len(stk.src))) == [x.src[1].arg for x in stk.src] else None), + lambda src,stk: src if stk.shape == src.shape and list(range(len(stk.src))) == [x.src[1].val for x in stk.src] else None), # const INDEX into STACK is src (UPat(Ops.INDEX, src=(UPat(Ops.STACK, name="a"), UPat.cvar("i")), name="idx", allow_any_len=True), - lambda a,i,idx: a.src[i.arg] if len(idx.src) <= 2 else a.src[i.arg].index(*idx.src[2:])), + lambda a,i,idx: a.src[i.val] if len(idx.src) <= 2 else a.src[i.val].index(*idx.src[2:])), # INDEX on INDEX is INDEX (UPat(Ops.INDEX, src=(UPat(Ops.INDEX, name="idx1", allow_any_len=True),), allow_any_len=True, name="idx2"), lambda idx1,idx2: idx1.src[0].index(*idx1.src[1:], *idx2.src[1:]) if all(x.shape == () for x in idx1.src[1:]+idx2.src[1:]) else None), diff --git a/tinygrad/uop/ops.py b/tinygrad/uop/ops.py index 17899ce79f..027508e2be 100644 --- a/tinygrad/uop/ops.py +++ b/tinygrad/uop/ops.py @@ -567,7 +567,7 @@ class UOp(RandMixin, metaclass=UOpMetaClass): return UOp(Ops.GROUP, src=tuple([x for x in srcs if x is not None])) def index(self, *srcs:UOp|int|None, **kwargs): new_srcs: list[UOp] = [UOp.const(x) if isinstance(x, int) else x for x in srcs if x is not None] - if len(new_srcs) == 1 and new_srcs[0].op is Ops.CONST and self.op is Ops.STACK: return self.src[new_srcs[0].arg] + if len(new_srcs) == 1 and new_srcs[0].op is Ops.CONST and self.op is Ops.STACK: return self.src[new_srcs[0].val] return UOp(Ops.INDEX, src=(self,)+tuple(new_srcs), **kwargs) def __getitem__(self, idx): # buffers index into INDEX UOps (scalar lookup); everything else uses the shared mixin view path @@ -931,7 +931,7 @@ class UOp(RandMixin, metaclass=UOpMetaClass): if self.op is Ops.SLICE: if (cret:=buffers.get(self)) is not None: return cret buf = self.src[0].buffer - offset = self.src[1].arg + offset = self.src[1].val if isinstance(buf, MultiBuffer): mbuf = MultiBuffer.__new__(MultiBuffer) mbuf.bufs = [b.view(self.arg, self.dtype, offset * self.src[0].dtype.itemsize) for b in buf.bufs] @@ -1787,7 +1787,7 @@ def lower_weak_srcs(ctx:dict[UOp, UOp]|None, u:UOp) -> UOp|None: def commit_weak(s:UOp, dt:DType) -> UOp: # a bare weak CONST commits directly (its number must fit), a weak non-const src takes the demand cast - return UOp.const(s.arg, dt) if s.op is Ops.CONST else s.cast(dt) + return UOp.const(s.val, dt) if s.op is Ops.CONST else s.cast(dt) def commit_weak_srcs(u:UOp) -> UOp|None: if (dt:=least_upper_dtype(*(s.dtype for s in u.src))) in dtypes.weaks: return None diff --git a/tinygrad/uop/render.py b/tinygrad/uop/render.py index 97ebe79db4..c8467ae654 100644 --- a/tinygrad/uop/render.py +++ b/tinygrad/uop/render.py @@ -18,7 +18,7 @@ def pretty_print(x:UOp, cache=None, d=0)->str: def print_uops(uops:list[UOp]): uops_index = {u:i for i,u in enumerate(uops)} for i,u in enumerate(uops): - formatted_srcs = [(uops_index[x] if x.op is not Ops.CONST else f"{x.arg}") if x in uops else "--" for x in u.src] + formatted_srcs = [(uops_index[x] if x.op is not Ops.CONST else f"{x.val}") if x in uops else "--" for x in u.src] print(f"{i:4d} {str(u.op):20s}: {multirange_str(u.ranges, color=True, pad=10)} {str(u.dtype):40s} " f"{str(formatted_srcs):32s} {u.arg}") # for debug @@ -36,7 +36,7 @@ renderer = PatternMatcher([ (UPat((Ops.SPECIAL), name="x"), lambda x: x.arg), (UPat(Ops.RANGE, dtypes.void, name="x"), lambda x: f"loop{x.arg[0]}"), (UPat(Ops.RANGE, name="x"), lambda x: f"r{range_str(x)}"), - (UPat(Ops.CONST, name="x"), lambda x: str(x.arg)), + (UPat(Ops.CONST, name="x"), lambda x: str(x.val)), (UPat(Ops.CAST, name="x"), lambda ctx,x: f"({str(x.dtype)[7:]})({ctx[x.src[0]]})"), (UPat(Ops.BIND, name="x"), lambda ctx,x: ctx[x.src[0]]), (UPat(Ops.NEG, name="x"), lambda ctx,x: f"(-{ctx[x.src[0]]})"), @@ -79,9 +79,9 @@ def render_marg(ctx,x:UOp): sugar = {Ops.SINK, Ops.END, Ops.STORE, Ops.LOAD, Ops.SQRT, Ops.INDEX, Ops.REDUCE, Ops.AFTER, Ops.THREEFRY, Ops.RECIPROCAL, Ops.EXP2, Ops.LOG2, Ops.SIN, Ops.CONTIGUOUS, Ops.BARRIER, Ops.DETACH} pm_pyrender_extra = PatternMatcher([ - (UPat(Ops.CONST, src=(), name="x"), lambda x: f"UOp.const({x.arg}, {x.dtype})"), + (UPat(Ops.CONST, src=(), name="x"), lambda x: f"UOp.const({x.val}, {x.dtype})"), (UPat((Ops.CAST, Ops.BITCAST), name="x"), lambda ctx,x: f"{ctx[x.src[0]]}.{x.op.name.lower()}({x.dtype})"), - (UPat(Ops.SPECIAL, src=(UPat(Ops.CONST),), name="x"), lambda x: f"UOp.special({x.src[0].arg}, {repr(x.arg)}, dtype={x.dtype})"), + (UPat(Ops.SPECIAL, src=(UPat(Ops.CONST),), name="x"), lambda x: f"UOp.special({x.src[0].val}, {repr(x.arg)}, dtype={x.dtype})"), (UPat(Ops.BUFFER, src=(UPat(),), name="x"), lambda x: f"UOp.new_buffer({repr(x.arg.device)}, {x.max_numel()}, {x.dtype}, {x.arg.slot})" if isinstance(x.arg, ParamArg) and x.addrspace is AddrSpace.GLOBAL else None), @@ -90,7 +90,7 @@ pm_pyrender_extra = PatternMatcher([ (UPat(Ops.REDUCE, name="r"), lambda ctx,r: f"{ctx[r.src[0]]}._rop({r.arg[0]}, {tuple(range(r.arg[1]))})" if r.arg[1] else None), # NOTE: range has srcs sometimes after control flow (UPat(Ops.RANGE, src=(UPat(Ops.CONST, name="c"),), allow_any_len=True, name="x"), lambda ctx,x,c: - "UOp.range("+', '.join([str(c.arg)] + [repr(y) for y in x.arg])+ + "UOp.range("+', '.join([str(c.val)] + [repr(y) for y in x.arg])+ (f', src={srcs(ctx, x.src[1:])}' if len(x.src) > 1 else '')+(', dtype='+str(x.dtype) if x.dtype is not dtypes.weakint else '')+")"), # TODO: index shouldn't mismatch dtype (UPat(Ops.INDEX, src=(UPat(), UPat()), allow_any_len=True, name="x"), lambda ctx,x: diff --git a/tinygrad/uop/spec.py b/tinygrad/uop/spec.py index b61b971f14..4f07447a60 100644 --- a/tinygrad/uop/spec.py +++ b/tinygrad/uop/spec.py @@ -10,7 +10,7 @@ from tinygrad.helpers import DEBUG, Context, SPEC, Metadata, panic, CHECK_OOB, a def validate_index(uidx:UOp, gate:UOp|None=None): if len(uidx.src) != 2: return True # skip for non final index. TODO: check more complex index with shape buf,idx = uidx.src - if idx.op is Ops.CONST and idx.arg is Invalid: return True + if idx.op is Ops.CONST and idx.val is Invalid: return True if gate is None: gate = UOp.const(True) # TODO: check for overflow if not CHECK_OOB or is_image_shape(buf._shape): return True @@ -54,7 +54,7 @@ spec_shared = PatternMatcher([ (UPat(Ops.NOOP), lambda: True), # CONST is everywhere - (UPat(Ops.CONST, src=(), name="x"), lambda x: type(x.arg) is type(x.dtype.const(x.arg))), + (UPat(Ops.CONST, src=(), name="x"), lambda x: type(x.val) is type(x.dtype.const(x.val))), # STACK is everywhere too (UPat(Ops.STACK, dtype=dtypes.void, src=()), lambda: True), diff --git a/tinygrad/uop/symbolic.py b/tinygrad/uop/symbolic.py index 3e21f6eb0b..4e9a5e1c66 100644 --- a/tinygrad/uop/symbolic.py +++ b/tinygrad/uop/symbolic.py @@ -23,11 +23,11 @@ def fold_bitcast(root:UOp, c:UOp) -> UOp|None: if (from_fmt:=c.dtype.fmt) is None or (to_fmt:=root.dtype.fmt) is None: return None if c.dtype.itemsize != root.dtype.itemsize: return None def convert(v:ConstType) -> ConstType: return struct.unpack(to_fmt, struct.pack(from_fmt, v))[0] - return root.const_like(convert(c.arg)) + return root.const_like(convert(c.val)) def const_arg(u:UOp) -> ConstType|tuple[ConstType, ...]|None: - if u.op is Ops.CONST: return u.arg - if u.op is Ops.STACK and all(s.op is Ops.CONST for s in u.src): return tuple(s.arg for s in u.src) + if u.op is Ops.CONST: return u.val + if u.op is Ops.STACK and all(s.op is Ops.CONST for s in u.src): return tuple(s.val for s in u.src) return None def fold_const_alu(a:UOp) -> UOp|None: @@ -39,8 +39,8 @@ def _quotient_base(q:UOp, base:UOp, div:int) -> UOp|None: # moves consts freely: the quotient may be merged ((x//c + a)//div -> (x + a*c)//(c*div) for div>0) and shifted ((y + k*D)//D == y//D + k) (q, s), (num, a) = q.pop_const(), base.pop_const() if q.op is not Ops.FLOORDIV or q.src[1].op is not Ops.CONST: return None - if div > 0 and num.op is Ops.FLOORDIV and num.src[1].op is Ops.CONST and q.src[1].arg == (c:=num.src[1].arg)*div: num, a, D = num.src[0], a*c, c*div - elif q.src[1].arg == div: D = div + if div > 0 and num.op is Ops.FLOORDIV and num.src[1].op is Ops.CONST and q.src[1].val == (c:=num.src[1].val)*div: num, a, D = num.src[0], a*c, c*div + elif q.src[1].val == div: D = div else: return None (x, xa), (p, pa) = num.pop_const(), q.src[0].pop_const() if p is not x or (t:=xa + a - pa) % D: return None @@ -54,13 +54,13 @@ def fold_add_divmod_recombine(x:UOp) -> UOp|None: for i,u in enumerate(terms): mod, mul = u.pop_const(Ops.MUL) if mod.op is not Ops.FLOORMOD or mod.src[1].op is not Ops.CONST: continue - base, div = mod.src[0], mod.src[1].arg + base, div = mod.src[0], mod.src[1].val for j,v in enumerate(terms): q, scale = v.pop_const(Ops.MUL) if i == j or scale != div*mul: continue rest = [t for k,t in enumerate(terms) if k not in (i,j)] if (b:=_quotient_base(q, base, div)) is not None: return (b*mul).usum(*rest) - if q.op is Ops.FLOORMOD and q.src[1].op is Ops.CONST and (d:=q.src[1].arg) > 0 and (b:=_quotient_base(q.src[0], base, div)) is not None: + if q.op is Ops.FLOORMOD and q.src[1].op is Ops.CONST and (d:=q.src[1].val) > 0 and (b:=_quotient_base(q.src[0], base, div)) is not None: return ((b % (div*d))*mul).usum(*rest) return None @@ -119,7 +119,7 @@ symbolic_simple = pm_data_invalid + PatternMatcher([ (UPat.var("x", dtype=dtypes.bool).where(UPat.const(False, dtypes.bool), UPat.const(True, dtypes.bool)), lambda x: x.logical_not()), # CAST(bool -> int) != const — CAST(True)=1, CAST(False)=0, so fold based on const value (UPat.var("x", dtype=dtypes.bool).cast(dtypes.ints+(dtypes.weakint,)) != UPat.cvar("c"), - lambda x,c: x if c.arg == 0 else x.logical_not() if c.arg == 1 else x.const_like(True)), + lambda x,c: x if c.val == 0 else x.logical_not() if c.val == 1 else x.const_like(True)), (UPat.var("x", dtype=dtypes.ints+(dtypes.bool, dtypes.weakint)).trunc(), lambda x: x), # ** zero folding ** (UPat.var("x") < UPat.var("x"), lambda x: x.const_like(False, dtypes.bool)), # x < x -> False @@ -129,9 +129,9 @@ symbolic_simple = pm_data_invalid + PatternMatcher([ # (x&mask)>>k -> x>>k when mask only clears bits below k # TODO: combine this with "# rules for threefry" below ((UPat.var("x") & UPat.cvar("mask")) >> UPat.cvar("k"), - lambda x,mask,k: x >> k.arg if mask.arg | ((1 << k.arg) - 1) == -1 else None), + lambda x,mask,k: x >> k.val if mask.val | ((1 << k.val) - 1) == -1 else None), ((UPat.var("x") & UPat.cvar("mask")) // UPat.cvar("c"), - lambda x,mask,c: x // c.arg if c.arg > 0 and c.arg & (c.arg-1) == 0 and mask.arg | (c.arg-1) == -1 else None), + lambda x,mask,c: x // c.val if c.val > 0 and c.val & (c.val-1) == 0 and mask.val | (c.val-1) == -1 else None), (UPat.var("x", dtype=dtypes.ints+(dtypes.bool, dtypes.weakint)) != UPat.var("x"), lambda x: x.const_like(False, dtypes.bool)), # x != x -> False (only ints) # ** constant folding ** @@ -199,7 +199,7 @@ def canonicalize_simplex(X:UOp) -> UOp|None: changed, ret = False, [] for u in X.split_uop(Ops.ADD): # assumed the const is the last src of MUL - if u.op is Ops.MUL and u.src[1].op is Ops.CONST and u.src[1].arg > 0: + if u.op is Ops.MUL and u.src[1].op is Ops.CONST and u.src[1].val > 0: changed = True u = u.src[0] if not (u.op in GroupOp.Irreducible and u.vmin >= 0): return None @@ -265,10 +265,10 @@ symbolic = symbolic_simple+commutative+PatternMatcher([ # ** lt ** # c0*x sign(c0)*x < ceil(c1/abs(c0)) ((UPat.cvar("c0")*UPat.var("x", dtype=dtypes.weakint)) 0 else -x)<-(-c1.arg//abs(c0.arg)) if abs(c0.arg) > 1 else None), + lambda x,c0,c1: (x if c0.val > 0 else -x)<-(-c1.val//abs(c0.val)) if abs(c0.val) > 1 else None), # x//d x0, and -> c*d 0 else (x>c.arg*d.arg) if d.arg < 0 else None), + lambda x,d,c: (x 0 else (x>c.val*d.val) if d.val < 0 else None), # ** move add/mul consts to end (NOTE: this is still happening before constant folding) ** ((UPat.var("x") + UPat.cvar("c1")) + UPat.var("y"), lambda x,c1,y: (x+y)+c1 if y.op is not Ops.CONST else None), ((UPat.var("x") * UPat.cvar("c1")) * UPat.var("y"), lambda x,c1,y: (x*y)*c1 if y.op is not Ops.CONST else None), @@ -304,12 +304,12 @@ def parse_valid(v:UOp) -> tuple[UOp, bool, int]|None: # if it's X <= c, returns X, True, c # if it's X >= c, returns X, False, c - if v.op is Ops.CMPNE and v.src[1].op is Ops.CONST and v.src[1].arg == 1 and (s0:=v.src[0]).op is Ops.CMPLT and dtypes.is_int(s0.src[0].dtype): + if v.op is Ops.CMPNE and v.src[1].op is Ops.CONST and v.src[1].val == 1 and (s0:=v.src[0]).op is Ops.CMPLT and dtypes.is_int(s0.src[0].dtype): # (X < c).ne(True) -> X >= c return s0.src[0], False, int(s0.src[1].vmin) if v.op is Ops.CMPLT and dtypes.is_int(v.src[0].dtype): # c < X -> X >= c+1 (a const on the left is a lower bound on the right) - if v.src[0].op is Ops.CONST: return v.src[1], False, int(v.src[0].arg)+1 + if v.src[0].op is Ops.CONST: return v.src[1], False, int(v.src[0].val)+1 # X < c -> X <= c-1 return v.src[0], True, int((v.src[1]).vmax)-1 return None diff --git a/tinygrad/uop/validate.py b/tinygrad/uop/validate.py index 7c50b76ea4..d87fce5a31 100644 --- a/tinygrad/uop/validate.py +++ b/tinygrad/uop/validate.py @@ -44,8 +44,8 @@ z3_renderer = PatternMatcher([ (UPat((Ops.LOAD, Ops.INDEX), dtypes.bool), lambda ctx: (z3.Bool(f"load{len(ctx[1])}", ctx=ctx[0]), None)), # constants (UPat(Ops.CONST, arg=Invalid), lambda ctx: (z3.Int("Invalid", ctx=ctx[0]), None)), - (UPat(Ops.CONST, dtypes.ints+(dtypes.weakint,), name="x"), lambda x,ctx: (z3.IntVal(x.arg, ctx=ctx[0]), None)), - (UPat(Ops.CONST, dtypes.bool, name="x"), lambda x,ctx: (z3.BoolVal(x.arg, ctx=ctx[0]), None)), + (UPat(Ops.CONST, dtypes.ints+(dtypes.weakint,), name="x"), lambda x,ctx: (z3.IntVal(x.val, ctx=ctx[0]), None)), + (UPat(Ops.CONST, dtypes.bool, name="x"), lambda x,ctx: (z3.BoolVal(x.val, ctx=ctx[0]), None)), # casts from floats create new variables (UPat(Ops.CAST, dtypes.ints+(dtypes.weakint,), src=(UPat(dtype=dtypes.floats),), name="x"), lambda x,ctx: create_bounded(f"cast{len(ctx[1])}", x.dtype.min, x.dtype.max, ctx[0])), diff --git a/tinygrad/viz/serve.py b/tinygrad/viz/serve.py index d9c32c0a31..8574bb1a71 100755 --- a/tinygrad/viz/serve.py +++ b/tinygrad/viz/serve.py @@ -129,7 +129,7 @@ def uop_to_json(data:VizData, x:UOp) -> dict[int, dict]: with soft_err(): if u.op in GroupOp.Movement and u.marg: argst = (mask_to_str if u.op in {Ops.SHRINK, Ops.PAD} else shape_to_str)(u.marg) if u.op is Ops.BINARY: argst = f"<{len(u.arg)} bytes>" - if u.op is Ops.CONST and dtypes.is_float(u.dtype): argst = f"{u.arg:g}" + if u.op is Ops.CONST and dtypes.is_float(u.dtype): argst = f"{u.val:g}" wrap_len = 200 if u.op is Ops.SOURCE else 80 label = f"{str(u.op).split('.')[1]}{(chr(10)+word_wrap(argst.replace(':', ''), wrap=wrap_len)) if u.arg is not None else ''}" if u.dtype != dtypes.void: label += f"\n{u.dtype}" @@ -138,7 +138,7 @@ def uop_to_json(data:VizData, x:UOp) -> dict[int, dict]: # walk through excluded movement ops to find the underlying CONST cx = x while cx.op in GroupOp.Movement and len(cx.src) >= 1 and cx.src[0] in excluded: cx = cx.src[0] - arg = f"{cx.arg:g}" if cx.op is Ops.CONST and dtypes.is_float(cx.dtype) else cx.render() if cx.op is Ops.STACK else f"{cx.arg}" + arg = f"{cx.val:g}" if cx.op is Ops.CONST and dtypes.is_float(cx.dtype) else cx.render() if cx.op is Ops.STACK else f"{cx.arg}" label += f"\n{cx.op.name}{idx} {arg}" + (f" {cx.src[0].op}" if len(cx.src) else "") try: if len(rngs:=u.ranges): From 161783d8f754c4719209ccfb54467580c8f03e93 Mon Sep 17 00:00:00 2001 From: qazal <77887910+Qazalin@users.noreply.github.com> Date: Sat, 1 Aug 2026 13:56:50 +0800 Subject: [PATCH 37/44] add _device_num back to ast.variables (kimi) (#17327) --- test/backend/test_multitensor.py | 1 - tinygrad/uop/ops.py | 5 +++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/test/backend/test_multitensor.py b/test/backend/test_multitensor.py index 5ee6a0cc3b..5e461236c1 100644 --- a/test/backend/test_multitensor.py +++ b/test/backend/test_multitensor.py @@ -76,7 +76,6 @@ class TestMultiTensor(unittest.TestCase): run_linear(linear) self.assertEqual(len(set(names)), 1, "function was relinearized") - @unittest.expectedFailure def test_shard_beam(self): cpu_2 = ("CPU:1", "CPU:2") src = Tensor.ones(16).shard(cpu_2, 0).realize() diff --git a/tinygrad/uop/ops.py b/tinygrad/uop/ops.py index 027508e2be..712242630e 100644 --- a/tinygrad/uop/ops.py +++ b/tinygrad/uop/ops.py @@ -993,8 +993,9 @@ class UOp(RandMixin, metaclass=UOpMetaClass): ret:dict[Variable, int] = {} return graph_rewrite(self, pm_unbind, ctx=ret), ret def variables(self) -> list[Variable]: - return sorted({x for x in self.backward_slice_with_self if x.op is Ops.PARAM and x.arg.addrspace is AddrSpace.ALU}, - key=lambda v: v.expr) + return sorted({x if x.op is Ops.PARAM else UOp.variable("_device_num", 0, x.vmax, dtype=x.dtype) + for x in self.backward_slice_with_self if (x.op is Ops.RANGE and x.arg[-1] is AxisType.DEVICE) or x.op is Ops.PARAM + and x.arg.addrspace is AddrSpace.ALU}, key=lambda v: v.expr) # *** uop symbolic stuff *** From b502fc1367c3fb996d599c4da72640e1afc5eb9d Mon Sep 17 00:00:00 2001 From: chenyu Date: Sat, 1 Aug 2026 02:18:10 -0400 Subject: [PATCH 38/44] more const arg -> val (#17350) --- extra/export_model.py | 2 +- test/amd/test_emu2_pcode.py | 34 +++++++-------- test/backend/test_isel.py | 2 +- test/backend/test_linearizer.py | 6 +-- test/mockgpu/amd/emu.py | 2 +- test/mockgpu/amd/pcode.py | 28 ++++++------ test/null/test_const_folding.py | 10 ++--- test/null/test_gpudims.py | 2 +- test/null/test_graph_rewrite.py | 64 ++++++++++++++-------------- test/null/test_memory_planner.py | 2 +- test/null/test_pattern_matcher.py | 2 +- test/null/test_simplify_valid_idx.py | 14 +++--- test/null/test_uop_graph.py | 54 +++++++++++------------ test/null/test_uop_symbolic.py | 2 +- test/null/test_viz.py | 4 +- test/unit/test_dtype_weak.py | 4 +- test/unit/test_invalid_tensor.py | 2 +- tinygrad/runtime/support/hcq2.py | 4 +- 18 files changed, 119 insertions(+), 119 deletions(-) diff --git a/extra/export_model.py b/extra/export_model.py index a9050b77e3..a0024494ba 100644 --- a/extra/export_model.py +++ b/extra/export_model.py @@ -264,7 +264,7 @@ def export_model(model, target:str, *inputs, model_name: Optional[str] = "model" if getattr(dim, "op", None) is Ops.ADD and len(dim.src) == 2 and \ any(s.op is Ops.PARAM and s.addrspace is AddrSpace.ALU for s in dim.src) and any(s.op is Ops.CONST for s in dim.src): name, val = dim.src if dim.src[1].op is Ops.CONST else reversed(dim.src) - global_size[j] = f"_{name.expr}[0] + {val.arg}" + global_size[j] = f"_{name.expr}[0] + {val.val}" prg = "" if target == "clang": diff --git a/test/amd/test_emu2_pcode.py b/test/amd/test_emu2_pcode.py index 5cdf9206e5..29dcfefdff 100644 --- a/test/amd/test_emu2_pcode.py +++ b/test/amd/test_emu2_pcode.py @@ -67,32 +67,32 @@ class TestParseExpr(unittest.TestCase): def test_integer_literals(self): """Test parsing integer literals.""" - self.assertEqual(parse_expr('0', {}).arg, 0) - self.assertEqual(parse_expr('42', {}).arg, 42) - self.assertEqual(parse_expr('42U', {}).arg, 42) + self.assertEqual(parse_expr('0', {}).val, 0) + self.assertEqual(parse_expr('42', {}).val, 42) + self.assertEqual(parse_expr('42U', {}).val, 42) def test_negative_integers(self): """Test parsing negative integer literals.""" result = parse_expr('-1', {}) - self.assertEqual(result.arg, -1) + self.assertEqual(result.val, -1) self.assertEqual(result.dtype, dtypes.int) def test_float_literals(self): """Test parsing float literals.""" result = parse_expr('1.0F', {}) - self.assertEqual(result.arg, 1.0) + self.assertEqual(result.val, 1.0) self.assertEqual(result.dtype, dtypes.float32) def test_hex_literals(self): """Test parsing hex literals.""" result = parse_expr('0xFF', {}) - self.assertEqual(result.arg, 255) + self.assertEqual(result.val, 255) def test_variable_lookup(self): """Test variable lookup in parse_expr.""" vrs = {'x': UOp.const(42, dtypes.uint32)} result = parse_expr('x', vrs) - self.assertEqual(result.arg, 42) + self.assertEqual(result.val, 42) def test_binary_ops(self): """Test parsing binary operations.""" @@ -105,7 +105,7 @@ class TestParseExpr(unittest.TestCase): # Subtraction with constant folding result = parse_expr('10 - 5', {}) self.assertEqual(result.op, Ops.CONST) - self.assertEqual(result.arg, 5) + self.assertEqual(result.val, 5) def test_ternary(self): """Test parsing ternary expressions.""" @@ -150,7 +150,7 @@ class TestForLoopParsing(unittest.TestCase): # Unwrap CAST if present while val.op == Ops.CAST: val = val.src[0] - self.assertEqual(val.arg, -1) + self.assertEqual(val.val, -1) def test_ctz_parsing(self): """Test CTZ pcode parsing.""" @@ -262,8 +262,8 @@ class TestDSPcodePatterns(unittest.TestCase): _, assigns = parse_pcode(pcode, srcs) # Check addresses: 100 + 2*4 = 108, 100 + 5*4 = 120 # assigns[i][1] is (addr, val) tuple for MEM writes; mypy sees UOp - self.assertEqual(assigns[0][1][0].simplify().arg, 108) # type: ignore[index] - self.assertEqual(assigns[1][1][0].simplify().arg, 120) # type: ignore[index] + self.assertEqual(assigns[0][1][0].simplify().val, 108) # type: ignore[index] + self.assertEqual(assigns[1][1][0].simplify().val, 120) # type: ignore[index] def test_ds_store_data_values(self): """Test DS_STORE_2ADDR_B32 uses correct data values.""" @@ -280,8 +280,8 @@ class TestDSPcodePatterns(unittest.TestCase): _, assigns = parse_pcode(pcode, srcs) # assigns[i][1] is (addr, val) tuple for MEM writes; mypy sees UOp # DATA[31:0] should preserve the value - self.assertEqual(assigns[0][1][1].simplify().arg, 0xAAAAAAAA) # type: ignore[index] - self.assertEqual(assigns[1][1][1].simplify().arg, 0xBBBBBBBB) # type: ignore[index] + self.assertEqual(assigns[0][1][1].simplify().val, 0xAAAAAAAA) # type: ignore[index] + self.assertEqual(assigns[1][1][1].simplify().val, 0xBBBBBBBB) # type: ignore[index] class TestConditionalParsing(unittest.TestCase): """Test conditional (if/elsif/else) pcode parsing.""" @@ -306,12 +306,12 @@ class TestConcatWidthParsing(unittest.TestCase): def test_permlanex16_altrow_concat(self): for row, expected in [(0, 1), (1, 0), (2, 3), (3, 2)]: parsed = parse_expr('{ row[1], ~row[0] }', {'row': UOp.const(row, dtypes.uint32)}) - self.assertEqual(parsed.simplify().arg, expected) + self.assertEqual(parsed.simplify().val, expected) def test_permlane64_altlane_concat(self): for lane, expected in [(0, 32), (1, 33), (31, 63), (32, 0), (63, 31)]: parsed = parse_expr('{ ~lane[5], lane[4:0] }', {'lane': UOp.const(lane, dtypes.uint32)}) - self.assertEqual(parsed.simplify().arg, expected) + self.assertEqual(parsed.simplify().val, expected) def test_permlane64_wave64_pcode_indices(self): vgpr = UOp.param(0, dtypes.uint32, (256,)) @@ -333,12 +333,12 @@ class TestConcatWidthParsing(unittest.TestCase): self.assertEqual(simp.src[0].op, Ops.INDEX) idx = simp.src[0].src[1].simplify() self.assertEqual(idx.op, Ops.CONST) - return idx.arg + return idx.val _, assigns = parse_pcode(PCODE[VOP1Op.V_PERMLANE64_B32_E32], srcs) self.assertEqual(len(assigns), 64) for lane, (dst_idx, src_idx) in {0: (64, 32), 31: (95, 63), 32: (96, 0), 63: (127, 31)}.items(): - self.assertEqual(assigns[lane][1][0].simplify().arg, dst_idx) # type: ignore[index] + self.assertEqual(assigns[lane][1][0].simplify().val, dst_idx) # type: ignore[index] self.assertEqual(load_idx(assigns[lane][1][1]), src_idx) # type: ignore[index] class TestAllPcode(unittest.TestCase): diff --git a/test/backend/test_isel.py b/test/backend/test_isel.py index cd357af06d..6965a5db17 100644 --- a/test/backend/test_isel.py +++ b/test/backend/test_isel.py @@ -49,7 +49,7 @@ class TestIselX86(unittest.TestCase): load = UOp.param(0, dtypes.int32, (16,)).index(a + 1).load() n = self.isel_rewrite(load) # displacement is the constant in "a" scaled to the buffer element size, dtype is int8 when the value fits otherwise int32 - self.assertTrue(n.src[2].op is Ops.CONST and n.src[2].dtype is dtypes.int8 and n.src[2].arg == 4) + self.assertTrue(n.src[2].op is Ops.CONST and n.src[2].dtype is dtypes.int8 and n.src[2].val == 4) if __name__ == "__main__": unittest.main() diff --git a/test/backend/test_linearizer.py b/test/backend/test_linearizer.py index 5dda898f0c..6d833cbeba 100644 --- a/test/backend/test_linearizer.py +++ b/test/backend/test_linearizer.py @@ -268,9 +268,9 @@ class TestLinearizer(unittest.TestCase): uops = tuple(to_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).src[1].src) idxs = dedup([uop for uop in uops if uop.op is Ops.SPECIAL]) idxs = sorted(idxs, key=lambda uop: uop.arg) - assert (idxs[0].arg, idxs[0].src[0].arg) == ('gidx0', 6), idxs[0] - assert (idxs[1].arg, idxs[1].src[0].arg) == ('gidx1', 5), idxs[1].arg - assert (idxs[2].arg, idxs[2].src[0].arg) == ('gidx2', 4), idxs[2].arg + assert (idxs[0].arg, idxs[0].src[0].val) == ('gidx0', 6), idxs[0] + assert (idxs[1].arg, idxs[1].src[0].val) == ('gidx1', 5), idxs[1].arg + assert (idxs[2].arg, idxs[2].src[0].val) == ('gidx2', 4), idxs[2].arg def test_sum_collapse(self): t = Tensor([2]).reshape(1, 1).expand(256, 256).sum() diff --git a/test/mockgpu/amd/emu.py b/test/mockgpu/amd/emu.py index c6add529c1..46925b876b 100644 --- a/test/mockgpu/amd/emu.py +++ b/test/mockgpu/amd/emu.py @@ -709,7 +709,7 @@ class _Ctx: # VGPR bit-slice assignment: VGPR[lane][reg][hi:lo] = (vgpr_idx, rhs_val, hi, lo[, cond]) -> read-modify-write if dest.startswith('VGPR[') and re.search(r'\[\d+:\d+\]', dest): # VGPR bit-slice: (vgpr_idx, rhs_val, hi_bit, lo_bit) - hi/lo are UOp constants - hi_bit, lo_bit = int(val[2].arg), int(val[3].arg) + hi_bit, lo_bit = int(val[2].val), int(val[3].val) width = hi_bit - lo_bit + 1 old = self.vgpr.index(val[0]).load() new_val = _set_bits(old, _val_to_bits(val[1]), width, lo_bit).cast(dtypes.uint32) diff --git a/test/mockgpu/amd/pcode.py b/test/mockgpu/amd/pcode.py index f7fbd8587b..bb66ffe03f 100644 --- a/test/mockgpu/amd/pcode.py +++ b/test/mockgpu/amd/pcode.py @@ -55,8 +55,8 @@ def _expr_bits(v: UOp) -> int: if v.op in (Ops.AND, Ops.XOR): widths: list[int] = [] for src in v.src: - if src.op == Ops.CONST and isinstance(src.arg, int) and src.arg > 0 and (src.arg & (src.arg + 1)) == 0: - widths.append(src.arg.bit_length()) + if src.op == Ops.CONST and isinstance(src.val, int) and src.val > 0 and (src.val & (src.val + 1)) == 0: + widths.append(src.val.bit_length()) if widths: return max(widths) return v.dtype.bitsize @@ -144,9 +144,9 @@ def _minmax_reduce(is_max: bool, dt, *args: UOp) -> UOp: def _find_two_pi_mul(x): if x.op != Ops.MUL or len(x.src) != 2: return None for i, s in enumerate(x.src): - if s.op == Ops.CONST and abs(s.arg - 6.283185307179586) < 1e-5: return (x.src[1-i], 6.283185307179586) + if s.op == Ops.CONST and abs(s.val - 6.283185307179586) < 1e-5: return (x.src[1-i], 6.283185307179586) if s.op == Ops.MUL and len(s.src) == 2: - vals = [ss.arg for ss in s.src if ss.op == Ops.CONST] + [ss.src[0].arg for ss in s.src if ss.op == Ops.CAST and ss.src[0].op == Ops.CONST] + vals = [ss.val for ss in s.src if ss.op == Ops.CONST] + [ss.src[0].val for ss in s.src if ss.op == Ops.CAST and ss.src[0].op == Ops.CONST] if len(vals) == 2 and abs(vals[0] * vals[1] - 6.283185307179586) < 1e-5: return (x.src[1-i], vals[0] * vals[1]) return None @@ -163,7 +163,7 @@ def _trig_reduce(x, phase=0.0): def _signext(val: UOp) -> UOp: for bits, mask, ext in [(4, 0xF, 0xFFFFFFF0), (8, 0xFF, 0xFFFFFF00), (16, 0xFFFF, 0xFFFF0000)]: - if (val.op == Ops.AND and len(val.src) == 2 and val.src[1].op == Ops.CONST and val.src[1].arg == mask) or val.dtype.itemsize == bits // 8: + if (val.op == Ops.AND and len(val.src) == 2 and val.src[1].op == Ops.CONST and val.src[1].val == mask) or val.dtype.itemsize == bits // 8: v32 = val.cast(dtypes.uint32) if val.dtype != dtypes.uint32 else val sb = (v32 >> _u32(bits - 1)) & _u32(1) return sb.ne(_u32(0)).where(v32 | _u32(ext), v32).cast(dtypes.int) @@ -497,7 +497,7 @@ class Parser: if not dtypes.is_int(right.dtype): right = right.cast(dtypes.uint32) return (left >> right) if op == '>>' else (left << right) case '+' | '-': - if op == '-' and left.op == Ops.CONST and right.op == Ops.CONST: return _const(left.dtype, left.arg - right.arg) + if op == '-' and left.op == Ops.CONST and right.op == Ops.CONST: return _const(left.dtype, left.val - right.val) return (left + right) if op == '+' else (left - right) case '*' | '/': # Integer promotion: promote 16-bit integers to 32-bit before multiply to avoid overflow @@ -507,7 +507,7 @@ class Parser: left, right = left.cast(pdt), right.cast(pdt) if op == '*': return left * right return (left // right) if dtypes.is_int(left.dtype) else (left / right) - case '**': return UOp(Ops.EXP2, src=(right.cast(left.dtype),)) if left.op == Ops.CONST and left.arg == 2.0 else left + case '**': return UOp(Ops.EXP2, src=(right.cast(left.dtype),)) if left.op == Ops.CONST and left.val == 2.0 else left _PREC = [('||',), ('&&',), ('|',), ('^',), ('&',), ('==', '!=', '<>'), ('>=', '<=', '>', '<'), ('>>', '<<'), ('+', '-'), ('*', '/'), ('**',)] @@ -530,7 +530,7 @@ class Parser: if self.try_eat_val('-', 'OP'): inner = self.unary() if inner.op == Ops.CONST: - return _const(dtypes.int if inner.dtype == dtypes.uint32 else inner.dtype, -inner.arg) + return _const(dtypes.int if inner.dtype == dtypes.uint32 else inner.dtype, -inner.val) return inner.neg() if self.try_eat_val('+', 'OP'): return self.unary() return self.postfix() @@ -670,14 +670,14 @@ class Parser: width = self.parse() self.eat('RBRACKET') if width.op == Ops.CONST: - w = int(width.arg) + w = int(width.val) return (base >> _to_u32(first)) & _const(base.dtype, (1 << w) - 1) return base if self.try_eat('COLON'): second = self.parse() self.eat('RBRACKET') if first.op == Ops.CONST and second.op == Ops.CONST: - a, b = int(first.arg), int(second.arg) + a, b = int(first.val), int(second.val) if a < b: return _bitreverse(base, b - a + 1) hi, lo = a, b if lo >= base.dtype.itemsize * 8: @@ -699,7 +699,7 @@ class Parser: if var_name is None: var_name = self._find_var_name(base) if first.op == Ops.CONST: - idx = int(first.arg) + idx = int(first.val) # Check for array element (var@idx) if var_name and f'{var_name}@{idx}' in self.vars: v = self.vars[f'{var_name}@{idx}'] @@ -872,7 +872,7 @@ class Parser: def _coerce_cmp(self, l: UOp, r: UOp) -> tuple[UOp, UOp]: if l.dtype != r.dtype: - if r.dtype == dtypes.int and r.op == Ops.CONST and r.arg < 0: l = l.cast(dtypes.int) + if r.dtype == dtypes.int and r.op == Ops.CONST and r.val < 0: l = l.cast(dtypes.int) else: r = r.cast(l.dtype) return l, r @@ -970,7 +970,7 @@ def parse_block(lines: list[str], start: int, env: dict[str, VarVal], funcs: dic if p.at('NUM'): return int(p.eat('NUM').val.rstrip('UuLl')) expr = p.parse().simplify() assert expr.op == Ops.CONST, f"loop bound must be constant, got {expr}" - return int(expr.arg) + return int(expr.val) start_val = parse_bound() p.eat('COLON') end_val = parse_bound() @@ -1258,7 +1258,7 @@ def parse_block(lines: list[str], start: int, env: dict[str, VarVal], funcs: dic def parse_cond(s, kw): ll = s.lower() return _to_bool(parse_expr(s[ll.find(kw) + len(kw):ll.rfind('then')].strip(), env, funcs)) - def is_const(c, v): return c.op == Ops.CONST and c.arg is v + def is_const(c, v): return c.op == Ops.CONST and c.val is v cond = parse_cond(line, 'if') conditions: list[tuple[UOp, UOp | dict[str, VarVal] | None]] = [(cond, None)] if not is_const(cond, False) else [] branch_assigns: list[tuple[UOp, list]] = [] # (cond, assigns_list) for side-effect merging diff --git a/test/null/test_const_folding.py b/test/null/test_const_folding.py index 373174528c..058aa80f6c 100644 --- a/test/null/test_const_folding.py +++ b/test/null/test_const_folding.py @@ -37,7 +37,7 @@ class TestUnaryOpsConstFolding(unittest.TestCase): class TestWeakConstFolding(unittest.TestCase): def test_weakint_math(self): out = (UOp.const(2**40) + UOp.const(2**40)).simplify() - self.assertEqual((out.op, out.dtype, out.arg), (Ops.CONST, dtypes.weakint, 2**41)) + self.assertEqual((out.op, out.dtype, out.val), (Ops.CONST, dtypes.weakint, 2**41)) def test_float_unaries(self): for op in (Ops.SIN, Ops.LOG2, Ops.EXP2, Ops.SQRT, Ops.RECIPROCAL): @@ -46,10 +46,10 @@ class TestWeakConstFolding(unittest.TestCase): def test_weakfloat_math(self): out = (UOp.const(1.25) + UOp.const(2.5)).simplify() - self.assertEqual((out.op, out.dtype, out.arg), (Ops.CONST, dtypes.weakfloat, 3.75)) + self.assertEqual((out.op, out.dtype, out.val), (Ops.CONST, dtypes.weakfloat, 3.75)) def test_invalid_poison(self): - self.assertIs(UOp.invalid().alu(Ops.CDIV, UOp.const(0)).simplify().arg, Invalid) + self.assertIs(UOp.invalid().alu(Ops.CDIV, UOp.const(0)).simplify().val, Invalid) class TestBinaryOpsConstFolding(unittest.TestCase): def test_add_literal_zero(self): @@ -123,7 +123,7 @@ class TestBitcastConstFolding(unittest.TestCase): r = full_rewrite(UOp.const(from_v, from_dt).bitcast(to_dt).sink()).src[0] self.assertEqual(r.op, Ops.CONST, msg:=f"{from_dt} -> {to_dt} ({from_v} -> {to_v})") self.assertEqual(r.dtype, to_dt, msg) - np.testing.assert_equal(r.arg, to_v, msg) + np.testing.assert_equal(r.val, to_v, msg) t({dtypes.int8: 0, dtypes.uint8: 0, dtypes.bool: False}) t({dtypes.int8: 1, dtypes.uint8: 1, dtypes.bool: True}) @@ -146,7 +146,7 @@ class TestBitcastConstFolding(unittest.TestCase): with Context(SPEC=0): srcs = full_rewrite(UOp.const((-1, -2**31, 75), dtypes.int32).bitcast(dtypes.uint32).sink()).src self.assertTrue(all(r.op is Ops.CONST and r.dtype == dtypes.uint32 for r in srcs)) - self.assertEqual(tuple(x.arg for x in srcs), (2**32-1, 2**31, 75)) + self.assertEqual(tuple(x.val for x in srcs), (2**32-1, 2**31, 75)) # folds advance indexing into basic indexing class TestIndexingConstFolding(unittest.TestCase): diff --git a/test/null/test_gpudims.py b/test/null/test_gpudims.py index 06c50e4655..8d73d907bc 100644 --- a/test/null/test_gpudims.py +++ b/test/null/test_gpudims.py @@ -12,7 +12,7 @@ class TestGroupedDims(unittest.TestCase): idxs = get_grouped_dims(prefix, dims, max_sizes, reverse) loop_idxs = dedup(flatten([[y for y in x.toposort() if y.op is Ops.SPECIAL] for x in idxs])) loop_idxs = sorted(loop_idxs, key=lambda uop: uop.arg) - sizes = [x.src[0].arg for x in loop_idxs] + sizes = [x.src[0].val for x in loop_idxs] assert len(idxs) == len(dims), f"expected idxs to have same length as dims {len(dims)}, got {len(idxs)}" if assert_same_length: assert len(loop_idxs) == min(len(sizes), len(dims)), f"expected idxs to have length {min(len(sizes), len(dims))}, got {len(loop_idxs)}" diff --git a/test/null/test_graph_rewrite.py b/test/null/test_graph_rewrite.py index 6b57882d26..e9d26d0113 100644 --- a/test/null/test_graph_rewrite.py +++ b/test/null/test_graph_rewrite.py @@ -15,13 +15,13 @@ def apply_rewrite(expr): def apply_rewrite_values(expr): srcs = full_rewrite(expr.sink()).src if len(srcs) == 1: - if srcs[0].op is Ops.CONST: return (srcs[0].arg,) if not isinstance(srcs[0].arg, tuple) else srcs[0].arg - if srcs[0].op is Ops.STACK: return tuple(s.arg for s in srcs[0].src) - return tuple(s.arg for s in srcs) + if srcs[0].op is Ops.CONST: return (srcs[0].val,) + if srcs[0].op is Ops.STACK: return tuple(s.val for s in srcs[0].src) + return tuple(s.val for s in srcs) def evaluate_uop(uop, variables): if uop.op == Ops.CONST: - return uop.arg + return uop.val elif uop.op == Ops.PARAM and uop.arg.addrspace is AddrSpace.ALU: return variables[uop.expr] elif uop.op in GroupOp.ALU: @@ -34,12 +34,12 @@ class TestArithmeticSimplifications(unittest.TestCase): def test_full_graph_rewrite_division_by_zero(self): optimized_div_uop = apply_rewrite(UOp.const(10.0) / UOp.const(0.0)) self.assertEqual(optimized_div_uop.op, Ops.CONST) - self.assertTrue(math.isinf(optimized_div_uop.arg) or math.isnan(optimized_div_uop.arg)) + self.assertTrue(math.isinf(optimized_div_uop.val) or math.isnan(optimized_div_uop.val)) def test_full_graph_rewrite_redundant_operations(self): optimized_uop = apply_rewrite((UOp.const(10.0) + UOp.const(0.0)) * UOp.const(1.0)) self.assertEqual(optimized_uop.op, Ops.CONST) - self.assertEqual(optimized_uop.arg, 10.0) + self.assertEqual(optimized_uop.val, 10.0) def test_full_graph_rewrite_large_graph(self): prev_uop = UOp.const(0) @@ -47,17 +47,17 @@ class TestArithmeticSimplifications(unittest.TestCase): prev_uop += UOp.const(i) optimized_uop = apply_rewrite(prev_uop) self.assertEqual(optimized_uop.op, Ops.CONST) - self.assertEqual(optimized_uop.arg, sum(range(1, 101))) + self.assertEqual(optimized_uop.val, sum(range(1, 101))) def test_full_graph_rewrite_division_by_one(self): optimized_uop = apply_rewrite(UOp.const(42.0) / UOp.const(1.0)) self.assertEqual(optimized_uop.op, Ops.CONST) - self.assertEqual(optimized_uop.arg, 42.0) + self.assertEqual(optimized_uop.val, 42.0) def test_full_graph_rewrite_modulo_by_one(self): optimized_uop = apply_rewrite(UOp.const(42) % UOp.const(1)) self.assertEqual(optimized_uop.op, Ops.CONST) - self.assertEqual(optimized_uop.arg, 0) + self.assertEqual(optimized_uop.val, 0) class TestFoldingAndReduction(unittest.TestCase): @@ -68,7 +68,7 @@ class TestFoldingAndReduction(unittest.TestCase): const3 = UOp.const(20) optimized_sink = apply_rewrite((const1 + const2 + const3).reduce(Ops.ADD)) expected_sum = 5 + 10 + 20 - self.assertEqual(optimized_sink.arg, expected_sum) + self.assertEqual(optimized_sink.val, expected_sum) @unittest.skip("reduce is removed now") def test_full_graph_rewrite_reduction_with_unused_range(self): @@ -77,14 +77,14 @@ class TestFoldingAndReduction(unittest.TestCase): rng = UOp.range(10, idx=0) optimized_sink = apply_rewrite((const1 + const2).reduce(Ops.ADD, rng)) expected_sum = 10 * (15 + 25) - self.assertEqual(optimized_sink.arg, expected_sum) + self.assertEqual(optimized_sink.val, expected_sum) @unittest.skip("currently failing") def test_full_graph_rewrite_range_reduction(self): simple_range = UOp.range(5, idx=0) optimized_sink = apply_rewrite(simple_range.reduce(Ops.ADD, simple_range)) expected_sum = sum(range(5)) - self.assertEqual(optimized_sink.arg, expected_sum) + self.assertEqual(optimized_sink.val, expected_sum) @unittest.skip("currently failing") def test_full_graph_rewrite_simple_reduction_folding(self): @@ -92,7 +92,7 @@ class TestFoldingAndReduction(unittest.TestCase): add_uop = simple_range + UOp.const(1) optimized_sink = apply_rewrite(add_uop.reduce(Ops.ADD, simple_range)) expected_sum = sum(i + 1 for i in range(4)) - self.assertEqual(optimized_sink.arg, expected_sum) + self.assertEqual(optimized_sink.val, expected_sum) @unittest.skip("currently failing") def test_full_graph_rewrite_nested_loop_collapse(self): @@ -101,7 +101,7 @@ class TestFoldingAndReduction(unittest.TestCase): expr = (outer_range * 10) + inner_range optimized_reduce_uop = apply_rewrite(expr.reduce(Ops.ADD, outer_range, inner_range)) self.assertEqual(optimized_reduce_uop.op, Ops.CONST) - self.assertEqual(optimized_reduce_uop.arg, sum((i * 10) + j for i in range(8) for j in range(4))) + self.assertEqual(optimized_reduce_uop.val, sum((i * 10) + j for i in range(8) for j in range(4))) class TestModuloAndDivisionFolding(unittest.TestCase): @@ -110,21 +110,21 @@ class TestModuloAndDivisionFolding(unittest.TestCase): x_var_uop = UOp.variable('x', 0, 100).cast(dtypes.weakint) optimized_mod_uop = apply_rewrite(((x_var_uop * 4) + 2) % 4) self.assertEqual(optimized_mod_uop.op, Ops.CONST) - self.assertEqual(optimized_mod_uop.arg, 2) + self.assertEqual(optimized_mod_uop.val, 2) def test_full_graph_rewrite_division_folding_with_define_var(self): # index dtype because div-mod rules only work on index n_var_uop = UOp.variable('n', 1, 1000).cast(dtypes.weakint) optimized_div_uop = apply_rewrite((n_var_uop * 6) // 3) self.assertEqual(optimized_div_uop.op, Ops.MUL) - self.assertEqual(optimized_div_uop.src[1].arg, 2) + self.assertEqual(optimized_div_uop.src[1].val, 2) def test_full_graph_rewrite_complex_mod_div_folding(self): # index dtype because div-mod rules only work on index k_var_uop = UOp.variable('k', 0, 50).cast(dtypes.weakint) optimized_div_uop = apply_rewrite(((k_var_uop * 12 + 8) % 6) // 2) self.assertEqual(optimized_div_uop.op, Ops.CONST) - self.assertEqual(optimized_div_uop.arg, 1) + self.assertEqual(optimized_div_uop.val, 1) def test_graph_rewrite_div_folding_bug(self): lhs = UOp(Ops.ADD, src=( @@ -161,9 +161,9 @@ class TestEdgeCasesAndSpecialOperations(unittest.TestCase): def test_full_graph_rewrite_transcendental_edge_cases(self): optimized_sink = full_rewrite(UOp.const(-1.0).log2().sink(UOp.const(0.0).reciprocal())) optimized_log2_neg, optimized_recip_zero = optimized_sink.src - self.assertTrue(math.isnan(optimized_log2_neg.arg), f"Expected NaN for log2(-1.0), got {optimized_log2_neg.arg}") - self.assertTrue(math.isinf(optimized_recip_zero.arg) and optimized_recip_zero.arg > 0, - f"Expected +inf for reciprocal(0.0), got {optimized_recip_zero.arg}") + self.assertTrue(math.isnan(optimized_log2_neg.val), f"Expected NaN for log2(-1.0), got {optimized_log2_neg.val}") + self.assertTrue(math.isinf(optimized_recip_zero.val) and optimized_recip_zero.val > 0, + f"Expected +inf for reciprocal(0.0), got {optimized_recip_zero.val}") @unittest.skip("broken") def test_full_graph_rewrite_modulo_negative_dividend(self): @@ -183,7 +183,7 @@ class TestGEPAndVectorizeRewrite(unittest.TestCase): def test_gep_single_element_extraction(self): # GEP on a vector dtype to extract a single element base_vector = UOp.const((1.0, 2.0, 3.0, 4.0)) - self.assertEqual(apply_rewrite(base_vector.index(2)).arg, 3.0) + self.assertEqual(apply_rewrite(base_vector.index(2)).val, 3.0) def test_gep_tuple_extraction(self): # GEP on a vector dtype to extract multiple elements as a vector @@ -193,7 +193,7 @@ class TestGEPAndVectorizeRewrite(unittest.TestCase): def test_gep_on_const_stack(self): # GEP on a const STACK to extract a single element const_stack = UOp.const((1.0, 2.0, 3.0, 4.0)) - self.assertEqual(apply_rewrite(const_stack.index(2)).arg, 3.0) + self.assertEqual(apply_rewrite(const_stack.index(2)).val, 3.0) def test_gep_tuple_on_const_stack(self): # GEP on a const STACK using a tuple to extract multiple elements @@ -322,7 +322,7 @@ class TestRecurse(unittest.TestCase): with self.assertRaises(RuntimeError): graph_rewrite(a, pm, bottom_up=True) -def bidir_append(ctx, x, b): ctx.append((x.arg if x.op is Ops.CONST else "+", b)) +def bidir_append(ctx, x, b): ctx.append((x.val if x.op is Ops.CONST else "+", b)) class TestBidirectional(unittest.TestCase): def test_simple(self): a = UOp.const(1) @@ -342,8 +342,8 @@ class TestStopEarly(unittest.TestCase): cn = UOp.const(7) d = UOp.const(2) def visit_const(c:UOp): - print(f"visit {c.arg}") - assert c.arg not in (3,4) + print(f"visit {c.val}") + assert c.val not in (3,4) pm_cvisit = PatternMatcher([(UPat(Ops.CONST, name="c"), visit_const),]) ret = (c+d).substitute({c:cn}, extra_pm=pm_cvisit) assert ret == cn+d @@ -418,7 +418,7 @@ class TestWalkRewrite(unittest.TestCase): """Top-down walk fires pm after children are processed (post-order).""" visited = [] def track_visit(ctx, x): - ctx.append(x.arg if x.op is Ops.CONST else x.op) + ctx.append(x.val if x.op is Ops.CONST else x.op) return None pm = PatternMatcher([(UPat(GroupOp.All, name="x"), track_visit)]) a = UOp.const(1) @@ -466,7 +466,7 @@ class TestWalkRewrite(unittest.TestCase): """Bottom-up walk fires bpm before descending (pre-order).""" visited = [] def track_visit(ctx, x): - ctx.append(x.arg if x.op is Ops.CONST else x.op) + ctx.append(x.val if x.op is Ops.CONST else x.op) return None pm = PatternMatcher([(UPat(GroupOp.All, name="x"), track_visit)]) a = UOp.const(1) @@ -490,10 +490,10 @@ class TestWalkRewrite(unittest.TestCase): """Bidirectional walk: bpm fires pre-order, pm fires post-order.""" visited = [] def bpm_visit(ctx, x): - ctx.append((x.arg if x.op is Ops.CONST else x.op, "bpm")) + ctx.append((x.val if x.op is Ops.CONST else x.op, "bpm")) return None def pm_visit(ctx, x): - ctx.append((x.arg if x.op is Ops.CONST else x.op, "pm")) + ctx.append((x.val if x.op is Ops.CONST else x.op, "pm")) return None bpm = PatternMatcher([(UPat(GroupOp.All, name="x"), bpm_visit)]) pm = PatternMatcher([(UPat(GroupOp.All, name="x"), pm_visit)]) @@ -509,12 +509,12 @@ class TestWalkRewrite(unittest.TestCase): """If bpm matches, children are skipped and pm never fires on that node.""" visited = [] def bpm_match(ctx, x): - ctx.append((x.arg if x.op is Ops.CONST else x.op, "bpm")) + ctx.append((x.val if x.op is Ops.CONST else x.op, "bpm")) # rewrite const(1) -> const(10), short-circuiting its subtree - if x.op is Ops.CONST and x.arg == 1: return x.replace(arg=10) + if x.op is Ops.CONST and x.val == 1: return x.replace(arg=10) return None def pm_match(ctx, x): - ctx.append((x.arg if x.op is Ops.CONST else x.op, "pm")) + ctx.append((x.val if x.op is Ops.CONST else x.op, "pm")) return None bpm = PatternMatcher([(UPat(GroupOp.All, name="x"), bpm_match)]) pm = PatternMatcher([(UPat(GroupOp.All, name="x"), pm_match)]) diff --git a/test/null/test_memory_planner.py b/test/null/test_memory_planner.py index 4361f268c7..63be203320 100644 --- a/test/null/test_memory_planner.py +++ b/test/null/test_memory_planner.py @@ -42,7 +42,7 @@ def check_assign(buffer_lists, copies=None): for orig_si, new_si in zip(linear.src, result.src): for orig, new in zip(orig_si.src[1:], new_si.src[1:]): if new.op is Ops.SLICE and id(orig) not in replace_map: - replace_map[id(orig)] = (new.src[0], new.src[1].arg * new.src[0].dtype.itemsize, new.arg * new.dtype.itemsize) + replace_map[id(orig)] = (new.src[0], new.src[1].val * new.src[0].dtype.itemsize, new.arg * new.dtype.itemsize) # verify pinned buffers are not planned for buf in held_bufs: diff --git a/test/null/test_pattern_matcher.py b/test/null/test_pattern_matcher.py index 01378bfba3..ca9b6bf0ef 100644 --- a/test/null/test_pattern_matcher.py +++ b/test/null/test_pattern_matcher.py @@ -96,7 +96,7 @@ class TestPatternMatcher(unittest.TestCase): def test_filter_arg(self): matcher = PatternMatcher([ (UPat(Ops.MUL, src=[UPat(Ops.CONST, name="c"), UPat(Ops.CONST, arg=2)], name="x"), - lambda x,c: x.rtag() if c.arg in {1, -1} else None) + lambda x,c: x.rtag() if c.val in {1, -1} else None) ]) y1 = UOp.const(1) y2 = UOp.const(2) diff --git a/test/null/test_simplify_valid_idx.py b/test/null/test_simplify_valid_idx.py index 20f378000d..15a9ce9171 100644 --- a/test/null/test_simplify_valid_idx.py +++ b/test/null/test_simplify_valid_idx.py @@ -529,7 +529,7 @@ class TestRangeShrink(unittest.TestCase): load = get_gated_load_uop(r < UOp.const(4), r) ranges = self.get_ranges(load.sink()) self.assertEqual(len(ranges), 1) - self.assertEqual(ranges[0].src[0].arg, 4) + self.assertEqual(ranges[0].src[0].val, 4) def test_range_shrink_picks_max_guard(self): # two loads guard the same range with r < 4 and r < 8 -> shrink to max(4, 8) = 8 @@ -538,7 +538,7 @@ class TestRangeShrink(unittest.TestCase): load2 = get_gated_load_uop(r < UOp.const(8), r) ranges = self.get_ranges(UOp.sink(load1, load2)) self.assertEqual(len(ranges), 1) - self.assertEqual(ranges[0].src[0].arg, 8) + self.assertEqual(ranges[0].src[0].val, 8) def test_range_no_shrink_guard_ge_max(self): # guard r < 300 with range max 204 -> no shrink (guard doesn't constrain) @@ -546,7 +546,7 @@ class TestRangeShrink(unittest.TestCase): load = get_gated_load_uop(r < UOp.const(300), r) ranges = self.get_ranges(load.sink()) self.assertEqual(len(ranges), 1) - self.assertEqual(ranges[0].src[0].arg, 204) + self.assertEqual(ranges[0].src[0].val, 204) def test_range_no_shrink_when_unguarded_elsewhere(self): # one load guards r < 4, but another load uses r without a gate -> no shrink @@ -555,7 +555,7 @@ class TestRangeShrink(unittest.TestCase): load2 = UOp(Ops.LOAD, src=(UOp.param(1, dtypes.float, (204,)).index(r),)) ranges = self.get_ranges(UOp.sink(load1, load2)) self.assertEqual(len(ranges), 1) - self.assertEqual(ranges[0].src[0].arg, 204) + self.assertEqual(ranges[0].src[0].val, 204) def test_range_no_shrink_when_used_in_reduce(self): # range used in both a gated load AND directly in the reduce expression -> no shrink @@ -564,7 +564,7 @@ class TestRangeShrink(unittest.TestCase): red = (r.cast(dtypes.float) + gated_load).reduce(r, arg=Ops.ADD) ranges = self.get_ranges(red.sink()) self.assertEqual(len(ranges), 1) - self.assertEqual(ranges[0].src[0].arg, 204) + self.assertEqual(ranges[0].src[0].val, 204) def test_range_shrink_to_single_iteration(self): # guard r < 1 shrinks range to 1 -> single iteration, range eliminated entirely @@ -580,7 +580,7 @@ class TestRangeShrink(unittest.TestCase): x = (r < 4).where(UOp.const(1.0), Invalid) ranges = self.get_ranges(UOp.param(0, dtypes.float, (204,)).index(r).store((r < 4).where(x, Invalid)).sink()) self.assertEqual(len(ranges), 1) - self.assertEqual(ranges[0].src[0].arg, 4) + self.assertEqual(ranges[0].src[0].val, 4) def test_range_shrink_store_where_invalid_flipped(self): # above, but flipped @@ -589,7 +589,7 @@ class TestRangeShrink(unittest.TestCase): x = (r < 4).where(UOp.const(1.0), Invalid) ranges = self.get_ranges(UOp.param(0, dtypes.float, (204,)).index(r).store((r >= 4).where(Invalid, x)).sink()) self.assertEqual(len(ranges), 1) - self.assertEqual(ranges[0].src[0].arg, 4) + self.assertEqual(ranges[0].src[0].val, 4) if __name__ == '__main__': unittest.main() diff --git a/test/null/test_uop_graph.py b/test/null/test_uop_graph.py index 48471de469..e4ce98c0db 100644 --- a/test/null/test_uop_graph.py +++ b/test/null/test_uop_graph.py @@ -8,14 +8,14 @@ from test.helpers import to_uops_list simple_pm = PatternMatcher([ (UPat.cvar('x', dtypes.weakint), lambda x: UOp.const(1.0) + UOp.const(2.0)), - (UPat.cvar('x') + UPat.cvar('y'), lambda x,y: UOp.const(x.arg+y.arg)), - (UPat.cvar('x') * UPat.cvar('y') * UPat.cvar('z'), lambda x,y,z: UOp.const(x.arg*y.arg*z.arg)), - ((UPat.var('x') + UPat.cvar('c1')) + UPat.cvar('c2'), lambda x,c1,c2: x + (c1.arg+c2.arg)), + (UPat.cvar('x') + UPat.cvar('y'), lambda x,y: UOp.const(x.val+y.val)), + (UPat.cvar('x') * UPat.cvar('y') * UPat.cvar('z'), lambda x,y,z: UOp.const(x.val*y.val*z.val)), + ((UPat.var('x') + UPat.cvar('c1')) + UPat.cvar('c2'), lambda x,c1,c2: x + (c1.val+c2.val)), ]) def const_values(u:UOp): - if u.op is Ops.CONST: return (u.arg,) if not isinstance(u.arg, tuple) else u.arg - if u.op is Ops.STACK: return tuple(x.arg for x in u.src) + if u.op is Ops.CONST: return (u.val,) + if u.op is Ops.STACK: return tuple(x.val for x in u.src) raise AssertionError(f"expected const-like UOp, got {u.op}") class TestGraphRewriteConst(unittest.TestCase): @@ -24,7 +24,7 @@ class TestGraphRewriteConst(unittest.TestCase): v2 = v1.index(1) ret = graph_rewrite(v2, sym) self.assertEqual(ret.dtype, dtypes.int) - self.assertEqual(ret.arg, 1) + self.assertEqual(ret.val, 1) def test_add_const(self): v1 = UOp.const((0,1,2)) @@ -49,7 +49,7 @@ class TestModularWraparound(unittest.TestCase): self.assertEqual(len(results), 2) # +1 for SINK self.assertEqual(results[0].op, Ops.CONST) self.assertEqual(results[0].dtype, uop.dtype) - self.assertEqual(results[0].arg, expected) + self.assertEqual(results[0].val, expected) @xfail_broken_const_wraparound def test_cast(self): @@ -114,7 +114,7 @@ class TestGraphRewrite(unittest.TestCase): c2 = UOp.const(2.0) nout = graph_rewrite(c1+c2, simple_pm) self.assertEqual(nout.op, Ops.CONST) - self.assertEqual(nout.arg, 3.0) + self.assertEqual(nout.val, 3.0) def test_depth_2_late(self): c1 = UOp.const(1.0) @@ -122,7 +122,7 @@ class TestGraphRewrite(unittest.TestCase): c3 = UOp.const(3.0) nout = graph_rewrite(c1*c2*(c3+c3), simple_pm) self.assertEqual(nout.op, Ops.CONST) - self.assertEqual(nout.arg, 12.0) + self.assertEqual(nout.val, 12.0) def test_double(self): c1 = UOp.const(1.0) @@ -130,7 +130,7 @@ class TestGraphRewrite(unittest.TestCase): c3 = UOp.const(3.0) nout = graph_rewrite(c1+c2+c3, simple_pm) self.assertEqual(nout.op, Ops.CONST) - self.assertEqual(nout.arg, 6.0) + self.assertEqual(nout.val, 6.0) def test_triple(self): c1 = UOp.const(1.0) @@ -139,7 +139,7 @@ class TestGraphRewrite(unittest.TestCase): c4 = UOp.const(4.0) nout = graph_rewrite(c1+c2+c3+c4, simple_pm) self.assertEqual(nout.op, Ops.CONST) - self.assertEqual(nout.arg, 10.0) + self.assertEqual(nout.val, 10.0) def test_diamond(self): c1 = UOp.const(1.0) @@ -147,13 +147,13 @@ class TestGraphRewrite(unittest.TestCase): c3 = UOp.const(3.0) nout = graph_rewrite((c1+c2)+(c1+c3), simple_pm) self.assertEqual(nout.op, Ops.CONST) - self.assertEqual(nout.arg, 7.0) + self.assertEqual(nout.val, 7.0) def test_magic_4(self): c1 = UOp.const(4) nout = graph_rewrite(c1, simple_pm) self.assertEqual(nout.op, Ops.CONST) - self.assertEqual(nout.arg, 3.0) + self.assertEqual(nout.val, 3.0) def test_depth_2_fold(self): v = UOp.variable("v", 0, 1, dtypes.float) @@ -163,7 +163,7 @@ class TestGraphRewrite(unittest.TestCase): self.assertEqual(nout.op, Ops.ADD) self.assertEqual(nout.src[0].op, Ops.PARAM) self.assertEqual(nout.src[1].op, Ops.CONST) - self.assertEqual(nout.src[1].arg, 3.0) + self.assertEqual(nout.src[1].val, 3.0) def test_commutative_work(self): a = UOp.variable('a', 0, 1) @@ -198,7 +198,7 @@ class TestUOpGraph(unittest.TestCase): self.assertEqual(len(uops), 2) # +1 for SINK out = uops[-2] self.assertEqual(out.op, Ops.CONST) - self.assertEqual(out.arg, 3.0) + self.assertEqual(out.val, 3.0) def test_where_same_fold(self): v = UOp.variable('tmp', 0, 1) @@ -210,7 +210,7 @@ class TestUOpGraph(unittest.TestCase): self.assertEqual(len(uops), 2) # +1 for SINK out = uops[-2] self.assertEqual(out.op, Ops.CONST) - self.assertEqual(out.arg, 1.0) + self.assertEqual(out.val, 1.0) def test_where_const_fold(self): bf = UOp.const(False) @@ -221,7 +221,7 @@ class TestUOpGraph(unittest.TestCase): self.assertEqual(len(uops), 2) # +1 for SINK out = uops[-2] self.assertEqual(out.op, Ops.CONST) - self.assertEqual(out.arg, 2.0) + self.assertEqual(out.val, 2.0) def test_const_cast(self): bf = UOp.const(False) @@ -230,7 +230,7 @@ class TestUOpGraph(unittest.TestCase): self.assertEqual(len(uops), 2) # +1 for SINK out = uops[-2] self.assertEqual(out.op, Ops.CONST) - self.assertEqual(out.arg, 0) + self.assertEqual(out.val, 0) def test_const_bitcast(self): bf = UOp.const(1.0, dtypes.float) @@ -239,7 +239,7 @@ class TestUOpGraph(unittest.TestCase): self.assertEqual(len(uops), 2) # +1 for SINK out = uops[-2] self.assertEqual(out.op, Ops.CONST) - self.assertEqual(out.arg, 0x3F800000) + self.assertEqual(out.val, 0x3F800000) @unittest.expectedFailure def test_const_shape_change_bitcast(self): @@ -348,7 +348,7 @@ class TestUOpGraph(unittest.TestCase): out = uops[-2] # -2 to skip SINK self.assertEqual(out.op, Ops.ADD) self.assertEqual(out.src[1].op, Ops.CONST) - self.assertEqual(out.src[1].arg, 6) + self.assertEqual(out.src[1].val, 6) def test_bitcast_to_same_dtype_fold(self): for dt in dtypes.ints + dtypes.floats + (dtypes.bool,): @@ -372,7 +372,7 @@ class TestUOpGraph(unittest.TestCase): uops = to_uops_list([out.index(ridx0).store(w)]) for u in uops: assert u.op is not Ops.WHERE - if u.op is Ops.LOAD and u.src[0].src[0].op is Ops.PARAM: assert u.src[1].arg==5 + if u.op is Ops.LOAD and u.src[0].src[0].op is Ops.PARAM: assert u.src[1].val==5 def test_where_on_gated_load_folds_swapped_branches(self): ridx0 = UOp.range(100, 0) @@ -382,7 +382,7 @@ class TestUOpGraph(unittest.TestCase): uops = to_uops_list([w]) for u in uops: assert u.op is not Ops.WHERE - if u.op is Ops.LOAD: assert u.src[1].arg==5 + if u.op is Ops.LOAD: assert u.src[1].val==5 def test_where_on_gated_load_with_cast(self): ridx0 = UOp.range(100, 0) @@ -394,7 +394,7 @@ class TestUOpGraph(unittest.TestCase): uops = to_uops_list([out.index(ridx0).store(w)]) for u in uops: assert u.op is not Ops.WHERE - if u.op is Ops.LOAD and u.src[0].src[0].op is Ops.PARAM: assert u.src[1].arg == 5 + if u.op is Ops.LOAD and u.src[0].src[0].op is Ops.PARAM: assert u.src[1].val == 5 def test_where_on_casted_gated_load_extra_cond(self): ridx0 = UOp.range(100, 0) @@ -426,7 +426,7 @@ class TestUOpGraph(unittest.TestCase): uops = to_uops_list([st]) for u in uops: assert u.op is not Ops.WHERE - if u.op is Ops.STORE: assert u.src[1].arg==5 + if u.op is Ops.STORE: assert u.src[1].val==5 def test_load_idx_becomes_int(self): # mnist indexing with split reduceop @@ -571,7 +571,7 @@ class TestConstBufferize(unittest.TestCase): result = graph_rewrite(bufferize_with_range, pm_const_buffer_folding, name='test') # BUFFERIZE should be removed, result is const broadcast to shape self.assertNotEqual(result.op, Ops.STAGE) - const_vals = [u.arg for u in result.toposort() if u.op is Ops.CONST and u.dtype is dtypes.weakfloat] + const_vals = [u.val for u in result.toposort() if u.op is Ops.CONST and u.dtype is dtypes.weakfloat] self.assertIn(42.0, const_vals) def test_const_bufferize_with_multiple_ranges(self): @@ -586,14 +586,14 @@ class TestConstBufferize(unittest.TestCase): result = graph_rewrite(bufferize_with_ranges, pm_const_buffer_folding, name='test') # BUFFERIZE should be removed self.assertNotEqual(result.op, Ops.STAGE) - const_vals = [u.arg for u in result.toposort() if u.op is Ops.CONST and u.dtype is dtypes.weakfloat] + const_vals = [u.val for u in result.toposort() if u.op is Ops.CONST and u.dtype is dtypes.weakfloat] self.assertIn(3.14, const_vals) class TestUOpTags(unittest.TestCase): def test_inc_by_one(self): g = UOp.const(1) + UOp.const(1) assert g.ssimplify() == 2 - pm_plus_1 = PatternMatcher([(UPat(Ops.CONST, name="x"), lambda x: x.replace(arg=x.arg+1, tag=1) if x.tag is None else None)]) + pm_plus_1 = PatternMatcher([(UPat(Ops.CONST, name="x"), lambda x: x.replace(arg=x.val+1, tag=1) if x.tag is None else None)]) pm_strip_tags = PatternMatcher([(UPat(GroupOp.All, name="x"), lambda x: x.replace(tag=None) if x.tag is not None else None)]) g = graph_rewrite(g, pm_plus_1) assert g.ssimplify() == 4 diff --git a/test/null/test_uop_symbolic.py b/test/null/test_uop_symbolic.py index 183175ef09..36629e025a 100644 --- a/test/null/test_uop_symbolic.py +++ b/test/null/test_uop_symbolic.py @@ -1178,7 +1178,7 @@ class TestSymbolicVariables(unittest.TestCase): b = Variable("x", 1, 1).bind(1) s = b.simplify() self.assertEqual(s.op, Ops.CONST) - self.assertEqual(s.arg, 1) + self.assertEqual(s.val, 1) class TestSymInfer(unittest.TestCase): def test_sym_infer(self): diff --git a/test/null/test_viz.py b/test/null/test_viz.py index 2fa147d94a..fad98fc4a4 100644 --- a/test/null/test_viz.py +++ b/test/null/test_viz.py @@ -96,8 +96,8 @@ class TestViz(unittest.TestCase): def test_exceptions(self): # VIZ tracks rewrites up to and including the error def count_3(x:UOp): - assert x.arg <= 3 - return x.replace(arg=x.arg+1) + assert x.val <= 3 + return x.replace(arg=x.val+1) err_pm = PatternMatcher([(UPat.cvar("x"), count_3),]) a = UOp.const(1) with save_viz() as viz: diff --git a/test/unit/test_dtype_weak.py b/test/unit/test_dtype_weak.py index 5ae7cca6b0..30ce55fe65 100644 --- a/test/unit/test_dtype_weak.py +++ b/test/unit/test_dtype_weak.py @@ -55,7 +55,7 @@ class TestWeakPromotion(unittest.TestCase): x, y = Tensor([1], dtype=dtypes.int8)._broadcasted(0.5) self.assertEqual((y._uop.base.op, y.dtype, x.dtype), (Ops.CONST, dtypes.weakfloat, dtypes.weakfloat)) x, y = Tensor.const(1).reshape(1)._broadcasted(Tensor([1.0], dtype=dtypes.float32)) - self.assertEqual((x._uop.base.op, x._uop.base.arg, x.dtype, x.shape, y.dtype), + self.assertEqual((x._uop.base.op, x._uop.base.val, x.dtype, x.shape, y.dtype), (Ops.CONST, 1, dtypes.weakfloat, (1,), dtypes.float32)) def test_uop_scalar_const_lifts_kind(self): @@ -68,7 +68,7 @@ class TestWeakPromotion(unittest.TestCase): # the kind lift converts the VALUE too (the arg is the only dtype carrier once UOp.const loses its dtype arg), # and a bare weak const UOp is the same spelling as the python scalar: both lift to the same node x = UOp.variable("x", 0.0, 1.0, dtypes.float32) - self.assertIsInstance((x + 2).src[1].arg, float) + self.assertIsInstance((x + 2).src[1].val, float) self.assertIs(x + UOp.const(2), x + 2) def test_index_dtype_ignores_weakness(self): diff --git a/test/unit/test_invalid_tensor.py b/test/unit/test_invalid_tensor.py index a23c0df55a..a3f0647efa 100644 --- a/test/unit/test_invalid_tensor.py +++ b/test/unit/test_invalid_tensor.py @@ -141,7 +141,7 @@ class TestInvalidTensor(unittest.TestCase): self.assertIs(idx.op, Ops.STACK) self.assertIs(out.op, Ops.WHERE) self.assertIs(out.src[2].op, Ops.CONST) - self.assertIs(out.src[2].arg, Invalid) + self.assertIs(out.src[2].val, Invalid) if __name__ == '__main__': unittest.main() diff --git a/tinygrad/runtime/support/hcq2.py b/tinygrad/runtime/support/hcq2.py index 0f36f516de..61a8d3ebd3 100644 --- a/tinygrad/runtime/support/hcq2.py +++ b/tinygrad/runtime/support/hcq2.py @@ -419,8 +419,8 @@ def fold_binary(buf:UOp, blob:UOp) -> UOp: def fold_const_store(buf:UOp, off:UOp, val:UOp) -> UOp: for off,val in zip(off.src, val.src): for b,v in zip((bs:=mb.bufs if isinstance((mb:=buf.buffer), MultiBuffer) else (mb,)), val.src if val.op is Ops.STACK else (val,)*len(bs)): - data = struct.pack(f'<{v.dtype.fmt}', truncate[v.dtype](v.arg)) - b.ensure_allocated().as_memoryview(force_zero_copy=True, no_sync=True).cast('B')[(bo:=off.arg*buf.dtype.itemsize):bo+len(data)] = data + data = struct.pack(f'<{v.dtype.fmt}', truncate[v.dtype](v.val)) + b.ensure_allocated().as_memoryview(force_zero_copy=True, no_sync=True).cast('B')[(bo:=off.val*buf.dtype.itemsize):bo+len(data)] = data return UOp(Ops.NOOP) def resolve_getaddr(buf:UOp, g:UOp) -> UOp: From 6c0ec39279ba2f0bac964c6204911226517effc4 Mon Sep 17 00:00:00 2001 From: chenyu Date: Sat, 1 Aug 2026 02:36:32 -0400 Subject: [PATCH 39/44] UOp.is_invalid [PR] (#17351) helper to prep ConstArg --- test/null/test_const_folding.py | 4 ++-- test/null/test_uops.py | 2 +- test/unit/test_invalid_tensor.py | 2 +- tinygrad/codegen/__init__.py | 6 +++--- tinygrad/codegen/late/gater.py | 2 +- tinygrad/function.py | 4 +--- tinygrad/mixin/dtype.py | 4 ++-- tinygrad/schedule/rangeify.py | 4 ++-- tinygrad/uop/ops.py | 10 ++++++---- tinygrad/uop/spec.py | 6 +++--- tinygrad/uop/symbolic.py | 12 ++++++------ 11 files changed, 28 insertions(+), 28 deletions(-) diff --git a/test/null/test_const_folding.py b/test/null/test_const_folding.py index 058aa80f6c..0f217e8fa5 100644 --- a/test/null/test_const_folding.py +++ b/test/null/test_const_folding.py @@ -1,6 +1,6 @@ import unittest, itertools, math from tinygrad import Tensor, dtypes, Context -from tinygrad.dtype import DType, ConstType, Invalid +from tinygrad.dtype import DType, ConstType from tinygrad.uop.ops import Ops, UOp from test.helpers import full_rewrite import numpy as np @@ -49,7 +49,7 @@ class TestWeakConstFolding(unittest.TestCase): self.assertEqual((out.op, out.dtype, out.val), (Ops.CONST, dtypes.weakfloat, 3.75)) def test_invalid_poison(self): - self.assertIs(UOp.invalid().alu(Ops.CDIV, UOp.const(0)).simplify().val, Invalid) + self.assertTrue(UOp.invalid().alu(Ops.CDIV, UOp.const(0)).simplify().is_invalid) class TestBinaryOpsConstFolding(unittest.TestCase): def test_add_literal_zero(self): diff --git a/test/null/test_uops.py b/test/null/test_uops.py index 95362b767d..ea701a30ca 100644 --- a/test/null/test_uops.py +++ b/test/null/test_uops.py @@ -52,7 +52,7 @@ class TestDTypeFromUOp(unittest.TestCase): self.assertIs((moved:=invalid.reshape((1,))).cast(dtypes.float32), moved) scratch = Tensor.invalids(4, dtype=dtypes.float32) self.assertEqual((scratch.dtype, next(u.dtype for u in scratch.uop.toposort() if u.op is Ops.BUFFER), next(u.dtype for u in scratch.uop.toposort() - if u.arg is Invalid)), (dtypes.float32, dtypes.float32, dtypes.bool)) + if u.is_invalid)), (dtypes.float32, dtypes.float32, dtypes.bool)) invalid, value = UOp.invalid(), UOp.const(1, dtypes.float32) for u in (UOp(Ops.STACK, dtypes.float32, src=(value, invalid)), UOp(Ops.ADD, dtypes.float32, src=(value, invalid)), UOp.const(True).where(value, invalid), UOp(Ops.CMPLT, src=(invalid, value)), UOp(Ops.CMPLT, src=(value, invalid)), diff --git a/test/unit/test_invalid_tensor.py b/test/unit/test_invalid_tensor.py index a3f0647efa..1c8630af36 100644 --- a/test/unit/test_invalid_tensor.py +++ b/test/unit/test_invalid_tensor.py @@ -141,7 +141,7 @@ class TestInvalidTensor(unittest.TestCase): self.assertIs(idx.op, Ops.STACK) self.assertIs(out.op, Ops.WHERE) self.assertIs(out.src[2].op, Ops.CONST) - self.assertIs(out.src[2].val, Invalid) + self.assertTrue(out.src[2].is_invalid) if __name__ == '__main__': unittest.main() diff --git a/tinygrad/codegen/__init__.py b/tinygrad/codegen/__init__.py index fffce829fa..be9dd54220 100644 --- a/tinygrad/codegen/__init__.py +++ b/tinygrad/codegen/__init__.py @@ -8,7 +8,7 @@ from tinygrad.uop.render import pyrender from tinygrad.uop.spec import type_verify, spec_tensor, spec_program from tinygrad.renderer import Renderer, Estimates from tinygrad.renderer.isa import ISARenderer, IselContext, PreRegAllocContext -from tinygrad.dtype import dtypes, AddrSpace, Invalid +from tinygrad.dtype import dtypes, AddrSpace # import all pattern matchers here from tinygrad.codegen.gpudims import pm_add_gpudims @@ -123,11 +123,11 @@ pm_expand_broadcast = pm_wmma_add+PatternMatcher([ def do_devectorize(b:UOp): if b.shape == (): return None # broadcasting needs to be already unpacked, Invalid matches any dtype and shape - if not all(x.shape == b.shape or x.base.arg is Invalid for x in b.src): return None + if not all(x.shape == b.shape or x.base.is_invalid for x in b.src): return None src = [] for idx in itertools.product(*[range(x) for x in b.shape]): idx_c = [UOp.const(i) for i in idx] - src.append(b.replace(dtype=None, src=tuple(x.base if x.base.arg is Invalid else x.index(*idx_c) for x in b.src))) + src.append(b.replace(dtype=None, src=tuple(x.base if x.base.is_invalid else x.index(*idx_c) for x in b.src))) return UOp.stack(*src).reshape(b.shape) if b.op is not Ops.STORE else UOp.group(*src) def do_stack_wmma(u:UOp): diff --git a/tinygrad/codegen/late/gater.py b/tinygrad/codegen/late/gater.py index ca70eeb6c1..b9b3a4efb7 100644 --- a/tinygrad/codegen/late/gater.py +++ b/tinygrad/codegen/late/gater.py @@ -3,7 +3,7 @@ from tinygrad.uop.ops import PatternMatcher, UPat, Ops from tinygrad.dtype import Invalid, dtypes def move_where_load(gate, l, a, w): - return l.replace(src=(l.src[0], l.vconst_like(0) if a.arg is Invalid else + return l.replace(src=(l.src[0], l.vconst_like(0) if a.is_invalid else a.src[0] if a.op is Ops.CAST and a.src[0].dtype == l.dtype else a.cast(l.dtype), l.src[2])).cast(w.dtype) pm_move_gates_from_index = PatternMatcher([ diff --git a/tinygrad/function.py b/tinygrad/function.py index 8b540e3381..edb1691d9d 100644 --- a/tinygrad/function.py +++ b/tinygrad/function.py @@ -1,7 +1,6 @@ import functools, time from typing import Generic, TypeVar, Callable, cast, overload from tinygrad.helpers import Context, dedup, getenv, DEBUG -from tinygrad.dtype import Invalid from tinygrad.uop.ops import UOp, Ops, graph_rewrite, PatternMatcher, UPat from tinygrad.tensor import Tensor from tinygrad.nn.state import get_state_dict @@ -22,8 +21,7 @@ def invalid_outputs(uret:UOp) -> set[UOp]: # invalids() returns fresh write-only scratch: a clone storing CONST(Invalid) # don't capture it as an input; only skip fresh buffers, not realized ones return {u.src[0].buf_uop for u in uret.backward_slice_with_self - if u.op is Ops.STORE and u.src[1].base.op is Ops.CONST and u.src[1].base.val is Invalid - and not u.src[0].buf_uop.is_realized} + if u.op is Ops.STORE and u.src[1].base.is_invalid and not u.src[0].buf_uop.is_realized} ReturnType = TypeVar('ReturnType') class _function(Generic[ReturnType]): diff --git a/tinygrad/mixin/dtype.py b/tinygrad/mixin/dtype.py index a96e9c7028..08d112da74 100644 --- a/tinygrad/mixin/dtype.py +++ b/tinygrad/mixin/dtype.py @@ -1,5 +1,5 @@ from typing import TYPE_CHECKING, Self -from tinygrad.dtype import DType, DTypeLike, dtypes, to_dtype, Invalid +from tinygrad.dtype import DType, DTypeLike, dtypes, to_dtype from tinygrad.uop import Ops if TYPE_CHECKING: @@ -30,7 +30,7 @@ class DTypeMixin: print(t.dtype, t.numpy()) ``` """ - return self if self.dtype == (dt:=to_dtype(dtype)) or self._uop.base.arg is Invalid else self._wrap_uop(self._uop.alu(Ops.CAST, arg=dt)) + return self if self.dtype == (dt:=to_dtype(dtype)) or self._uop.base.is_invalid else self._wrap_uop(self._uop.alu(Ops.CAST, arg=dt)) def bitcast(self, dtype:DTypeLike) -> Self: """ diff --git a/tinygrad/schedule/rangeify.py b/tinygrad/schedule/rangeify.py index 0e6b44b376..90ca0938ae 100644 --- a/tinygrad/schedule/rangeify.py +++ b/tinygrad/schedule/rangeify.py @@ -23,7 +23,7 @@ def found_after(ctx:dict[UOp, UOp], after:UOp, src:UOp): while True: if x.op is Ops.PERMUTE: x, after = x.src[0], after.permute(argsort(x.marg)) elif x.op is Ops.RESHAPE: x, after = x.src[0], after.reshape(x.src[0].shape) - elif x.op is Ops.WHERE and x.src[2].base.arg == Invalid and x.src[1].op is Ops.PAD: + elif x.op is Ops.WHERE and x.src[2].base.is_invalid and x.src[1].op is Ops.PAD: x, after = x.src[1].src[0], after.shrink(tuple((o, s+o) for (o,_),s in zip(x.src[1].marg, x.src[1].src[0].shape))) else: break ctx[x] = after @@ -293,7 +293,7 @@ def remove_noop_bufferize(idx,b2): def after_all_invalid(after:UOp): buf = after.src[0].buf_uop # check all ranges are used (no expand), and same size (no pad and shrink) - return all(s.op is Ops.END and (st:=s.src[0]).op is Ops.STORE and st.src[1].base.arg is Invalid and st.src[0].buf_uop is buf + return all(s.op is Ops.END and (st:=s.src[0]).op is Ops.STORE and st.src[1].base.is_invalid and st.src[0].buf_uop is buf and all(r in st.src[0].ranges for r in s.ended_ranges) and resolve(cast(UOp, prod(r.src[0] for r in s.ended_ranges)).eq(buf.numel()), False) for s in after.src[1:]) diff --git a/tinygrad/uop/ops.py b/tinygrad/uop/ops.py index 712242630e..8341eca529 100644 --- a/tinygrad/uop/ops.py +++ b/tinygrad/uop/ops.py @@ -195,7 +195,7 @@ class UOpMetaClass(type): # TODO: delete this once the dtype field is removed, for now it just re-implements spec.py # an INDEX presents its access dtype, which a still-weak source matches up to weakness if SPEC == 2 and op is not Ops.CONST and \ - not any(s.base.arg is Invalid for s in src) and (expected_dtype:=dtype_from_uop(op, src, arg)) is not None and expected_dtype != dtype and \ + not any(s.base.is_invalid for s in src) and (expected_dtype:=dtype_from_uop(op, src, arg)) is not None and expected_dtype != dtype and \ not (op is Ops.INDEX and weak_dtype(expected_dtype) == weak_dtype(dtype)): raise RuntimeError(f"bad dtype {dtype}, expected {expected_dtype} on {op}") if (wret:=UOpMetaClass.ucache.get(key:=(op, dtype, src, arg, tag), None)) is not None and (ret:=wret()) is not None: return ret @@ -263,6 +263,8 @@ class UOp(RandMixin, metaclass=UOpMetaClass): def val(self): assert self.op is Ops.CONST, f"val is only for CONST, got {self.op}" return self.arg + @property + def is_invalid(self) -> bool: return self.op is Ops.CONST and self.val is Invalid @recursive_property def key(self) -> bytes: return hashlib.sha256(str((self.op, self.dtype, self.arg)).encode() + b"".join([s.key for s in self.src])).digest() @@ -648,10 +650,10 @@ class UOp(RandMixin, metaclass=UOpMetaClass): return cond.where(self, self.const_like(Invalid)) def get_idx(self) -> UOp: if self.op is Ops.STACK: return UOp.stack(*(x.get_idx() for x in self.src)) - return self.src[1] if self.op is Ops.WHERE and self.src[2].arg is Invalid else self + return self.src[1] if self.op is Ops.WHERE and self.src[2].is_invalid else self def get_valid(self) -> UOp: if self.op is Ops.STACK: return UOp.stack(*(x.get_valid() for x in self.src)) - return self.src[0] if self.op is Ops.WHERE and self.src[2].arg is Invalid else UOp.const(self.arg is not Invalid) + return self.src[0] if self.op is Ops.WHERE and self.src[2].is_invalid else UOp.const(not self.is_invalid) def reduce(self, *src:UOp, **kwargs): arg = kwargs.pop('arg', None) if isinstance(arg, Ops): arg = (arg, 0) @@ -1747,7 +1749,7 @@ def _rebuild_dtype(n:UOp, new_src:tuple[UOp,...]) -> DType: # TODO: delete this once the dtype field is removed, every rebuild will re-derive # TODO: these ops keep their stored dtype until dtype_from_uop works if n.op in {Ops.INDEX, Ops.CUSTOM, Ops.CUSTOMI, Ops.PYLITERAL} or \ - all(a.dtype is b.dtype or b.base.arg is Invalid for a,b in zip(n.src, new_src)): return n.dtype + all(a.dtype is b.dtype or b.base.is_invalid for a,b in zip(n.src, new_src)): return n.dtype return dtype_from_uop(n.op, new_src, n.arg) or n.dtype def sint_to_uop(x:sint, dtype=dtypes.weakint) -> UOp: return UOp.const(x, dtype) diff --git a/tinygrad/uop/spec.py b/tinygrad/uop/spec.py index 4f07447a60..69b13657d0 100644 --- a/tinygrad/uop/spec.py +++ b/tinygrad/uop/spec.py @@ -10,7 +10,7 @@ from tinygrad.helpers import DEBUG, Context, SPEC, Metadata, panic, CHECK_OOB, a def validate_index(uidx:UOp, gate:UOp|None=None): if len(uidx.src) != 2: return True # skip for non final index. TODO: check more complex index with shape buf,idx = uidx.src - if idx.op is Ops.CONST and idx.val is Invalid: return True + if idx.is_invalid: return True if gate is None: gate = UOp.const(True) # TODO: check for overflow if not CHECK_OOB or is_image_shape(buf._shape): return True @@ -44,7 +44,7 @@ def type_verify(ast:UOp|list[UOp], check_spec:PatternMatcher): raise RuntimeError(f"UOp verification failed at {i} on {u.op} {u.dtype} {len(u.src)} {[(x.op, x.dtype, x.arg) for x in u.src]} {u.arg}") # ***** new specs ***** -def matches_dtype(x:UOp, dtype:DType) -> bool: return x.dtype == dtype or x.base.arg is Invalid # Invalid matches any dtype +def matches_dtype(x:UOp, dtype:DType) -> bool: return x.dtype == dtype or x.base.is_invalid # Invalid matches any dtype # these ops can be used in the tensor graph and programs spec_shared = PatternMatcher([ # NOTE: for testing, we let sinks be anything @@ -79,7 +79,7 @@ spec_shared = PatternMatcher([ (UPat(Ops.RANGE, src=(UPat.var("x"),), allow_any_len=True, name="rng"), lambda rng,x: matches_dtype(x, rng.dtype) and isinstance(rng.arg, tuple) and len(rng.arg) >= 2 and \ all(isinstance(ra, int) for ra in rng.arg[0:-1]) and isinstance(rng.arg[-1], AxisType)), - (UPat(Ops.INDEX, name="x"), lambda x: len(x.src)>0 and all(dtypes.is_int(y.dtype) or y.base.arg is Invalid for y in x.src[1:]) or None), + (UPat(Ops.INDEX, name="x"), lambda x: len(x.src)>0 and all(dtypes.is_int(y.dtype) or y.base.is_invalid for y in x.src[1:]) or None), # END closes RANGEs (UPat(Ops.END, src=(UPat(),), allow_any_len=True, name="x"), lambda x: all(u.op is Ops.RANGE for u in x.src[1:]) or None), # a loop-ended END requires a trailing bool condition for the backedge (loop again while true) diff --git a/tinygrad/uop/symbolic.py b/tinygrad/uop/symbolic.py index 4e9a5e1c66..8830abcd70 100644 --- a/tinygrad/uop/symbolic.py +++ b/tinygrad/uop/symbolic.py @@ -81,11 +81,11 @@ pm_data_invalid = PatternMatcher([ (invalid_pat.where(UPat(), UPat()), lambda i: i), (invalid_gate.where(UPat.var("a"), UPat.var("b")), lambda cond,x,i,a,b: cond.where(x.where(a,b), i)), # normalize where(cond, Invalid, val) -> where(~cond, val, Invalid) - (UPat.var("cond").where(invalid_pat, UPat.var("val")), lambda cond, i, val: cond.logical_not().where(val, i) if val.arg != Invalid else i), + (UPat.var("cond").where(invalid_pat, UPat.var("val")), lambda cond, i, val: cond.logical_not().where(val, i) if not val.is_invalid else i), # lift Invalid out: a.where(cond.where(x, Invalid), c) -> (~a|cond).where(a.where(x, c), Invalid) (UPat.var("a").where(invalid_gate, UPat.var("c")), lambda cond,i,x,a,c: - (a.logical_not()|cond).where(a.where(x,c), i) if c.arg != Invalid else None), - (UPat.var("a").where(UPat.var("b"), invalid_gate), lambda cond,i,x,a,b: (a|cond).where(a.where(b, x), i) if b.arg != Invalid else None), + (a.logical_not()|cond).where(a.where(x,c), i) if not c.is_invalid else None), + (UPat.var("a").where(UPat.var("b"), invalid_gate), lambda cond,i,x,a,b: (a|cond).where(a.where(b, x), i) if not b.is_invalid else None), # fold gated LOAD/STORE (UPat(Ops.STORE, src=(UPat(Ops.INDEX, src=(UPat(), invalid_pat), allow_any_len=True).or_casted(), UPat())), lambda i: UOp(Ops.NOOP)), (UPat(Ops.LOAD, src=(UPat(Ops.INDEX, src=(UPat(), invalid_pat), allow_any_len=True).or_casted(),), allow_any_len=True, name="x"), @@ -94,8 +94,8 @@ pm_data_invalid = PatternMatcher([ pm_remove_invalid = PatternMatcher([ (invalid_gate.named("w"), lambda cond,x,i,w: w.replace(src=(cond,x,w.const_like(0)))), - (UPat(Ops.STACK, name="s"), lambda s: s.replace(src=tuple(UOp.const(0, s.dtype) if x.arg is Invalid else x for x in s.src)) - if any(x.arg is Invalid for x in s.src) else None), + (UPat(Ops.STACK, name="s"), lambda s: s.replace(src=tuple(UOp.const(0, s.dtype) if x.is_invalid else x for x in s.src)) + if any(x.is_invalid for x in s.src) else None), ]) symbolic_simple = pm_data_invalid + PatternMatcher([ @@ -237,7 +237,7 @@ symbolic = symbolic_simple+commutative+PatternMatcher([ (UPat.cvar("y") * (UPat.var("x", dtype=dtypes.weakint) + UPat.cvar("c")), lambda x,y,c: (y*x)+(y*c)), # y*(x+c) -> y*x + y*c # ** where folding ** (UPat.var("cond", dtype=dtypes.bool).logical_not().where(UPat.var("t"), UPat.var("f")), - lambda cond, t, f: cond.where(f,t) if f.arg is not Invalid else None), + lambda cond, t, f: cond.where(f,t) if not f.is_invalid else None), # in cond.where(t, f), uses of cond fold to True within t and False within f (UPat.var("cond", dtype=dtypes.bool).where(UPat.var("t"), UPat.var("f")), fold_where_closure), # alu of two where with same conds can combine, only do if true branch or false branch is const From 665822ab34dc3a47c1c82e5898be0194336beec8 Mon Sep 17 00:00:00 2001 From: nimlgen <138685161+nimlgen@users.noreply.github.com> Date: Sat, 1 Aug 2026 16:32:04 +0300 Subject: [PATCH 40/44] hcq2: faster replace (#17353) * hcq2: use runtime device for submit params * hcq2: parameterize buffers in linear time --- tinygrad/runtime/support/hcq2.py | 14 +++++++++----- 1 file changed, 9 insertions(+), 5 deletions(-) diff --git a/tinygrad/runtime/support/hcq2.py b/tinygrad/runtime/support/hcq2.py index 61a8d3ebd3..008167ae3f 100644 --- a/tinygrad/runtime/support/hcq2.py +++ b/tinygrad/runtime/support/hcq2.py @@ -77,9 +77,11 @@ def encode_kernargs_clike(call:UOp, prg:UOp, devs:str|tuple[str, ...]) -> UOp: # ***************** # 0.1. prep: replace buffers with params -def replace_call_buffers(ctx:list[UOp], call:UOp) -> UOp|None: - ctx += [s for s in call.src[1:] if s not in ctx and s.op not in (Ops.PARAM, Ops.BIND)] - return call.replace(src=call.src[:1] + tuple(s if s.op in (Ops.PARAM, Ops.BIND) else s.param_like(ctx.index(s)) for s in call.src[1:])) +def replace_call_buffers(ctx:tuple[list[UOp], dict[UOp, int]], call:UOp) -> UOp|None: + bufs, slots = ctx + for s in call.src[1:]: + if s.op not in (Ops.PARAM, Ops.BIND) and slots.setdefault(s, len(bufs)) == len(bufs): bufs.append(s) + return call.replace(src=call.src[:1] + tuple(s if s.op in (Ops.PARAM, Ops.BIND) else s.param_like(slots[s]) for s in call.src[1:])) pm_replace_buffers = PatternMatcher([(UPat(Ops.CALL, name="call"), replace_call_buffers)]) # ***************** @@ -321,7 +323,7 @@ def replace_params(call:UOp) -> UOp|None: addrs = dedup([g.src[0].without_after for x in call.src for g in x.toposort() if g.op is Ops.GETADDR]) refhold += [a for a in addrs if a not in held and all(b.op is not Ops.PARAM or b.tag is not None for b in unwrap_mstack(a))] - sub = {(b:=u.without_after): UOp.param(i, u.dtype, shape=b.shape, device=u.device, volatile=b.op is Ops.PARAM and b.arg.volatile) + sub = {(b:=u.without_after): UOp.param(i, u.dtype, shape=b.shape, device=HCQ_RUNTIME_DEV.value, volatile=b.op is Ops.PARAM and b.arg.volatile) for i,u in enumerate(c_args)} | {v: v.replace(arg=replace(v.arg, slot=-1)) for v in variables if v.op is Ops.PARAM} info = replace(call.arg.aux, inputs=next((i for i,u in enumerate(c_args) if u.tag == "inputs"), None)) return call.replace(src=(body.substitute(sub).replace(arg="hcq_args"), *c_args, *refhold), @@ -374,7 +376,9 @@ hcq_compile_cache:dict[tuple[bytes, bool], UOp] = {} @track_rewrites(lambda linear,input_uops,jit,ret: f"HCQ Compile {pluralize('Kernel', len(ret.src))}") def hcq_compile(linear:UOp, input_uops:list[UOp]|None=None, jit=False) -> UOp: - if input_uops is not None: linear = graph_rewrite(linear, pm_replace_buffers, ctx=input_uops, walk=True, name="replace buffer") + if input_uops is not None: + slots = {u:i for i,u in reversed(tuple(enumerate(input_uops)))} + linear = graph_rewrite(linear, pm_replace_buffers, ctx=(input_uops, slots), walk=True, name="replace buffer") if (final_linear:=(hcq_compile_cache.get(cache_key:=(linear.key, jit)))) is None: # prep From 8e524ca467ebc854b80536611cc29bf9432fe152 Mon Sep 17 00:00:00 2001 From: chenyu Date: Sat, 1 Aug 2026 09:54:28 -0400 Subject: [PATCH 41/44] CONST related cleanups [pr] (#17352) ConstFloat(nan) != nan should be False, and some Invalid bool cleanups --- test/null/test_uops.py | 22 +++++++++++++++++++++- tinygrad/dtype.py | 4 ++-- tinygrad/uop/ops.py | 3 +-- tinygrad/uop/spec.py | 4 ++-- 4 files changed, 26 insertions(+), 7 deletions(-) diff --git a/test/null/test_uops.py b/test/null/test_uops.py index ea701a30ca..09f6e9a55a 100644 --- a/test/null/test_uops.py +++ b/test/null/test_uops.py @@ -5,7 +5,7 @@ from tinygrad.tensor import Tensor from tinygrad.helpers import Timing, Context, cdiv from tinygrad.dtype import dtypes, AddrSpace, ConstFloat, Invalid # noqa: F401 from tinygrad.device import Device -from tinygrad.uop.ops import Ops, ParamArg, UOp, UPat, dtype_from_uop, exec_alu, graph_rewrite, pm_lower_index_dtype # noqa: F401 # ParamArg used by eval(str(uop)) roundtrip tests +from tinygrad.uop.ops import Ops, ParamArg, PatternMatcher, UOp, UPat, dtype_from_uop, exec_alu, graph_rewrite, pm_lower_index_dtype # noqa: F401 # ParamArg used by eval(str(uop)) roundtrip tests from tinygrad.uop.spec import spec_program, spec_shared, type_verify from tinygrad.uop.symbolic import sym, pm_remove_invalid from test.helpers import eval_uop, to_uops_list @@ -45,6 +45,12 @@ class TestDTypeFromUOp(unittest.TestCase): with self.assertRaises(RuntimeError): type_verify(UOp.const(value, weak).sink(), spec_program) type_verify(UOp.const(value, concrete).sink(), spec_program) + def test_invalid_stated_dtype(self): + # UOp.const normalizes a stated dtype away (const_like/full pass their position's); the core constructor does not, + # and the spec is what rejects a non-bool Invalid + self.assertIs(UOp.const(Invalid, dtypes.float32), UOp.invalid()) + with self.assertRaises(RuntimeError): type_verify(UOp(Ops.CONST, dtypes.float32, arg=Invalid), spec_shared) + def test_invalid_dtype_and_consumers(self): invalid = UOp.invalid() self.assertIs(invalid.dtype, dtypes.bool) @@ -112,6 +118,20 @@ class TestSafeCast(unittest.TestCase): self.assertEqual(a.cast(dtypes.int8).cast(dtypes.int64).simplify(), a.cast(dtypes.int64)) self.assertEqual(a.cast(dtypes.int8).cast(dtypes.float).simplify(), a.cast(dtypes.float)) +class TestConstFloatEq(unittest.TestCase): + def test_nan_eq_ne_agree(self): + nan = dtypes.float32.const(math.nan) + self.assertTrue(nan == math.nan) + self.assertFalse(nan != math.nan) # float.__ne__ would say True here + self.assertFalse(nan == Invalid) + self.assertTrue(nan != Invalid) # __ne__ must defer to the reflected eq, not swallow NotImplemented + + def test_matchers_agree_on_nan(self): + n = UOp.const(math.nan, dtypes.float32) + for compiled in (False, True): + pm = PatternMatcher([(UPat(Ops.CONST, arg=math.nan), lambda: True)], compiled=compiled) + self.assertTrue(pm.rewrite(n), f"{compiled=}") + class TestExecALU(unittest.TestCase): def test_sqrt(self): self.assertEqual(exec_alu(Ops.SQRT, dtypes.float, (0.0,)), 0.0) diff --git a/tinygrad/dtype.py b/tinygrad/dtype.py index af8f7c7267..9272b7e617 100644 --- a/tinygrad/dtype.py +++ b/tinygrad/dtype.py @@ -17,6 +17,7 @@ class ConstFloat(float): if self is other: return True if isinstance(other, float) and math.isnan(self) and math.isnan(other): return True return float.__eq__(self, other) + def __ne__(self, other): return res if (res:=self.__eq__(other)) is NotImplemented else not res # float.__ne__ disagrees with __eq__ on nan def __hash__(self): return hash(self.bits) def __repr__(self): return f"ConstFloat({float.__repr__(self)})" def __str__(self): return float.__repr__(self) @@ -74,8 +75,7 @@ class DType(metaclass=DTypeMetaClass): def max(self): if dtypes.is_int(self): return 2**(self.bitsize)-1+self.min return float("inf") if dtypes.is_float(self) else True - def const(self, val: tuple[ConstType, ...]|ConstType): - if isinstance(val, tuple): return tuple(map(self.const, val)) + def const(self, val: ConstType): if isinstance(val, InvalidType): return val # NOTE: float('nan') != float('nan'), so we canonicalize here if isinstance(val, float) and math.isnan(val): val = math.nan diff --git a/tinygrad/uop/ops.py b/tinygrad/uop/ops.py index 8341eca529..4168089812 100644 --- a/tinygrad/uop/ops.py +++ b/tinygrad/uop/ops.py @@ -189,7 +189,6 @@ class UOpMetaClass(type): ucache:dict[tuple, weakref.ReferenceType[UOp]] = {} def __call__(cls, op:Ops, dtype:DType|None=None, src:tuple[UOp,...]=tuple(), arg:Any=None, tag:Any=None, metadata:tuple[Metadata,...]|None=None, _buffer:Buffer|None=None): - if op is Ops.CONST and arg is Invalid: dtype = dtypes.bool if dtype is None: dtype = dtype_from_uop(op, src, arg) or dtypes.void # CONST derives its dtype by value only when the constructor omits one # TODO: delete this once the dtype field is removed, for now it just re-implements spec.py @@ -618,7 +617,7 @@ class UOp(RandMixin, metaclass=UOpMetaClass): def alu(self, op, *src:UOp, **kwargs): return UOp(op, src=(self, *src), **kwargs) @staticmethod def const(b:ConstLike, dtype:DType|None=None): - if dtype is None: dtype = dtypes.from_py(b) + if dtype is None or b is Invalid: dtype = dtypes.from_py(b) if isinstance(b, UOp): return b.cast(dtype) # NOTE: it always has to be STACK now, even if they are all the same if isinstance(b, tuple): return UOp.stack(*[UOp.const(c, dtype) for c in b]) diff --git a/tinygrad/uop/spec.py b/tinygrad/uop/spec.py index 69b13657d0..a6af821b61 100644 --- a/tinygrad/uop/spec.py +++ b/tinygrad/uop/spec.py @@ -53,8 +53,8 @@ spec_shared = PatternMatcher([ # NOOP. TODO: remove this (UPat(Ops.NOOP), lambda: True), - # CONST is everywhere - (UPat(Ops.CONST, src=(), name="x"), lambda x: type(x.val) is type(x.dtype.const(x.val))), + # CONST is everywhere; Invalid is a bool const + (UPat(Ops.CONST, src=(), name="x"), lambda x: x.dtype is dtypes.bool if x.is_invalid else type(x.val) is type(x.dtype.const(x.val))), # STACK is everywhere too (UPat(Ops.STACK, dtype=dtypes.void, src=()), lambda: True), From 98b700bad17a371521da5dd083396cb1bbf9d4ac Mon Sep 17 00:00:00 2001 From: wozeparrot Date: Sun, 2 Aug 2026 00:10:06 +0800 Subject: [PATCH 42/44] gptoss optim fixes (#17356) --- examples/mlperf/optim.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/examples/mlperf/optim.py b/examples/mlperf/optim.py index dece377013..f93ccef119 100644 --- a/examples/mlperf/optim.py +++ b/examples/mlperf/optim.py @@ -2,7 +2,7 @@ from tinygrad.tensor import Tensor from tinygrad.dtype import dtypes from tinygrad.nn.optim import Optimizer, OptimizerGroup from tinygrad.helpers import FUSE_OPTIM, getenv -from tinygrad.uop.ops import UOp, Ops +from tinygrad.uop.ops import UOp, Ops, AxisType STOCHASTIC_ROUND = getenv("STOCHASTIC_ROUND", 0) MASTER_WEIGHTS = getenv("MASTER_WEIGHTS", 0) @@ -42,8 +42,8 @@ class GradAccClipAdamW(Optimizer): self.master_params = None def _zero_shard(self, t:Tensor) -> Tensor: - if not self.zero or (t.shape[0] % len(self.device)) != 0: return t - return Tensor(t.uop._shard(0, len(self.device)).unshard(0)).clone() + if not self.zero or t.ndim < 2 or (t.shape[0] % len(self.device)) != 0: return t + return Tensor(t.uop._shard(0, UOp.range(len(self.device), -1, AxisType.DEVICE)).unshard(0)).clone() def _zero_gather(self, t:Tensor) -> Tensor: if not isinstance(t.device, tuple) or t.uop.axis != 0: return t @@ -123,6 +123,9 @@ class GradAccClipAdamW(Optimizer): return out.shard_like(t) if offloaded else out class GradAccClipAdamWGroup(OptimizerGroup): + def __init__(self, *optimizers:GradAccClipAdamW): + super().__init__(*optimizers) + for o in self.optimizers[1:]: o.lr = self.optimizers[0].lr def fstep(self, grads:list[Tensor], grad_norm:Tensor|None=None): offset = 0 to_realize = [] From a54bb3b795791e4f85627df5e718d831a5f79d4b Mon Sep 17 00:00:00 2001 From: geohot Date: Thu, 30 Jul 2026 09:00:54 -0700 Subject: [PATCH 43/44] fix szdiff workflow for Gitea: use local master and REST API comments Two bugs prevented the 'Core Library Line Count' bot from working: 1. checkbranch fetched master from github.com/tinygrad/tinygrad instead of the local Gitea instance. Since the Gitea mirror lags behind GitHub, every PR appeared 'behind' and the szdiff job was always skipped. 2. marocchino/sticky-pull-request-comment@v3 uses GraphQL to find existing comments, but Gitea has no GraphQL API (returns 404). Both the szdiff and rebase comment steps failed with '404 page not found'. Fix 1: fetch origin/master (the local Gitea repo) instead of adding a remote to github.com. Fix 2: replace the third-party action with curl calls to the Gitea REST API (GET/POST/PATCH /repos/.../issues/.../comments), which fully supports the sticky-comment pattern (find existing, update or create). --- .github/workflows/szdiff.yml | 54 ++++++++++++++++++++++++------------ 1 file changed, 37 insertions(+), 17 deletions(-) diff --git a/.github/workflows/szdiff.yml b/.github/workflows/szdiff.yml index 37a985f0c5..7db9c3c79b 100644 --- a/.github/workflows/szdiff.yml +++ b/.github/workflows/szdiff.yml @@ -26,14 +26,13 @@ jobs: - name: Check whether branch is up-to-date id: brstat run: | - git remote add tinygrad https://github.com/tinygrad/tinygrad - git fetch tinygrad master + git fetch origin master echo "${{ github.event.pull_request.head.sha }}" - git rev-list --left-right --count tinygrad/master...${{ github.event.pull_request.head.sha }} | awk '{print "Behind "$1" - Ahead "$2""}' - count=$(git rev-list --left-right --count tinygrad/master...${{ github.event.pull_request.head.sha }} | awk '{print $1}') + git rev-list --left-right --count origin/master...${{ github.event.pull_request.head.sha }} | awk '{print "Behind "$1" - Ahead "$2""}' + count=$(git rev-list --left-right --count origin/master...${{ github.event.pull_request.head.sha }} | awk '{print $1}') if [ $count -gt 0 ] then - echo "Current branch is behind tinygrad master branch!" + echo "Current branch is behind master branch!" echo "stat=true" >> "$GITHUB_OUTPUT" else echo "stat=false" >> "$GITHUB_OUTPUT" @@ -75,13 +74,23 @@ jobs: python sz.py "$BASE" "$PR" > loc_content.txt - name: Comment Code Line Diff continue-on-error: false - uses: marocchino/sticky-pull-request-comment@v3 - with: + env: GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} - ignore_empty: true - skip_unchanged: true - recreate: true - path: loc_content.txt + run: | + set -eu + API="${GITHUB_API_URL}/repos/${GITHUB_REPOSITORY}/issues/${{ github.event.pull_request.number }}/comments" + BODY=$(python3 -c "import json; print(json.dumps(open('loc_content.txt').read()))") + # find existing sticky comment (contains "### Changes") + EXISTING_ID=$(curl -sf -H "Authorization: token $GITHUB_TOKEN" "$API" \ + | python3 -c "import sys,json; cs=json.load(sys.stdin); print(next((c['id'] for c in cs if '### Changes' in c.get('body','')), ''))") + if [ -n "$EXISTING_ID" ]; then + curl -sf -X PATCH -H "Authorization: token $GITHUB_TOKEN" -H "Content-Type: application/json" \ + "${GITHUB_API_URL}/repos/${GITHUB_REPOSITORY}/issues/comments/$EXISTING_ID" \ + -d "{\"body\": $BODY}" + else + curl -sf -X POST -H "Authorization: token $GITHUB_TOKEN" -H "Content-Type: application/json" \ + "$API" -d "{\"body\": $BODY}" + fi rebase: name: Core Library Line Difference @@ -93,10 +102,21 @@ jobs: steps: - name: Comment Rebase continue-on-error: false - uses: marocchino/sticky-pull-request-comment@v3 - with: + env: GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} - skip_unchanged: true - recreate: true - message: | - This branch currently is behind tinygrad/master. The line count difference bot is disabled. + run: | + set -eu + API="${GITHUB_API_URL}/repos/${GITHUB_REPOSITORY}/issues/${{ github.event.pull_request.number }}/comments" + MSG="This branch currently is behind master. The line count difference bot is disabled." + BODY=$(python3 -c "import json; print(json.dumps('''$MSG'''))") + # find existing sticky comment (contains "line count difference bot is disabled") + EXISTING_ID=$(curl -sf -H "Authorization: token $GITHUB_TOKEN" "$API" \ + | python3 -c "import sys,json; cs=json.load(sys.stdin); print(next((c['id'] for c in cs if 'line count difference bot is disabled' in c.get('body','')), ''))") + if [ -n "$EXISTING_ID" ]; then + curl -sf -X PATCH -H "Authorization: token $GITHUB_TOKEN" -H "Content-Type: application/json" \ + "${GITHUB_API_URL}/repos/${GITHUB_REPOSITORY}/issues/comments/$EXISTING_ID" \ + -d "{\"body\": $BODY}" + else + curl -sf -X POST -H "Authorization: token $GITHUB_TOKEN" -H "Content-Type: application/json" \ + "$API" -d "{\"body\": $BODY}" + fi From 60996454a3aeab1270ce0bc1a80a5e84f3893b2d Mon Sep 17 00:00:00 2001 From: George Hotz Date: Sat, 1 Aug 2026 10:44:41 -0700 Subject: [PATCH 44/44] szdiff: make comment steps work on both github and gitea - checkbranch: fetch master from the base repo clone_url instead of the PR head remote (origin), which on GitHub is the fork for fork PRs - replace the curl comment snippets with sticky_comment.py: stdlib-only, works with the REST API on both platforms, and keeps the skip_unchanged/ignore_empty behavior of the old action --- .github/workflows/sticky_comment.py | 34 +++++++++++++++++++ .github/workflows/szdiff.yml | 52 ++++++++++------------------- 2 files changed, 51 insertions(+), 35 deletions(-) create mode 100644 .github/workflows/sticky_comment.py diff --git a/.github/workflows/sticky_comment.py b/.github/workflows/sticky_comment.py new file mode 100644 index 0000000000..6b43ccfc8a --- /dev/null +++ b/.github/workflows/sticky_comment.py @@ -0,0 +1,34 @@ +#!/usr/bin/env python3 +# Sticky PR comment via the REST API: find an existing comment containing MARKER and PATCH it, or POST a new one. +# Works on GitHub and Gitea (stdlib only, replaces marocchino/sticky-pull-request-comment which needs GraphQL). +# Env vars: GITHUB_TOKEN, GITHUB_API_URL, GITHUB_REPOSITORY (set by the runner), PR_NUMBER, MARKER, and BODY_FILE or MESSAGE. +import json, os, sys, urllib.request + +api, repo = os.environ["GITHUB_API_URL"], os.environ["GITHUB_REPOSITORY"] +pr, marker = os.environ["PR_NUMBER"], os.environ["MARKER"] +body = open(os.environ["BODY_FILE"]).read() if os.environ.get("BODY_FILE") else os.environ["MESSAGE"] + +if not body.strip(): + print("comment body is empty, not posting") + sys.exit(0) + +def req(url, method="GET", payload=None): + r = urllib.request.Request(url, data=None if payload is None else json.dumps(payload).encode(), method=method, + headers={"Authorization": f"token {os.environ['GITHUB_TOKEN']}", "Accept": "application/json", "Content-Type": "application/json"}) + return json.load(urllib.request.urlopen(r)) + +# find the latest sticky comment (paginate, 100 comments per page) +existing, page = None, 1 +while True: + comments = req(f"{api}/repos/{repo}/issues/{pr}/comments?per_page=100&page={page}") + stickies = [c for c in comments if marker in (c.get("body") or "")] + if stickies: existing = stickies[-1] + if not comments or len(comments) < 100: break + page += 1 + +if existing is not None and existing["body"] == body: + print("comment is already up to date") + sys.exit(0) +url = f"{api}/repos/{repo}/issues/comments/{existing['id']}" if existing is not None else f"{api}/repos/{repo}/issues/{pr}/comments" +resp = req(url, 'PATCH' if existing is not None else 'POST', {'body': body}) +print(f"{'updated' if existing is not None else 'created'} comment {resp['id']}") diff --git a/.github/workflows/szdiff.yml b/.github/workflows/szdiff.yml index 7db9c3c79b..efeb92c3df 100644 --- a/.github/workflows/szdiff.yml +++ b/.github/workflows/szdiff.yml @@ -26,13 +26,14 @@ jobs: - name: Check whether branch is up-to-date id: brstat run: | - git fetch origin master + # fetch master from the base repo (tinygrad/tinygrad on GitHub, the mirror on Gitea), not the PR head remote + git fetch "${{ github.event.pull_request.base.repo.clone_url }}" master echo "${{ github.event.pull_request.head.sha }}" - git rev-list --left-right --count origin/master...${{ github.event.pull_request.head.sha }} | awk '{print "Behind "$1" - Ahead "$2""}' - count=$(git rev-list --left-right --count origin/master...${{ github.event.pull_request.head.sha }} | awk '{print $1}') + git rev-list --left-right --count FETCH_HEAD...${{ github.event.pull_request.head.sha }} | awk '{print "Behind "$1" - Ahead "$2""}' + count=$(git rev-list --left-right --count FETCH_HEAD...${{ github.event.pull_request.head.sha }} | awk '{print $1}') if [ $count -gt 0 ] then - echo "Current branch is behind master branch!" + echo "Current branch is behind ${{ github.event.pull_request.base.repo.full_name }} master branch!" echo "stat=true" >> "$GITHUB_OUTPUT" else echo "stat=false" >> "$GITHUB_OUTPUT" @@ -76,21 +77,11 @@ jobs: continue-on-error: false env: GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} - run: | - set -eu - API="${GITHUB_API_URL}/repos/${GITHUB_REPOSITORY}/issues/${{ github.event.pull_request.number }}/comments" - BODY=$(python3 -c "import json; print(json.dumps(open('loc_content.txt').read()))") - # find existing sticky comment (contains "### Changes") - EXISTING_ID=$(curl -sf -H "Authorization: token $GITHUB_TOKEN" "$API" \ - | python3 -c "import sys,json; cs=json.load(sys.stdin); print(next((c['id'] for c in cs if '### Changes' in c.get('body','')), ''))") - if [ -n "$EXISTING_ID" ]; then - curl -sf -X PATCH -H "Authorization: token $GITHUB_TOKEN" -H "Content-Type: application/json" \ - "${GITHUB_API_URL}/repos/${GITHUB_REPOSITORY}/issues/comments/$EXISTING_ID" \ - -d "{\"body\": $BODY}" - else - curl -sf -X POST -H "Authorization: token $GITHUB_TOKEN" -H "Content-Type: application/json" \ - "$API" -d "{\"body\": $BODY}" - fi + PR_NUMBER: ${{ github.event.pull_request.number }} + MARKER: "### Changes" + BODY_FILE: loc_content.txt + # note: run the script from the base checkout, never from the PR checkout + run: python3 "$GITHUB_WORKSPACE/base/.github/workflows/sticky_comment.py" rebase: name: Core Library Line Difference @@ -100,23 +91,14 @@ jobs: needs: checkbranch if: needs.checkbranch.outputs.branchstat == 'true' steps: + # pull_request_target: a plain checkout gets the base repo, so no PR code is executed + - uses: actions/checkout@v6 - name: Comment Rebase continue-on-error: false env: GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} - run: | - set -eu - API="${GITHUB_API_URL}/repos/${GITHUB_REPOSITORY}/issues/${{ github.event.pull_request.number }}/comments" - MSG="This branch currently is behind master. The line count difference bot is disabled." - BODY=$(python3 -c "import json; print(json.dumps('''$MSG'''))") - # find existing sticky comment (contains "line count difference bot is disabled") - EXISTING_ID=$(curl -sf -H "Authorization: token $GITHUB_TOKEN" "$API" \ - | python3 -c "import sys,json; cs=json.load(sys.stdin); print(next((c['id'] for c in cs if 'line count difference bot is disabled' in c.get('body','')), ''))") - if [ -n "$EXISTING_ID" ]; then - curl -sf -X PATCH -H "Authorization: token $GITHUB_TOKEN" -H "Content-Type: application/json" \ - "${GITHUB_API_URL}/repos/${GITHUB_REPOSITORY}/issues/comments/$EXISTING_ID" \ - -d "{\"body\": $BODY}" - else - curl -sf -X POST -H "Authorization: token $GITHUB_TOKEN" -H "Content-Type: application/json" \ - "$API" -d "{\"body\": $BODY}" - fi + PR_NUMBER: ${{ github.event.pull_request.number }} + MARKER: "line count difference bot is disabled" + MESSAGE: | + This branch currently is behind ${{ github.event.pull_request.base.repo.full_name }} master. The line count difference bot is disabled. + run: python3 .github/workflows/sticky_comment.py