From 0b00981cd1a082fcdb1e6b6c8de0d4ead9cdb71b Mon Sep 17 00:00:00 2001 From: George Hotz Date: Wed, 15 Oct 2025 09:38:07 +0800 Subject: [PATCH] fix wmma --- tinygrad/uop/ops.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tinygrad/uop/ops.py b/tinygrad/uop/ops.py index b11c734cc2..517554c7bb 100644 --- a/tinygrad/uop/ops.py +++ b/tinygrad/uop/ops.py @@ -264,7 +264,9 @@ class UOp(MathTrait, metaclass=UOpMetaClass): # NOTE: ssimplify is required because the shape needs to be canonical for broadcasting and same shape checking if self.op in GroupOp.Movement.union({Ops.MULTI, Ops.REDUCE_AXIS, Ops.WMMA}): ps = self.src[0]._shape - if ps is None: raise RuntimeError(f"movement op {self.op} requires shape on {self}") + # TODO: WMMA is used for both axis WMMA and op WMMA. fix this and remove this hack. tested by BERT on AMD LLVM + if ps is None and self.op is Ops.WMMA: return None + if ps is None: raise RuntimeError(f"movement op {self.op} requires shape") match self.op: case Ops.RESHAPE: if not all(x >= 0 for x in self.arg): raise ValueError(f"shape can't contain negative numbers {self.arg}")