pytorch - 💡(How to fix) Fix Improve the reinplace FX pass to handle in-place mutations on views of graph inputs [1 pull requests]

def store(buf, indices, values):
    buf.view(-1)[indices] = values
    return buf

---

%view = reshape(%arg0_1, [-1])
%index_put = index_put(%view, [%arg1_1], %arg2_1)
%view_1 = reshape(%index_put, [1024, 1024])
%copy_ = copy_(%arg0_1, %view_1)
return (copy_,)

---

%view = reshape(%arg0_1, [-1])
%index_put = index_put_(%view, [%arg1_1], %arg2_1)
return (%index_put,)

---

Buffer: 128 MiB, Scatter: 33.0 KiB
  Eager:   0.0042 ms
  Compile: 0.6744 ms
  Ratio:   161.21x

---

import torch


def store_via_view(buf: torch.Tensor, indices: torch.Tensor, values: torch.Tensor):
    buf_flat = buf.view(-1)
    buf_flat[indices] = values
    return buf


def time_cuda(fn, warmup=20, iters=200):
    for _ in range(warmup):
        fn()
    torch.cuda.synchronize()
    start = torch.cuda.Event(enable_timing=True)
    end = torch.cuda.Event(enable_timing=True)
    start.record()
    for _ in range(iters):
        fn()
    end.record()
    end.synchronize()
    return start.elapsed_time(end) / iters


def main():
    device = torch.device("cuda")
    num_tokens = 256
    token_bytes = 132
    scatter_bytes = num_tokens * token_bytes
    buf_size = 128 * 1024 * 1024  # 128 MiB

    buf = torch.zeros(1, buf_size, dtype=torch.uint8, device=device)
    indices = torch.randint(0, buf_size, (scatter_bytes,), dtype=torch.long, device=device)
    values = torch.ones(scatter_bytes, dtype=torch.uint8, device=device)

    # Eager
    eager_ms = time_cuda(lambda: store_via_view(buf, indices, values))

    # Compile
    torch._dynamo.reset()
    compiled = torch.compile(store_via_view, fullgraph=True)
    compiled(buf, indices, values)
    torch.cuda.synchronize()
    compile_ms = time_cuda(lambda: compiled(buf, indices, values))

    print(f"Buffer: 128 MiB, Scatter: {scatter_bytes / 1024:.1f} KiB")
    print(f"  Eager:   {eager_ms:.4f} ms")
    print(f"  Compile: {compile_ms:.4f} ms")
    print(f"  Ratio:   {compile_ms / eager_ms:.2f}x")


if __name__ == "__main__":
    main()