1 #pragma once 2 3 #include <functional> 4 #include <memory> 5 6 #include <c10/core/Device.h> 7 #include <c10/util/strong_type.h> 8 #include <torch/csrc/Export.h> 9 10 struct CUevent_st; 11 12 namespace torch { 13 namespace profiler { 14 namespace impl { 15 16 // ---------------------------------------------------------------------------- 17 // -- Annotation -------------------------------------------------------------- 18 // ---------------------------------------------------------------------------- 19 using ProfilerEventStub = std::shared_ptr<CUevent_st>; 20 using ProfilerVoidEventStub = std::shared_ptr<void>; 21 22 struct TORCH_API ProfilerStubs { 23 virtual void record( 24 c10::DeviceIndex* device, 25 ProfilerVoidEventStub* event, 26 int64_t* cpu_ns) const = 0; 27 virtual float elapsed( 28 const ProfilerVoidEventStub* event, 29 const ProfilerVoidEventStub* event2) const = 0; 30 virtual void mark(const char* name) const = 0; 31 virtual void rangePush(const char* name) const = 0; 32 virtual void rangePop() const = 0; enabledProfilerStubs33 virtual bool enabled() const { 34 return false; 35 } 36 virtual void onEachDevice(std::function<void(int)> op) const = 0; 37 virtual void synchronize() const = 0; 38 virtual ~ProfilerStubs(); 39 }; 40 41 TORCH_API void registerCUDAMethods(ProfilerStubs* stubs); 42 TORCH_API const ProfilerStubs* cudaStubs(); 43 TORCH_API void registerITTMethods(ProfilerStubs* stubs); 44 TORCH_API const ProfilerStubs* ittStubs(); 45 TORCH_API void registerPrivateUse1Methods(ProfilerStubs* stubs); 46 TORCH_API const ProfilerStubs* privateuse1Stubs(); 47 48 using vulkan_id_t = strong::type< 49 int64_t, 50 struct _VulkanID, 51 strong::regular, 52 strong::convertible_to<int64_t>, 53 strong::hashable>; 54 55 } // namespace impl 56 } // namespace profiler 57 } // namespace torch 58