improve constant folding, detach before moving tensor
MODIFY
accel/lazy/ops_lazy.py
accel/lazy/ops_lazy.py
@@ -187,8 +187,9 @@ def _realize_binary_op(self:LazyBuffer) -> Tuple[gops.GPUBuffer, List[gops.GPUBu
lazy_srcs : List[LazyBuffer] = [seen.setdefault(x,x) for x in get_lazybuffers(self.op) if x not in seen]
real_dict : Dict[LazyBuffer, str] = {}
for s in lazy_srcs:
- if s.optype == MovementOps and s.realized is None:
- root = get_lazybuffers(s.op)[0]
+ if s.optype in [LoadOps, MovementOps] and s.realized is None:
+ if s.optype == MovementOps: root = get_lazybuffers(s.op)[0]
+ else: root = s
# NOTE: if this is used, there can be 0 sources for a kernel
if FOLD_CONSTANTS_INTO_KERNELS and root.realized is None and root.optype == LoadOps and root.op.op == LoadOps.FROMCPU and root.shape == (1,):
if not s.st.needs_valid():
@@ -198,6 +199,8 @@ def _realize_binary_op(self:LazyBuffer) -> Tuple[gops.GPUBuffer, List[gops.GPUBu
inline_valid = s.st.expr().replace("valid=valid && ", "").replace(";idx=0", "").replace("//", "/").replace("idx", "gid")
if ';' not in inline_valid:
real_dict[s] = f"(({inline_valid}) * {str(root.op.arg[0])}f)"
+ else:
+ print(f"couldn't fold {str(root.op.arg[0])} with expr {s.st.expr()}")
if s not in real_dict: # nicer way to write this?
real_dict[s] = f"arg_{len(real_srcs)}"
real_srcs.append((f"arg_{len(real_srcs)}", s.realize()))
@@ -207,6 +210,7 @@ def _realize_binary_op(self:LazyBuffer) -> Tuple[gops.GPUBuffer, List[gops.GPUBu
def _realize(self:LazyBuffer) -> Tuple[gops.GPUBuffer, List[gops.GPUBuffer]]:
if self.optype == LoadOps and self.op.op == LoadOps.FROMCPU:
#print("load", self, self.shape, self.op.arg if prod(self.shape) == 1 else "<data>")
+ if self.shape == (1,): print("NOTE: resolving unary Tensor", self)
return gops.GPUBuffer.fromCPU(self.op.arg), []
elif self.optype == LoadOps and self.op.op == LoadOps.CONTIGUOUS:
real_src = self.op.src[0].realize()

MODIFY
examples/benchmark_train_efficientnet.py
examples/benchmark_train_efficientnet.py
@@ -49,7 +49,7 @@ if __name__ == "__main__":
p.realize()
et = time.monotonic()
- loss = loss.cpu().data[0]
+ loss = loss.detach().cpu().data[0]
print(f"{(et-st)*1000.0:7.2f} ms run, {(mt-st)*1000.0:7.2f} ms build, {loss:7.2f} loss, {tensors_allocated():4d} tensors, {get_memory_used()/1e9:.2f} GB used")

MODIFY
extra/training.py
extra/training.py
@@ -35,7 +35,7 @@ def train(model, X_train, Y_train, optim, steps, BS=128, lossfn=sparse_categoric
accuracy = (cat == y).mean()
# printing
- loss = loss.cpu().data
+ loss = loss.detach().cpu().data
losses.append(loss)
accuracies.append(accuracy)
t.set_description("loss %.2f accuracy %.2f" % (loss, accuracy))

MODIFY
test/test_mnist.py
test/test_mnist.py
@@ -50,6 +50,8 @@ class TestMNIST(unittest.TestCase):
model = TinyBobNet()
optimizer = optim.SGD(model.parameters(), lr=0.001)
train(model, X_train, Y_train, optimizer, BS=69, steps=1)
+ for p in model.parameters():
+ p.realize()
def test_sgd_threestep(self):
np.random.seed(1337)

improve constant folding, detach before moving tensor