xref: /aosp_15_r20/external/pytorch/aten/src/ATen/cuda/tunable/StreamTimer.h (revision da0073e96a02ea20f0ac840b70461e3646d07c45)
1 // Original TunableOp is from onnxruntime.
2 // https://github.com/microsoft/onnxruntime/blob/main/onnxruntime/core/framework/tunable.h
3 // https://github.com/microsoft/onnxruntime/tree/main/onnxruntime/core/providers/rocm/tunable
4 // Copyright (c) Microsoft Corporation.
5 // Licensed under the MIT license.
6 //
7 // Adapting TunableOp into PyTorch
8 // Copyright (c) Advanced Micro Devices, Inc.
9 //
10 #pragma once
11 
12 #include <cuda_runtime.h>
13 
14 #include <ATen/cuda/tunable/Tunable.h>
15 
16 namespace at::cuda::tunable {
17 
18 class StreamTimer : public ITimer {
19   public:
20     StreamTimer();
21     virtual ~StreamTimer() override;
22 
23     void Start() override;
24 
25     void End() override;
26 
27     float Duration() override;
28 
29   private:
30     cudaEvent_t start_;
31     cudaEvent_t end_;
32 };
33 
34 } // namespace at::cuda::tunable
35