mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-07 12:21:27 +01:00
0b81f700aa
Summary: We want to add compile IDs and frames to each Torch-Compiled Region in order to help users cross reference the section they are checking alongside data obtained from tools, such as tlparse. This diff operates on the assumption that each graph section will enter and exit a CompileContext before it is ran to either compile the graph or look it up in the cache. Based on this assuption, we can save the value of the graph section from the exited CompileContext in eval_frame.c using a Python C API. After this, we can create a new interface in cpp shim to wrap around the record_function in order to pass in the new keyword argument for "context". Test Plan: Enhance test_profiler_dynamo_compiled_region to look for kwinputs as well as a name to see that the context is now labeled. Also changed test to run graph with more contexts so that we test a wider range of profiling. Differential Revision: D60803317 Pull Request resolved: https://github.com/pytorch/pytorch/pull/132765 Approved by: https://github.com/anijain2305
40 lines
1.2 KiB
C++
40 lines
1.2 KiB
C++
#include <ATen/record_function.h>
|
|
#include <torch/csrc/dynamo/cpp_shim.h>
|
|
|
|
struct _PytorchRecordFunctionState {
|
|
at::RecordFunction guard;
|
|
|
|
_PytorchRecordFunctionState() : guard(at::RecordScope::FUNCTION) {}
|
|
};
|
|
|
|
_PytorchRecordFunctionState* _pytorch_record_function_enter(const char* name) {
|
|
_PytorchRecordFunctionState* state = new _PytorchRecordFunctionState();
|
|
state->guard.before(name);
|
|
return state;
|
|
}
|
|
|
|
static inline _PytorchRecordFunctionState*
|
|
_pytorch_record_function_enter_with_kwinputs(
|
|
const char* name,
|
|
const std::unordered_map<std::string, c10::IValue>* kwargs) {
|
|
_PytorchRecordFunctionState* state = new _PytorchRecordFunctionState();
|
|
std::vector<c10::IValue> args;
|
|
state->guard.before(name, &args, kwargs);
|
|
return state;
|
|
}
|
|
|
|
_PytorchRecordFunctionState* _pytorch_record_function_enter_with_context(
|
|
const char* name,
|
|
const char* context) {
|
|
auto map = std::unordered_map<std::string, c10::IValue>();
|
|
map.insert({"context", c10::IValue(context)});
|
|
return _pytorch_record_function_enter_with_kwinputs(name, &map);
|
|
}
|
|
|
|
void _pytorch_record_function_exit(_PytorchRecordFunctionState* state) {
|
|
if (state == nullptr) {
|
|
return;
|
|
}
|
|
delete state;
|
|
}
|