test_GPU_time() — pytorch Function Reference

Architecture documentation for the test_GPU_time() function in inductor_mm.py from the pytorch codebase.

Function python DynamoBenchmarks Configuration calls 4 called by 1

Entity Profile

DynamoBenchmarks→ Configuration→ test_GPU_time() — pytorch Function Reference

Dependency Diagram

graph TD
  dd08baca_e7b7_4c73_60cc_7d9f6fa55a3a["test_GPU_time()"]
  31b7d7cf_d369_349c_d3fb_295720c1291f["shapes()"]
  31b7d7cf_d369_349c_d3fb_295720c1291f -->|calls| dd08baca_e7b7_4c73_60cc_7d9f6fa55a3a
  3717ae06_9fef_f61d_b6ef_83761c73d71b["inductor_aten_mm()"]
  dd08baca_e7b7_4c73_60cc_7d9f6fa55a3a -->|calls| 3717ae06_9fef_f61d_b6ef_83761c73d71b
  f175abea_1f77_fbbd_4800_616258913520["inductor_triton_mm()"]
  dd08baca_e7b7_4c73_60cc_7d9f6fa55a3a -->|calls| f175abea_1f77_fbbd_4800_616258913520
  3994379c_800b_4e6e_b053_5276e27e0062["torch_mm()"]
  dd08baca_e7b7_4c73_60cc_7d9f6fa55a3a -->|calls| 3994379c_800b_4e6e_b053_5276e27e0062
  4571a7c8_7237_ace0_feed_bcdebd3d9115["triton_mm()"]
  dd08baca_e7b7_4c73_60cc_7d9f6fa55a3a -->|calls| 4571a7c8_7237_ace0_feed_bcdebd3d9115
  style dd08baca_e7b7_4c73_60cc_7d9f6fa55a3a fill:#6366f1,stroke:#818cf8,color:#fff

Relationship Graph

Source Code

benchmarks/dynamo/microbenchmarks/inductor_mm.py lines 64–86

def test_GPU_time(shapes):
    print("shape; torch mm; triton mm; inductor aten mm; inductor triton mm")
    for i in range(len(shapes)):
        a_shape, b_shape = shapes[i]
        print(a_shape, "x", b_shape, end="; ")
        a = torch.randn(a_shape, device="cuda", dtype=torch.float16)
        b = torch.randn(b_shape, device="cuda", dtype=a.dtype)

        config.triton.mm = "aten"
        inductor_aten_mm(a, b)

        config.triton.mm = "triton"
        inductor_triton_mm(a, b)

        torch_ms, _, _ = benchmarker.benchmark_gpu(lambda: torch_mm(a, b))
        triton_ms, _, _ = benchmarker.benchmark_gpu(lambda: triton_mm(a, b))
        ind_aten_ms, _, _ = benchmarker.benchmark_gpu(lambda: inductor_aten_mm(a, b))
        ind_triton_ms, _, _ = benchmarker.benchmark_gpu(
            lambda: inductor_triton_mm(a, b)
        )
        print(torch_ms, triton_ms, ind_aten_ms, ind_triton_ms, sep="; ")

        torch._dynamo.reset()