From 90e4dc7dc8a130d3c06488fb0f0385429157280e Mon Sep 17 00:00:00 2001 From: qazal Date: Tue, 28 May 2024 23:11:25 +0300 Subject: [PATCH] ugh --- tinygrad/codegen/linearizer.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tinygrad/codegen/linearizer.py b/tinygrad/codegen/linearizer.py index ef135217ad..60fbf191ea 100644 --- a/tinygrad/codegen/linearizer.py +++ b/tinygrad/codegen/linearizer.py @@ -258,7 +258,8 @@ class Linearizer(Kernel): assert not locals_to_store, "storing locals isn't supported here" # load earlybufs - loaded_buffers.update({b:self.global_load(self.bufs.index(self.local_alias[i]) if i in self.local_alias else i, + loaded_buffers.update({b:self.global_load(self.bufs.index([lb for aliases in self.local_alias.values() for lb in aliases.values() + if lb is b][0]) if any(i in aliases for aliases in self.local_alias.values()) else i, global_idxs+local_idxs+reduce_idxs+full_upcast_idxs) for i,b in enumerate(self.bufs) if b in self.earlybufs}) # run early AST (with reduce)