mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-07 12:21:27 +01:00
72b0447f8d
Summary:
This PR moves tracing logic out of the generated VariableType kernels, to associate it with a new dedicated dispatch key Tracer.
It also toggles the dispatch key set at various places to keep the semantics unchanged - see the inline [Tracing Mode Switches] note.
Sample generated code:
```
Tensor & __ilshift___Tensor(Tensor & self, const Tensor & other) {
#if !defined(PYTORCH_DISABLE_TRACING)
torch::jit::Node* node = nullptr;
std::shared_ptr<jit::tracer::TracingState> tracer_state;
if (jit::tracer::isTracing()) {
tracer_state = jit::tracer::getTracingState();
at::Symbol op_name;
op_name = jit::Symbol::fromQualString("aten::__ilshift__");
node = tracer_state->graph->create(op_name, /*num_outputs=*/0);
jit::tracer::recordSourceLocation(node);
jit::tracer::addInputs(node, "self", self);
jit::tracer::addInputs(node, "other", other);
tracer_state->graph->insertNode(node);
jit::tracer::setTracingState(nullptr);
}
#endif
static auto op = c10::Dispatcher::singleton().findSchemaOrThrow("aten::__ilshift__", "Tensor");
c10::Dispatcher::singleton().redispatch<Tensor &, Tensor &, const Tensor &>(op, c10::DispatchKey::Tracer, self, other);
#if !defined(PYTORCH_DISABLE_TRACING)
if (tracer_state) {
jit::tracer::setTracingState(std::move(tracer_state));
jit::tracer::addOutput(node, self);
}
#endif
return self;
}
```
Pull Request resolved: https://github.com/pytorch/pytorch/pull/38467
ghstack-source-id: 105215150
Test Plan: CI
Differential Revision: D21570684
fbshipit-source-id: 1a96761830307f9a934f38bfb9fe8b5b1763e0e0
69 lines
1.7 KiB
C++
69 lines
1.7 KiB
C++
#include <c10/core/DispatchKey.h>
|
|
|
|
namespace c10 {
|
|
|
|
const char* toString(DispatchKey t) {
|
|
switch (t) {
|
|
case DispatchKey::Undefined:
|
|
return "Undefined";
|
|
case DispatchKey::CPU:
|
|
return "CPU";
|
|
case DispatchKey::CUDA:
|
|
return "CUDA";
|
|
case DispatchKey::SparseCPU:
|
|
return "SparseCPU";
|
|
case DispatchKey::SparseCUDA:
|
|
return "SparseCUDA";
|
|
case DispatchKey::MKLDNN:
|
|
return "MKLDNN";
|
|
case DispatchKey::OpenGL:
|
|
return "OpenGL";
|
|
case DispatchKey::OpenCL:
|
|
return "OpenCL";
|
|
case DispatchKey::IDEEP:
|
|
return "IDEEP";
|
|
case DispatchKey::HIP:
|
|
return "HIP";
|
|
case DispatchKey::SparseHIP:
|
|
return "SparseHIP";
|
|
case DispatchKey::FPGA:
|
|
return "FPGA";
|
|
case DispatchKey::MSNPU:
|
|
return "MSNPU";
|
|
case DispatchKey::XLA:
|
|
return "XLA";
|
|
case DispatchKey::Vulkan:
|
|
return "Vulkan";
|
|
case DispatchKey::MkldnnCPU:
|
|
return "MkldnnCPU";
|
|
case DispatchKey::QuantizedCPU:
|
|
return "QuantizedCPU";
|
|
case DispatchKey::Autograd:
|
|
return "Autograd";
|
|
case DispatchKey::BackendSelect:
|
|
return "BackendSelect";
|
|
case DispatchKey::Batched:
|
|
return "Batched";
|
|
case DispatchKey::TESTING_ONLY_GenericMode:
|
|
return "TESTING_ONLY_GenericMode";
|
|
case DispatchKey::Autocast:
|
|
return "Autocast";
|
|
case DispatchKey::TESTING_ONLY_GenericWrapper:
|
|
return "TESTING_ONLY_GenericWrapper";
|
|
case DispatchKey::Profiler:
|
|
return "Profiler";
|
|
case DispatchKey::Named:
|
|
return "Named";
|
|
case DispatchKey::Tracer:
|
|
return "Tracer";
|
|
default:
|
|
return "UNKNOWN_TENSOR_TYPE_ID";
|
|
}
|
|
}
|
|
|
|
std::ostream& operator<<(std::ostream& str, DispatchKey rhs) {
|
|
return str << toString(rhs);
|
|
}
|
|
|
|
} // namespace c10
|