diff --git a/test/test_rangeify.py b/test/test_rangeify.py index 1ec988e3f9..a75ad1df9d 100644 --- a/test/test_rangeify.py +++ b/test/test_rangeify.py @@ -55,7 +55,17 @@ class TestDoubleMatmul(unittest.TestCase): def test_demote(self): self._test((Opt(OptOps.DEMOTE, 2, 8),)) @unittest.skipUnless(Device.DEFAULT == "METAL", "only for METAL TC") - def test_demote_tc(self): self._test((Opt(OptOps.DEMOTE, 2, 8), (Opt(OptOps.TC, 0, (0, 0, 1))))) + def test_demote_tc_top(self): + self._test((Opt(OptOps.DEMOTE, 2, 8), Opt(OptOps.TC, 0, (0, 0, 1, 0)))) + + @unittest.skipUnless(Device.DEFAULT == "METAL", "only for METAL TC") + def test_demote_tc_bottom(self): + self._test((Opt(OptOps.DEMOTE, 2, 8), Opt(OptOps.TC, 0, (0, 0, 1, 1)))) + + @unittest.skip("broken") + @unittest.skipUnless(Device.DEFAULT == "METAL", "only for METAL TC") + def test_demote_tc_both(self): + self._test((Opt(OptOps.DEMOTE, 2, 8), Opt(OptOps.TC, 0, (0, 0, 1, 0)), Opt(OptOps.TC, 0, (0, 0, 1, 1)))) class TestRangeifyAssign(unittest.TestCase): def test_assign_permuted(self): diff --git a/tinygrad/codegen/opt/postrange.py b/tinygrad/codegen/opt/postrange.py index dc7adfaff8..506de35221 100644 --- a/tinygrad/codegen/opt/postrange.py +++ b/tinygrad/codegen/opt/postrange.py @@ -194,11 +194,12 @@ class Scheduler: elif opt.op is OptOps.TC: #check(len(self.applied_opts) == 0, "tensor core opts must be first") # TODO: remove the need for this by having warps check(opt.axis is not None, "tensor core opts must have an axis") - check(opt.arg is not None and isinstance(opt.arg, tuple) and len(opt.arg) == 3, "tensor core opts must have valid arg") - check(-1 <= (tc_select:=cast(tuple, opt.arg)[0]) < len(self.ren.tensor_cores), "tensor core opts must have valid tc_select") - check(0 <= (tc_opt:=cast(tuple, opt.arg)[1]) <= 2, "tensor core opts must have valid tc_opt") - check(0 < (use_tensor_cores:=cast(tuple, opt.arg)[2]) <= 2, "use_tensor_cores value is not valid") - try: ret = self._apply_tc_opt(use_tensor_cores, cast(int, opt.axis), tc_select, tc_opt) + check(opt.arg is not None and isinstance(opt.arg, tuple) and len(opt.arg) >= 3, "tensor core opts must have valid arg") + assert isinstance(opt.arg, tuple) + check(-1 <= (tc_select:=opt.arg[0]) < len(self.ren.tensor_cores), "tensor core opts must have valid tc_select") + check(0 <= (tc_opt:=opt.arg[1]) <= 2, "tensor core opts must have valid tc_opt") + check(0 < (use_tensor_cores:=opt.arg[2]) <= 2, "use_tensor_cores value is not valid") + try: ret = self._apply_tc_opt(use_tensor_cores, cast(int, opt.axis), tc_select, tc_opt, opt.arg[3] if len(opt.arg) > 3 else 0) except ValueError as e: raise KernelOptError(str(e)) check(ret is not None, "no tensor core available") elif opt.op is OptOps.PADTO: @@ -233,10 +234,10 @@ class Scheduler: if append_opt: self.applied_opts.append(opt) return ret - def _apply_tc_opt(self, use_tensor_cores:int, axis:int, tc_select:int, opt_level:int) -> None|list[UOp]: + def _apply_tc_opt(self, use_tensor_cores:int, axis:int, tc_select:int, opt_level:int, reduce_choice:int) -> None|list[UOp]: reduceops = [x for x in self.ast.toposort() if x.op is Ops.REDUCE] if not len(reduceops): raise KernelOptError("no reduce ops for TensorCore") - reduceop = reduceops[0] + reduceop = reduceops[reduce_choice] if use_tensor_cores and reduceop is not None and reduceop.arg is Ops.ADD: mul = reduceop.src[0] if reduceop.src[0].op is not Ops.CAST else reduceop.src[0].src[0] if mul.op is not Ops.MUL: return None