From 363fbfc2e4dd7d67b49f75c2ccd8e06fefae4351 Mon Sep 17 00:00:00 2001 From: David Hou Date: Sun, 2 Jul 2023 18:33:57 -0700 Subject: [PATCH] do not emit loop end code for global+local loops in assembly kernel (#1100) --- tinygrad/codegen/assembly.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tinygrad/codegen/assembly.py b/tinygrad/codegen/assembly.py index df2ac19954..3997b349da 100644 --- a/tinygrad/codegen/assembly.py +++ b/tinygrad/codegen/assembly.py @@ -125,7 +125,7 @@ class AssemblyCodegen(Linearizer): ins.append(AssemblyInstruction(UOps.CONST, newreg(var, dtype=dtypes.int32, scalar=True), [], 0)) ins.append(AssemblyInstruction(UOps.LABEL, None, [], "$loop_"+var.expr)) elif uop == UOps.ENDLOOP: - if args[1] not in ["global", "local"]: + if args[1] not in ["global", "local", "global+local"]: for var in reversed(args[0]): if not isinstance(var, NumNode): # TODO: why is this coming through? ins.append(AssemblyInstruction(UOps.ALU, tor[var], [tor[var], 1], BinaryOps.ADD))