#pragma once
#include "megbrain/imperative/profiler.h"
#include "megbrain/utils/small_vector.h"
#include "../interpreter/stack_manager.h"
#include "../op_trait.h"
namespace mgb::imperative::profiler {
enum class TensorProp {
InvalidProp,
Device,
Shape,
DType,
DevValue,
HostValue,
};
using OpParams = std::unordered_map<std::string, std::string>;
}
namespace mgb::imperative {
template <>
struct ToStringTrait<profiler::TensorProp> {
using TensorProp = profiler::TensorProp;
std::string operator()(TensorProp prop) const {
switch (prop) {
case TensorProp::DType:
return "dtype";
case TensorProp::DevValue:
return "dev_value";
case TensorProp::Device:
return "device";
case TensorProp::HostValue:
return "host_value";
case TensorProp::Shape:
return "shape";
default:
return "unknown";
}
}
};
}
namespace mgb::imperative::profiler {
using Trace = interpreter::intl::StackManager::Trace;
struct ProfileOperatorState;
struct ProfileTensorState;
#define DEF_EVENT(X, ...) struct X##Event __VA_ARGS__;
#define DEF_DUR_EVENT(X, ...) \
struct X##Event __VA_ARGS__; \
struct X##FinishEvent __VA_ARGS__;
DEF_EVENT(OpDispatch, {
uint64_t op_id;
std::string op_name;
std::function<OpParams()> op_params;
SmallVector<uint64_t> inputs;
SmallVector<uint64_t> outputs;
Trace trace;
});
DEF_DUR_EVENT(OpInput, {
uint64_t tensor_id;
TensorShape shape;
});
DEF_DUR_EVENT(OpOutput, {
uint64_t tensor_id;
TensorShape shape;
});
DEF_DUR_EVENT(OpExecute, {
uint64_t op_id;
SmallVector<CompNode> device_list;
std::string reason;
});
DEF_DUR_EVENT(KernelLaunch, {
uint64_t op_id;
uint64_t kernel_id;
CompNode device;
});
DEF_EVENT(TensorDeclare, {
uint64_t tensor_id;
std::string name;
});
DEF_EVENT(TensorProduce, {
uint64_t tensor_id;
TensorLayout layout;
CompNode device;
void* ptr;
});
DEF_EVENT(TensorUsage, { uint64_t tensor_id; });
DEF_EVENT(TensorRelease, { uint64_t tensor_id; });
DEF_EVENT(TensorErase, {
uint64_t tensor_id;
size_t use_count;
});
DEF_EVENT(TensorGetProp, {
uint64_t tensor_id;
TensorProp prop;
});
DEF_EVENT(TensorNotifyProp, {
uint64_t tensor_id;
uint64_t wait_id;
TensorProp prop;
});
DEF_DUR_EVENT(TensorWaitProp, {
uint64_t tensor_id;
uint64_t wait_id;
TensorProp prop;
});
DEF_DUR_EVENT(SampleDevice, {
CompNode device;
size_t total_memory;
size_t free_memory;
});
DEF_EVENT(WorkerException, {});
DEF_EVENT(ShapeInfer, { bool success; });
DEF_DUR_EVENT(Scope, { std::string name; });
DEF_DUR_EVENT(Sync, { Trace trace; });
DEF_DUR_EVENT(StartProfile, { size_t capture_count; });
DEF_DUR_EVENT(StopProfile, { size_t escape_count; });
enum class TensorCommandKind { Put, Del, Drop, ReGen, RecFree, GetValue };
DEF_DUR_EVENT(TensorCommand, {
using Kind = TensorCommandKind;
uint64_t tensor_id;
Kind kind;
});
DEF_DUR_EVENT(AutoEvict, {});
DEF_DUR_EVENT(Custom, {
std::string title;
std::string content;
CompNode device;
});
DEF_EVENT(RecordDevice, { std::shared_ptr<CompNode::Event> event; });
DEF_DUR_EVENT(HostToDevice, {
TensorLayout layout;
CompNode device;
void* host_ptr;
void* device_ptr;
});
#undef DEF_EVENT
#undef DEF_DUR_EVENT
}