mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-07 12:21:27 +01:00
4fcdbb824b
Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/48178 I migrated all call sites that used the legacy dispatcher registration API (RegisterOperators()) to use the new API (TORCH_LIBRARY...). I found all call-sites by running `fbgs RegisterOperators()`. This includes several places, including other OSS code (nestedtensor, torchtext, torchvision). A few things to call out: For simple ops that only had one registered kernel without a dispatch key, I replaced them with: ``` TORCH_LIBRARY_FRAGMENT(ns, m) { m.def("opName", fn_name); } ``` For ops that registered to a specific dispatch key / had multiple kernels registered, I registered the common kernel (math/cpu) directly inside a `TORCH_LIBRARY_FRAGMENT` block, and registered any additional kernels from other files (e.g. cuda) in a separate `TORCH_LIBRARY_IMPL` block. ``` // cpu file TORCH_LIBRARY_FRAGMENT(ns, m) { m.def("opName(schema_inputs) -> schema_outputs"); m.impl("opName", torch::dispatch(c10::DispatchKey::CPU, TORCH_FN(cpu_kernel))); } // cuda file TORCH_LIBRARY_IMPL(ns, CUDA, m) { m.impl("opName", torch::dispatch(c10::DispatchKey::CUDA, TORCH_FN(cuda_kernel))); } ``` Special cases: I found a few ops that used a (legacy) `CPUTensorId`/`CUDATensorId` dispatch key. Updated those to use CPU/CUDA- this seems safe because the keys are aliased to one another in `DispatchKey.h` There were a handful of ops that registered a functor (function class) to the legacy API. As far as I could tell we don't allow this case in the new API, mainly because you can accomplish the same thing more cleanly with lambdas. Rather than delete the class I wrote a wrapper function on top of the class, which I passed to the new API. There were a handful of ops that were registered only to a CUDA dispatch key. I put them inside a TORCH_LIBRARY_FRAGMENT block, and used a `def()` and `impl()` call like in case two above. Test Plan: Imported from OSS Reviewed By: ezyang Differential Revision: D25056090 Pulled By: bdhirsh fbshipit-source-id: 8f868b45f545e5da2f21924046e786850eba70d9
70 lines
2.1 KiB
C++
70 lines
2.1 KiB
C++
#include <torch/script.h>
|
|
|
|
#include "op.h"
|
|
|
|
#include <cstddef>
|
|
#include <string>
|
|
|
|
torch::List<torch::Tensor> custom_op(
|
|
torch::Tensor tensor,
|
|
double scalar,
|
|
int64_t repeat) {
|
|
torch::List<torch::Tensor> output;
|
|
output.reserve(repeat);
|
|
for (int64_t i = 0; i < repeat; ++i) {
|
|
output.push_back(tensor * scalar);
|
|
}
|
|
return output;
|
|
}
|
|
|
|
int64_t custom_op2(std::string s1, std::string s2) {
|
|
return s1.compare(s2);
|
|
}
|
|
|
|
struct CustomOpAutogradFunction : public torch::autograd::Function<CustomOpAutogradFunction> {
|
|
static torch::Tensor forward(
|
|
torch::autograd::AutogradContext* ctx,
|
|
torch::Tensor var1,
|
|
int64_t mul,
|
|
torch::Tensor var2,
|
|
c10::optional<torch::Tensor> var3) {
|
|
ctx->saved_data["mul"] = mul;
|
|
ctx->saved_data["var3_has_value"] = var3.has_value();
|
|
ctx->save_for_backward({var1, var2});
|
|
if (var3) {
|
|
return var1 + mul * var2 + var1 * var2 + var3.value();
|
|
}
|
|
return var1 + mul*var2 + var1*var2;
|
|
}
|
|
|
|
static torch::autograd::variable_list backward(torch::autograd::AutogradContext *ctx, torch::autograd::variable_list grad_output) {
|
|
int mul = ctx->saved_data["mul"].toInt();
|
|
bool var3_has_value = ctx->saved_data["var3_has_value"].toBool();
|
|
auto saved = ctx->get_saved_variables();
|
|
auto var1 = saved[0];
|
|
auto var2 = saved[1];
|
|
auto var3_grad = var3_has_value ? grad_output[0] : torch::Tensor();
|
|
torch::autograd::variable_list output = {
|
|
grad_output[0] + grad_output[0] * var2,
|
|
torch::Tensor(),
|
|
grad_output[0] * mul + grad_output[0] * var1,
|
|
var3_grad};
|
|
return output;
|
|
}
|
|
};
|
|
|
|
torch::Tensor custom_op_with_autograd(
|
|
torch::Tensor var1,
|
|
int64_t mul,
|
|
torch::Tensor var2,
|
|
c10::optional<torch::Tensor> var3) {
|
|
return CustomOpAutogradFunction::apply(var1, mul, var2, var3);
|
|
}
|
|
|
|
TORCH_LIBRARY_FRAGMENT(custom, m) {
|
|
m.def("op", custom_op);
|
|
m.def("op2", custom_op2);
|
|
m.def("op_with_defaults(Tensor tensor, float scalar = 1, int repeat = 1) -> Tensor[]", custom_op);
|
|
m.def("op_with_autograd(Tensor var1, int mul, Tensor var2, Tensor? var3=None) -> Tensor", custom_op_with_autograd);
|
|
}
|