fix CUDA reduce
MODIFY
tinygrad/llops/ops_gpu.py
tinygrad/llops/ops_gpu.py
@@ -213,7 +213,7 @@ class CLASTKernel(ASTKernel):
print("old:", [x.shape for x in self.sts])
print("old:", [x.views[-1].strides for x in self.sts])
- if not CUDA: self.hand_coded_optimizations()
+ self.hand_coded_optimizations()
# add a local buffer for multistage reduce
if len(self.group_for_reduce):
@@ -260,8 +260,9 @@ class CLASTKernel(ASTKernel):
lidx, lvalid = self.sts[-1].expr_idxs()
assert str(lvalid) == "1", "local buffer must be valid"
- self.kernel.append(f"__local {accumulators[0].decltype()} {self.buftokens[-1].tok}[{prod(self.group_for_reduce)}]; // second stage\n")
- self.kernel.append(f"int mid_idx = {lidx.cl}; {self.buftokens[-1].tok}[mid_idx] = {accumulators[0].tok}; barrier(CLK_LOCAL_MEM_FENCE);\n")
+ self.kernel.append(("__shared__ " if CUDA else "__local ") + f"{accumulators[0].decltype()} temp[{prod(self.group_for_reduce)}]; // second stage\n")
+ self.kernel.append(f"int mid_idx = {lidx.cl}; {self.buftokens[-1].tok}[mid_idx] = {accumulators[0].tok};\n")
+ self.kernel.append("barrier(CLK_LOCAL_MEM_FENCE);\n" if not CUDA else "__syncthreads();\n")
if self.upcast_in_mid_reduce:
assert len(self.group_for_reduce) == 2

MODIFY
tinygrad/runtime/cuda.py
tinygrad/runtime/cuda.py
@@ -19,6 +19,9 @@ class CLProgram:
self.prg = SourceModule(prg).get_function(name)
def __call__(self, global_size, local_size, *args):
- global_size = global_size + [1] * (2 - len(global_size))
+ local_size = (local_size + [1] * (3 - len(local_size))) if local_size is not None else (1,1,1)
+ global_size = global_size + [1] * (3 - len(global_size))
+ assert all(x%y == 0 for x,y in zip(global_size, local_size)), f"local:{local_size} must divide global:{global_size}"
+ global_size = [x//y for x,y in zip(global_size, local_size)]
if DEBUG >= 2: print("CUDA launch", global_size, local_size)
- self.prg(*args, block=(1,1,1), grid=tuple(global_size))
\ No newline at end of file
+ self.prg(*args, block=tuple(local_size), grid=tuple(global_size))
\ No newline at end of file

fix CUDA reduce