From 11b5895c85ec3303151f59af8ee4e664b2e9ff71 Mon Sep 17 00:00:00 2001 From: George Hotz Date: Fri, 16 May 2025 20:09:41 -0700 Subject: [PATCH] hotfix: schedule timing in tensor.py --- tinygrad/engine/schedule.py | 3 +-- tinygrad/tensor.py | 7 +++++-- 2 files changed, 6 insertions(+), 4 deletions(-) diff --git a/tinygrad/engine/schedule.py b/tinygrad/engine/schedule.py index 27bf3dd44d..615abf3980 100644 --- a/tinygrad/engine/schedule.py +++ b/tinygrad/engine/schedule.py @@ -2,7 +2,7 @@ from dataclasses import dataclass, field from collections import deque, defaultdict from tinygrad.ops import UOp, Variable, Ops, UPat, PatternMatcher, graph_rewrite, buffers from tinygrad.device import Buffer -from tinygrad.helpers import Metadata, DEBUG, unwrap, merge_dicts +from tinygrad.helpers import Metadata, unwrap, merge_dicts # **** ScheduleItem return type @@ -66,7 +66,6 @@ def create_schedule_with_vars(sched_sink:UOp) -> tuple[list[ScheduleItem], dict[ # confirm everything was scheduled correctly assert len(schedule) == len(in_degree), f"Schedule length mistmatch {len(schedule)} != {len(in_degree)}" - if DEBUG >= 1 and len(schedule) >= 10: print(f"scheduled {len(schedule)} kernels") # map ASSIGN to BUFFER after ScheduleItems are constructed becomes_map = {u:u.buf_uop for u in toposort if u.op is Ops.ASSIGN} diff --git a/tinygrad/tensor.py b/tinygrad/tensor.py index fd4e32f98e..5f72c92d6e 100644 --- a/tinygrad/tensor.py +++ b/tinygrad/tensor.py @@ -6,7 +6,7 @@ from typing import Callable, ClassVar, Sequence, cast, get_args, Literal, Suppor from tinygrad.dtype import DType, DTypeLike, dtypes, ImageDType, ConstType, least_upper_float, least_upper_dtype, sum_acc_dtype, to_dtype, truncate from tinygrad.dtype import _from_np_dtype, _to_np_dtype from tinygrad.helpers import argfix, make_tuple, flatten, prod, all_int, round_up, merge_dicts, argsort, getenv, all_same, fully_flatten, dedup -from tinygrad.helpers import IMAGE, WINO, Metadata, TRACEMETA, ceildiv, fetch, polyN, unwrap +from tinygrad.helpers import IMAGE, WINO, Metadata, TRACEMETA, ceildiv, fetch, polyN, unwrap, DEBUG from tinygrad.engine.multi import get_multi_map from tinygrad.gradient import compute_gradient from tinygrad.ops import smax, smin, resolve, UOp, Ops, sint, Variable, SimpleMathTrait, identity_element, all_metadata @@ -254,10 +254,13 @@ class Tensor(SimpleMathTrait): NOTE: A Tensor can only be scheduled once. """ + st = time.perf_counter() self.kernelize(*lst) schedule, var_vals, becomes_map = create_schedule_with_vars(UOp.sink(*[x.lazydata for x in (self,)+lst])) _apply_map_to_tensors(becomes_map, name="Apply Schedule Map") - return memory_planner(schedule), var_vals + schedule = memory_planner(schedule) + if DEBUG >= 1 and len(schedule) >= 10: print(f"scheduled {len(schedule)} kernels in {(time.perf_counter()-st)*1000:.2f} ms") + return schedule, var_vals def schedule(self, *lst:Tensor) -> list[ScheduleItem]: """Creates the schedule needed to realize these Tensor(s)."""