From 1037defce56d239dc0fec502a2a7f5dfd76740f9 Mon Sep 17 00:00:00 2001 From: George Hotz Date: Sun, 24 Aug 2025 16:46:48 -0700 Subject: [PATCH] working WMMA --- tinygrad/codegen/opt/postrange.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/tinygrad/codegen/opt/postrange.py b/tinygrad/codegen/opt/postrange.py index d1e441dc05..f2d33cc232 100644 --- a/tinygrad/codegen/opt/postrange.py +++ b/tinygrad/codegen/opt/postrange.py @@ -57,9 +57,8 @@ class RKernel(Kernel): base_upcast_axes = tuple([(s,2) for s in self.shape_str_to_axis(tc.base_upcast_axes())]) tc_upcast_axes = tuple([base_upcast_axes[:int(math.log2(tc.elements_per_thread[i]))] for i in range(3)]) - tc_reduce_axes = tuple([self.rng[x].arg[0] for x in tc_reduce_axes]) - # TODO: remove tc_upcast_axes from the arg - #tc_upcast_axes = [(self.rng[x[0]].arg[0], x[1]) for x in tc_upcast_axes] + tc_upcast_axes = tuple([tuple([(self.rng[a].arg[0], sz) for a,sz in v]) for v in tc_upcast_axes]) + tc_reduce_axes = tuple([self.rng[a].arg[0] for a in tc_reduce_axes]) # construct the op wmma_arg = (str(tc), tc.dims, tc.dtype_in, tc.dtype_out, self.opts.device, tc.threads, tc_upcast_axes, tc_reduce_axes)