diff --git a/python/examples/pytorch_bench.py b/python/examples/pytorch_bench.py new file mode 100644 index 0000000..f62a7a5 --- /dev/null +++ b/python/examples/pytorch_bench.py @@ -0,0 +1,58 @@ +import sys + +import cuda.nvbench as nvbench +import torch + + +def as_torch_cuda_Stream( + cs: nvbench.CudaStream, dev: int | None +) -> torch.cuda.ExternalStream: + return torch.cuda.ExternalStream( + stream_ptr=cs.addressof(), device=torch.cuda.device(dev) + ) + + +def torch_bench(state: nvbench.State) -> None: + state.set_throttle_threshold(0.25) + + dev_id = state.get_device() + tc_s = as_torch_cuda_Stream(state.get_stream(), dev_id) + + dt = torch.float32 + scalar_shape: tuple = tuple() + n = 2**28 + with torch.cuda.stream(tc_s): + a3 = torch.randn(scalar_shape, dtype=dt) + a2 = torch.randn(scalar_shape, dtype=dt) + a1 = torch.randn(scalar_shape, dtype=dt) + a0 = torch.randn(scalar_shape, dtype=dt) + x = torch.linspace(-3, 3, n, dtype=dt) + y = torch.sin(x) + + learning_rate = 1e-4 + + def launcher(launch: nvbench.Launch) -> None: + tc_s = as_torch_cuda_Stream(launch.get_stream(), dev_id) + with torch.cuda.stream(tc_s): + x2 = torch.square(x) + y_pred = (a3 + x2 * a1) + x * (a2 + a0 * x2) + + _ = torch.square(y_pred - y).sum() + grad_y_pred = 2 * (y_pred - y) + grad_a3 = grad_y_pred.sum() + grad_a2 = (grad_y_pred * x).sum() + grad_a1 = (grad_y_pred * x2).sum() + grad_a0 = (grad_y_pred * x2 * x).sum() + + _ = a3 - grad_a3 * learning_rate + _ = a2 - grad_a2 * learning_rate + _ = a1 - grad_a1 * learning_rate + _ = a0 - grad_a0 * learning_rate + + state.exec(launcher, sync=True) + + +if __name__ == "__main__": + nvbench.register(torch_bench) + + nvbench.run_all_benchmarks(sys.argv)