diff --git a/tinygrad/codegen/gpudims.py b/tinygrad/codegen/gpudims.py index 08d20fa4ed..e2981ec90b 100644 --- a/tinygrad/codegen/gpudims.py +++ b/tinygrad/codegen/gpudims.py @@ -1,6 +1,6 @@ import math, functools, operator from tinygrad.uop.ops import UOp, Ops, sint, PatternMatcher, UPat, KernelInfo, ssimplify, AxisType, graph_rewrite -from tinygrad.helpers import all_int, partition, flatten, prod, dedup +from tinygrad.helpers import all_int, partition, flatten, prod, dedup, USE_TC from tinygrad.dtype import dtypes, AddrSpace from tinygrad.shape.view import get_contraction from tinygrad.renderer import Renderer @@ -136,6 +136,7 @@ pm_add_gpudims = PatternMatcher([ ]) def apply_tensor_cores(ctx:tuple[dict, Renderer], in0:UOp, in1:UOp, r_range:UOp, reduceop:UOp): + if not USE_TC: return None # tensor cores have three ranges. X, Y, and REDUCE in0_ranges = [u for u in in0.ranges if u not in in1.ranges] in1_ranges = [u for u in in1.ranges if u not in in0.ranges] @@ -202,5 +203,5 @@ pm_tensor_cores = PatternMatcher([ # replace range #(UPat(Ops.RANGE, name="r"), lambda ctx,r: ctx[0].get(r, None)), - (UPat(Ops.SINK, name="s"), lambda ctx,s: graph_rewrite(s.substitute(ctx[0]), pm_flatten_range)), + (UPat(Ops.SINK, name="s"), lambda ctx,s: graph_rewrite(s.substitute(ctx[0]), pm_flatten_range, name="flatten")), ]) \ No newline at end of file