custom_operator.h

#pragma once

#include <torch/csrc/jit/operator.h>
#include <ATen/core/stack.h>
#include <ATen/core/op_registration/op_registration.h>

namespace torch {
namespace jit {

/// Registration class for new operators. Effectively calls
/// `torch::jit::registerOperator` for every supplied operator, but allows doing
/// so in the global scope when a `RegisterOperators` object is assigned to a
/// static variable. Also handles registration of user-defined, "custom"
/// operators.
struct TORCH_API RegisterOperators {
  RegisterOperators() = default;

  /// Registers a vector of already created `Operator`s.
  RegisterOperators(std::vector<Operator> operators) {
    for (Operator& o : operators) {
      registerOperator(std::move(o));
    }
  }

  /// Calls `op(...)` with the given operator name and implementation.
  template <typename Implementation>
  C10_DEPRECATED_MESSAGE("torch::jit::RegisterOperators is deprecated. Please use torch::RegisterOperators instead.")
  RegisterOperators(const std::string& name, Implementation&& implementation) {
    op_(name, std::forward<Implementation>(implementation));
  }

  template <typename Implementation>
  C10_DEPRECATED_MESSAGE("torch::jit::RegisterOperators is deprecated. Please use torch::RegisterOperators instead.")
  RegisterOperators& op(
      const std::string& name,
      Implementation&& implementation) {
    op_(name, std::forward<Implementation>(implementation));

    return *this;
  }

private:

  template <typename Implementation>
  void op_(const std::string& name, Implementation&& implementation) {
    registrars_.emplace_back(std::make_shared<c10::RegisterOperators>(name, std::forward<Implementation>(implementation)));
  }

  // A c10::RegisterOperators instance is not copyable, so to make
  // torch::jit::RegisterOperators copyable, we use shared_ptrs.
  // We need to keep the c10::RegisterOperators instances around
  // because this is an RAII pattern. In the destructor, the registered
  // ops get de-registered.
  std::vector<std::shared_ptr<c10::RegisterOperators>> registrars_;
};

} // namespace jit

} // namespace torch