()
| 89 | |
| 90 | |
| 91 | def main() -> None: |
| 92 | if not hasattr(torch.Tensor, "__dlpack_c_exchange_api__"): |
| 93 | raise SystemExit("torch.Tensor.__dlpack_c_exchange_api__ not available") |
| 94 | |
| 95 | repeat = 10000 |
| 96 | invoke_n = _load_invoke_n() |
| 97 | a = torch.zeros(1, device="cuda:0") |
| 98 | b = torch.zeros(1, device="cuda:0") |
| 99 | c = torch.zeros(1, device="cuda:0") |
| 100 | |
| 101 | print("---------------------------------------------------") |
| 102 | print("Benchmark C++ -> Python callback with 3 torch.Tensor args") |
| 103 | print('Arguments: 3 x torch.zeros(1, device="cuda:0")') |
| 104 | print("---------------------------------------------------") |
| 105 | bench_pycallback_tensor_cls_torch(invoke_n, a, b, c, repeat) |
| 106 | bench_pycallback_from_dlpack(invoke_n, a, b, c, repeat) |
| 107 | print("---------------------------------------------------") |
| 108 | |
| 109 | |
| 110 | if __name__ == "__main__": |
no test coverage detected