a_shapes() — pytorch Function Reference

Architecture documentation for the a_shapes() function in matmul_relu.py from the pytorch codebase.

Function python DynamoBenchmarks Microbenchmarking calls 1

Entity Profile

DynamoBenchmarks→ Microbenchmarking→ a_shapes() — pytorch Function Reference

Dependency Diagram

graph TD
  30f99d63_57a2_51c8_7e20_9fb6ffcd7f76["a_shapes()"]
  fef5e6fa_eee5_e93a_0d8f_a4ef0477af02["time_with_torch_timer()"]
  30f99d63_57a2_51c8_7e20_9fb6ffcd7f76 -->|calls| fef5e6fa_eee5_e93a_0d8f_a4ef0477af02
  style 30f99d63_57a2_51c8_7e20_9fb6ffcd7f76 fill:#6366f1,stroke:#818cf8,color:#fff

Relationship Graph

Source Code

benchmarks/dynamo/microbenchmarks/matmul_relu.py lines 26–58

    a_shapes = [
        [2048, 768],
        [64, 1280],
        [2048, 768],
        [32, 2048],
        [1, 39200],
        [128, 3072],
        [16, 1280],
    ]
    b_shapes = [
        [768, 3072],
        [1280, 1000],
        [768, 768],
        [2048, 1000],
        [39200, 50],
        [3072, 1000],
        [1280, 1000],
    ]

    # Artificial larger shapes
    a_shapes += [[10240, 512], [10240, 1024]]
    b_shapes += [[512, 10240], [1024, 10240]]

    for i in range(len(a_shapes)):
        a_shape = a_shapes[i]
        b_shape = b_shapes[i]
        print("Shape:", a_shape, "x", b_shape)
        a = torch.randn(a_shape, device="cuda", dtype=torch.float16)
        b = torch.randn(b_shape, device="cuda", dtype=a.dtype)

        time_with_torch_timer(torch_mm, (a, b), string_id="torch mm")
        time_with_torch_timer(torch_mm_relu, (a, b), string_id="torch mm + relu")
        time_with_torch_timer(inductor_mm, (a, b), string_id="inductor mm")