Skip to content

Commit 1970a80

Browse files
Yuzhen11pytorchmergebot
authored andcommitted
Only print bw result for the first time we benchmark a kernel (pytorch#123568)
Summary: As title. Before this change, we use the benchmark result saved as cache and print out every time we call a kernel. The information is the same. Let's just print out at the first iteration. Test Plan: Local test. Differential Revision: D55878382 Pull Request resolved: pytorch#123568 Approved by: https://github.com/jackiexu1992
1 parent 5712c32 commit 1970a80

File tree

1 file changed

+7
-7
lines changed

1 file changed

+7
-7
lines changed

torch/_inductor/triton_heuristics.py

+7-7
Original file line numberDiff line numberDiff line change
@@ -845,13 +845,13 @@ def run(self, *args, grid, stream):
845845
if num_gb is None:
846846
num_gb = get_num_bytes(*args, num_in_out_args=num_in_out_ptrs) / 1e9
847847
gb_per_s = num_gb / (ms / 1e3)
848-
self.cached = (ms, num_gb, gb_per_s, kernel_name)
849-
else:
850-
ms, num_gb, gb_per_s, kernel_name = self.cached
851-
collected_calls.append((ms, num_gb, gb_per_s, kernel_name))
852-
print(
853-
create_bandwidth_info_str(ms, num_gb, gb_per_s, suffix=f" \t {kernel_name}")
854-
)
848+
self.cached = ms, num_gb, gb_per_s, kernel_name
849+
collected_calls.append((ms, num_gb, gb_per_s, kernel_name))
850+
print(
851+
create_bandwidth_info_str(
852+
ms, num_gb, gb_per_s, suffix=f" \t {kernel_name}"
853+
)
854+
)
855855

856856

857857
def hash_configs(configs: List[Config]):

0 commit comments

Comments
 (0)