custom_operator.h

#pragma once

#include <torch/csrc/jit/operator.h>
#include <ATen/core/stack.h>
#include <torch/csrc/jit/tracer.h>
#include <torch/csrc/utils/variadic.h>

#include <ATen/core/function_schema.h>
#include <c10/util/Metaprogramming.h>
#include <c10/util/TypeList.h>

namespace torch {
namespace jit {
namespace detail {

using ::c10::Argument;
using ::c10::FunctionSchema;

template <typename T>
void checkStaticTypes() {
  // Give nice error messages for some of the common error cases.
  // Use a LOUD ERROR MESSAGE SO USERS SEE THE STATIC_ASSERT
  static_assert(
      !std::is_integral<T>::value || std::is_same<T, int64_t>::value,
      "INVALID TYPE: Only int64_t is supported as an integral argument type");
  static_assert(
      !std::is_same<T, float>::value,
      "INVALID TYPE: float is not supported as an argument type, use double instead");
}

template <typename First, typename Second, typename... Rest>
void checkStaticTypes() {
  checkStaticTypes<First>();
  checkStaticTypes<Second, Rest...>();
}

template <typename... Ts, size_t... Is>
::std::vector<Argument> createArgumentVectorFromTypes(Indices<Is...> indices) {
  checkStaticTypes<decay_t<Ts>...>();
  // Arguments are named "_<index>"
  return {Argument("_" + std::to_string(Is), getTypePtr<decay_t<Ts>>())...};
}

template <typename... Ts, size_t... Is>
::std::vector<Argument> createReturns(Indices<Is...> indices) {
  return createArgumentVectorFromTypes<Ts..., Is...>();
}

template <typename... Ts>
::std::vector<Argument> createReturns(std::tuple<Ts...>* tuple) {
  // Create an index pack so we can call `get<Indices>` on the tuple next.
  return createReturns<Ts...>(typename MakeIndices<sizeof...(Ts)>::indices{});
}

template <typename ReturnType>
::std::vector<Argument> createReturns(ReturnType*) {
  checkStaticTypes<decay_t<ReturnType>>();
  return {Argument("_1", getTypePtr<decay_t<ReturnType>>())};
}

template <typename FunctionTraits, size_t... Is>
::std::vector<Argument> createArgumentVectorFromTraits(Indices<Is...> indices) {
  using ArgumentTypes = typename FunctionTraits::parameter_types;
  return createArgumentVectorFromTypes<
      c10::guts::typelist::element_t<Is, ArgumentTypes>...>(indices);
}

template <typename FunctionTraits>
FunctionSchema createFunctionSchemaFromTraits(const std::string& name) {
  using ReturnType = typename FunctionTraits::return_type;

  auto arguments = createArgumentVectorFromTraits<FunctionTraits>(
      typename MakeIndices<FunctionTraits::number_of_parameters>::indices{});
  auto returns = createReturns(static_cast<ReturnType*>(nullptr));

  return {name, "", arguments, returns};
}

template <size_t... Is, typename... Types>
Node* getTracedNode(
    const FunctionSchema& schema,
    const std::tuple<Types...>& tuple) {
  auto symbol = Symbol::fromQualString(schema.name());
  const auto& graph = tracer::getTracingState()->graph;
  Node* node = graph->create(symbol, /*num_outputs=*/0);
  tracer::recordSourceLocation(node);

  // Hack to call addInputs for the parameter pack in a sequenced fashion.
  // https://stackoverflow.com/questions/12030538/calling-a-function-for-each-variadic-template-argument-and-an-array
  int _[] = {
      (tracer::addInputs(
           node, schema.arguments()[Is].name().c_str(), std::get<Is>(tuple)),
       0)...};
  (void)_; // ignore

  graph->insertNode(node);

  return node;
}

template <typename Implementation, typename... Types, size_t... Is>
void callOperatorWithTuple(
    const FunctionSchema& schema,
    Implementation&& implementation,
    Stack& stack,
    std::tuple<Types...>& arguments,
    Indices<Is...>) {
  AT_ASSERT(stack.size() == sizeof...(Is));

  // Pop values from the stack into the elements of the tuple.
  pop(stack, std::get<Is>(arguments)...);

  Node* node = nullptr;
  if (jit::tracer::isTracing()) {
    node = getTracedNode<Is...>(schema, arguments);
  }

  // Call into the actual, original, user-supplied function.
  auto return_value =
      std::forward<Implementation>(implementation)(std::get<Is>(arguments)...);

  if (jit::tracer::isTracing()) {
    jit::tracer::addOutput(node, return_value);
  }

  // Push the return value back onto the stack.
  push(stack, IValue(std::move(return_value)));
}

inline void checkArgumentVector(
    const char* what,
    const std::vector<Argument>& inferred,
    const std::vector<Argument>& provided,
    const FunctionSchema& inferredSchema,
    const FunctionSchema& providedSchema) {
  // clang-format off
  AT_CHECK(
      inferred.size() == provided.size(),
      "Inferred ", inferred.size(), " ", what,
      "(s) for operator implementation, but the provided schema specified ",
      provided.size(), " ", what, "(s). Inferred schema: ", inferredSchema,
      " | Provided schema: ", providedSchema);
  // clang-format on
  for (size_t i = 0; i < provided.size(); ++i) {
    // clang-format off
    AT_CHECK(
        provided[i].type()->isSubtypeOf(inferred[i].type()),
        "Inferred type for ", what, " #", i, " was ", *inferred[i].type(),
        ", but the provided schema specified type ", *provided[i].type(),
        " for the ", what, " in that position. Inferred schema: ",
        inferredSchema, " | Provided schema: ", providedSchema);
    // clang-format on
  }
}

template <typename Traits>
FunctionSchema inferAndCheckSchema(const std::string& schemaOrName) {
  // If there is no '(' in the schema, we assume this is only the name (e.g.
  // "foo::bar").
  const auto bracketIndex = schemaOrName.find('(');
  if (bracketIndex == std::string::npos) {
    // Infer the full schema and we're good.
    return torch::jit::detail::createFunctionSchemaFromTraits<Traits>(
        /*name=*/schemaOrName);
  }

  // If the user provided her own schema, we need to infer it nevertheless and
  // check that it's correct. We return the user provided schema in the end
  // because it has proper argument names.

  auto providedSchema = parseSchema(schemaOrName);

  const auto inferredSchema =
      torch::jit::detail::createFunctionSchemaFromTraits<Traits>(
          providedSchema.name());
  checkArgumentVector(
      "argument",
      inferredSchema.arguments(),
      providedSchema.arguments(),
      inferredSchema,
      providedSchema);
  checkArgumentVector(
      "return value",
      inferredSchema.returns(),
      providedSchema.returns(),
      inferredSchema,
      providedSchema);
  return providedSchema;
}
} // namespace detail

template <typename Implementation>
Operator createOperator(
    const std::string& schemaOrName,
    Implementation&& implementation) {
  using Traits = c10::guts::infer_function_traits_t<Implementation>;
  using ArgumentTypes =
      c10::guts::typelist::map_t<decay_t, typename Traits::parameter_types>;
  using ArgumentTuple =
      typename c10::guts::typelist::to_tuple<ArgumentTypes>::type;
  static constexpr auto kNumberOfArguments =
      std::tuple_size<ArgumentTuple>::value;

  auto schema = torch::jit::detail::inferAndCheckSchema<Traits>(schemaOrName);

  return Operator(schema, [implementation, schema](Stack& stack) {
    ArgumentTuple tuple;
    torch::jit::detail::callOperatorWithTuple(
        schema,
        std::move(implementation), // NOLINT(bugprone-move-forwarding-reference)
        stack,
        tuple,
        typename MakeIndices<kNumberOfArguments>::indices{});
    return 0;
  });
}

struct TORCH_API RegisterOperators {
  RegisterOperators() = default;

  RegisterOperators(std::vector<Operator> operators) {
    for (Operator& o : operators) {
      registerOperator(std::move(o));
    }
  }

  template <typename Implementation>
  RegisterOperators(const std::string& name, Implementation&& implementation) {
    op(name, std::forward<Implementation>(implementation));
  }

  template <typename Implementation>
  RegisterOperators& op(
      const std::string& name,
      Implementation&& implementation) {
    registerOperator(
        createOperator(name, std::forward<Implementation>(implementation)));
    return *this;
  }
};

} // namespace jit
} // namespace torch