defer realize folding to kernel splitting [pr] (#7849)

* defer realize folding to schedule breaking [pr]

* this is init

* p2

* need to lookup edges

* refactor image cast folding [pr]

* Ops.LOAD diff

* image works

* refactor can_pad

* fix fold_img_cast
This commit is contained in:
qazal
2024-11-23 14:29:14 +08:00
committed by GitHub
parent 144e9f00df
commit 5b2c03e865
3 changed files with 27 additions and 21 deletions
+1 -1
View File
@@ -441,7 +441,7 @@ class Kernel:
check(not self.vars, "does not work with symbolic shape")
check(axis < self.first_upcast, "cannot pad upcasted")
# ok to pad SUM if all parent ALU ops have f(0) = 0
if (r:=self.reduceop) is not None and self.first_reduce <= axis: check(r.arg[0] is Ops.ADD and can_pad(r), f"cannot pad {r}")
if (r:=self.reduceop) is not None and self.first_reduce <= axis: check(r.arg[0] is Ops.ADD and can_pad(r, {}, set()), f"cannot pad {r}")
padded = False
for i,st in enumerate(self.sts):
if (s:=st.shape[axis]) == 1: continue # reduced