nimlgen/tinygrad
1
0
Fork 0
forked from tinygrad/tinygrad
Code Pull Requests Activity

fix POSTOPT pad (#11999)

Browse Source 
* fix POSTOPT=1

* fix some tests

* Revert "fix some tests"

This reverts commit 8ee058e206.

* fix padding restrictions

* cuda has two tensor cores

* Set POSTOPT ContextVar to 0 in helpers.py
This commit is contained in:
George Hotz
2025-09-04 14:28:58 -07:00
committed by GitHub
parent da61b40604
commit 30eb42a69e
2 changed files with 9 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/opt/test_tensor_cores.py
+1 -1
View File
@@ -114,7 +114,7 @@ class TestTensorCores(unittest.TestCase):
helper_tc_ensure_uops_and_opts_count(tc.dims[0]//4, tc.dims[1], tc.dims[2], tc.dtype_in, tc.dtype_out, tc_opt=2, ensure_triggered=False)
helper_tc_ensure_uops_and_opts_count(tc.dims[0], tc.dims[1]//4, tc.dims[2], tc.dtype_in, tc.dtype_out, tc_opt=2, ensure_triggered=False)
if not AMX: # AMX tc.dims[2] == 1
helper_tc_ensure_uops_and_opts_count(tc.dims[0], tc.dims[1], tc.dims[2]//4, tc.dtype_in, tc.dtype_out, tc_opt=2, ensure_triggered=False)
helper_tc_ensure_uops_and_opts_count(tc.dims[0], tc.dims[1], tc.dims[2]//8, tc.dtype_in, tc.dtype_out, tc_opt=2, ensure_triggered=False)
@unittest.skipIf(Device.DEFAULT == "PYTHON", "not generated on EMULATED device")
@unittest.skipIf(CI and Device.DEFAULT in {"AMD"}, "AMD CI is really slow here")
tinygrad/codegen/opt/postrange.py
+8 -4
View File
@@ -185,7 +185,9 @@ class Scheduler:
check(self._apply_tc_opt(use_tensor_cores, cast(int, opt.axis), tc_select, tc_opt), "no tensor core available")
elif opt.op is OptOps.PADTO:
check(rng.src[0].op is Ops.CONST, "only pad const")
replaced_rng = UOp.range(round_up(rng.vmax+1, cast(int, opt.arg)), *rng.arg)
new_sz = round_up(rng.vmax+1, cast(int, opt.arg))
check(rng.vmax+1 > new_sz//4, "pad adds more than quadruple the work")
replaced_rng = UOp.range(new_sz, *rng.arg)
replaces = {rng:replaced_rng}
for b in self.bufs:
if rng in b.src[1].sparents:
@@ -239,10 +241,12 @@ class Scheduler:
# do optimizations and save the ranges
try:
for i,a in enumerate(axes):
# apply_opt should return the updated range?
idx = self.rngs.index(a)
self.apply_opt(Opt(OptOps.PADTO, idx, tc.dims[i]), append_opt=False) # PADTO might fail
axes[i] = self.rngs[idx]
if (a.vmax+1) % tc.dims[i] != 0:
if opt_level < 2: raise KernelOptError("tc padding requires opt_level >= 2")
# apply_opt should return the updated range?
self.apply_opt(Opt(OptOps.PADTO, idx, tc.dims[i]), append_opt=False) # PADTO might fail
axes[i] = self.rngs[idx]
except KernelOptError: continue
ne: list[UOp] = []