FM / core

/usr/local/lib64/python3.6/site-packages/torch/include/caffe2/core/

/usr/local/lib64/python3.6/site-packages/torch/include/caffe2/core

Name	Size	Mode	Actions
allocator.h	136	0644	edit dl rm
blob.h	4168	0644	edit dl rm
blob_serialization.h	10791	0644	edit dl rm
blob_serializer_base.h	3905	0644	edit dl rm
blob_stats.h	1127	0644	edit dl rm
common.h	4329	0644	edit dl rm
common_cudnn.h	9893	0644	edit dl rm
common_gpu.h	21414	0644	edit dl rm
common_omp.h	156	0644	edit dl rm
context.h	6174	0644	edit dl rm
context_base.h	4382	0644	edit dl rm
context_gpu.h	11014	0644	edit dl rm
cudnn_wrappers.h	6956	0644	edit dl rm
db.h	9352	0644	edit dl rm
distributions_stubs.h	2161	0644	edit dl rm
event.h	12420	0644	edit dl rm
event_cpu.h	1192	0644	edit dl rm
export_c10_op_to_caffe2.h	9487	0644	edit dl rm
export_caffe2_op_to_c10.h	11101	0644	edit dl rm
flags.h	74	0644	edit dl rm
graph.h	5258	0644	edit dl rm
init.h	6496	0644	edit dl rm
logging.h	75	0644	edit dl rm
macros.h	3426	0644	edit dl rm
memonger.h	817	0644	edit dl rm
module.h	2473	0644	edit dl rm
net.h	4634	0644	edit dl rm
net_async_base.h	7397	0644	edit dl rm
net_async_scheduling.h	993	0644	edit dl rm
net_async_task.h	833	0644	edit dl rm
net_async_task_future.h	1925	0644	edit dl rm
net_async_task_graph.h	2253	0644	edit dl rm
net_async_tracing.h	5093	0644	edit dl rm
net_dag_utils.h	2146	0644	edit dl rm
net_parallel.h	2144	0644	edit dl rm
net_simple.h	2606	0644	edit dl rm
net_simple_refcount.h	2097	0644	edit dl rm
numa.h	72	0644	edit dl rm
observer.h	3809	0644	edit dl rm
operator.h	58872	0644	edit dl rm
operator_gradient.h	10222	0644	edit dl rm
operator_schema.h	18477	0644	edit dl rm
plan_executor.h	219	0644	edit dl rm
prof_dag_counters.h	2751	0644	edit dl rm
qtensor.h	6615	0644	edit dl rm
qtensor_serialization.h	2624	0644	edit dl rm
scope_guard.h	4675	0644	edit dl rm
static_tracepoint.h	398	0644	edit dl rm
static_tracepoint_elfx86.h	5555	0644	edit dl rm
stats.h	10365	0644	edit dl rm
storage.h	733	0644	edit dl rm
tensor.h	18668	0644	edit dl rm
tensor_impl.h	351	0644	edit dl rm
tensor_int8.h	450	0644	edit dl rm
test_utils.h	6285	0644	edit dl rm
timer.h	1218	0644	edit dl rm
transform.h	5741	0644	edit dl rm
types.h	2248	0644	edit dl rm
workspace.h	11305	0644	edit dl rm

Edit: /usr/local/lib64/python3.6/site-packages/torch/include/caffe2/core/export_c10_op_to_caffe2.h (9487B)

#pragma once

#include <c10/macros/Macros.h>
#include <c10/util/Registry.h>
#include "caffe2/core/operator.h"

// TODO Also register c10 operators on mobile
#if !defined(CAFFE2_IS_XPLAT_BUILD) && !defined(C10_MOBILE)
#include <ATen/core/dispatch/Dispatcher.h>
#include <ATen/core/ivalue.h>
#include <c10/util/ArrayRef.h>
#include <c10/util/C++17.h>
#include <c10/util/Metaprogramming.h>
#include "caffe2/core/export_caffe2_op_to_c10.h"

namespace caffe2 {

/**
 * To make a c10 operator "C10Add" callable from caffe2 as "C2MyAddOpName", just
 * write
 *
 *  To export the CPU kernel
 *     C10_EXPORT_C10_OP_TO_CAFFE2_CPU(C10Add, C2MyAddOp)
 *
 *  To export the CUDA kernel
 *     C10_EXPORT_C10_OP_TO_CAFFE2_CUDA(C10Add, C2MyAddOp)
 *
 */

namespace detail {
template <class Context>
class C10OperatorWrapper final : public Operator<Context> {
 public:
  USE_OPERATOR_CONTEXT_FUNCTIONS;

C10OperatorWrapper(
      const c10::OperatorHandle& op,
      const OperatorDef& operator_def,
      Workspace* ws)
      : Operator<Context>(operator_def, ws),
        op_(op),
        has_preallocated_outputs_(
            op_.schema().arguments().size() != 0 &&
            op_.schema().arguments().back().name() ==
                detail::PREALLOCATED_OUTPUT_ARGNAME) {
    AT_ASSERT(
        !has_preallocated_outputs_ ||
        op_.schema().arguments().back().type()->isSubtypeOf(
            OptionalType::create(ListType::ofTensors())));

// NOLINTNEXTLINE(clang-diagnostic-sign-compare)
    AT_ASSERT(operator_def.output_size() == op_.schema().returns().size());
    AT_ASSERT(
        operator_def.input_size() + (has_preallocated_outputs_ ? 1 : 0) <=
        op_.schema()
            .arguments()
            .size()); // '<=' because there might be caffe2 nontensor arguments
  }

bool RunOnDevice() override {
    // due to caching the stack_, concurrent calling is not allowed.
    // TODO thread_local might fix this
    std::lock_guard<std::mutex> lock(mutex_);

pushInputs_();
    callKernel_();
    popOutputs_();

return true;
  }

private:
  void pushInputs_() {
    AT_ASSERT(stack_.size() == 0);
    stack_.reserve(
        op_.schema().arguments().size() + (has_preallocated_outputs_ ? 1 : 0));

size_t input_tensor_index = 0;

for (const auto& argument : op_.schema().arguments()) {
      if (argument.name() == detail::PREALLOCATED_OUTPUT_ARGNAME) {
        // note: if detail::PREALLOCATED_OUTPUT_ARGNAME was at the end of the
        // argument list, then has_preallocated_outputs_ would be true.
        AT_ASSERTM(
            has_preallocated_outputs_,
            "Error in caffe2->c10 wrapper: Operator schema has a parameter named ",
            detail::PREALLOCATED_OUTPUT_ARGNAME,
            ", but it's not at the end of the argument list");

AT_ASSERTM(
            argument.type()->isSubtypeOf(
                OptionalType::create(ListType::ofTensors())),
            "Error in caffe2->c10 wrapper: Operator schema has a parameter named ",
            detail::PREALLOCATED_OUTPUT_ARGNAME,
            ", but it's not of type TensorList?");
        stack_.emplace_back(preallocated_outputs_());

} else if (argument.type()->isSubtypeOf(TensorType::get())) {
        AT_ASSERTM(
            // NOLINTNEXTLINE(clang-diagnostic-sign-compare)
            input_tensor_index < InputSize(),
            "Error in caffe2->c10 wrapper: Too few tensor arguments given (",
            InputSize(),
            "), operator schema expected more.");
        stack_.emplace_back(at::Tensor(Input(input_tensor_index++)));
      } else if (argument.type()->isSubtypeOf(OptionalType::ofTensor())) {
        if (input_tensor_index < InputSize()) {
          stack_.emplace_back(at::Tensor(Input(input_tensor_index++)));
        } else {
          stack_.emplace_back(IValue());
        }
      } else if (argument.type()->isSubtypeOf(ListType::ofTensors())) {
        AT_ASSERTM(
            input_tensor_index == 0,
            "Error in caffe2->c10 wrapper: Schema can only have either one or more Tensor inputs or one TensorList input.");
        stack_.emplace_back(array_inputs_());
        input_tensor_index = InputSize();

} else {
        stack_.emplace_back(get_nontensor_argument_(argument));
      }
    }
    AT_ASSERTM(
        input_tensor_index == InputSize(),
        "Error in caffe2->c10 wrapper: Number of caffe2 operator inputs (",
        InputSize(),
        ") doesn't match number of tensor arguments (",
        input_tensor_index,
        ") in the c10 operator schema.");
  }

void callKernel_() {
    AT_ASSERT(stack_.size() == op_.schema().arguments().size());
    op_.callBoxed(&stack_);
  }

void popOutputs_() {
    AT_ASSERT(stack_.size() == op_.schema().returns().size());
    for (size_t i = 0; i < op_.schema().returns().size(); ++i) {
      OperatorBase::SetOutputTensor(i, Tensor(std::move(stack_[i]).toTensor()));
    }
    stack_.clear();
  }

c10::List<at::Tensor> array_inputs_() {
    c10::List<at::Tensor> result;
    result.reserve(InputSize());
    // NOLINTNEXTLINE(clang-diagnostic-sign-compare)
    for (size_t i = 0; i < InputSize(); ++i) {
      result.emplace_back(Input(i));
    }
    return result;
  }

c10::List<at::Tensor> preallocated_outputs_() {
    c10::List<at::Tensor> result;
    result.reserve(OutputSize());
    // NOLINTNEXTLINE(clang-diagnostic-sign-compare)
    for (size_t i = 0; i < OutputSize(); ++i) {
      result.emplace_back(OperatorBase::OutputTensorOrUndefined(i));
    }
    return result;
  }

IValue get_nontensor_argument_(const c10::Argument& argument) {
    if (argument.type()->isSubtypeOf(IntType::get())) {
      return get_nontensor_argument_<int>(
          argument.name(), argument.default_value());
    } else if (argument.type()->isSubtypeOf(FloatType::get())) {
      return get_nontensor_argument_<double>(
          argument.name(), argument.default_value());
    } else if (argument.type()->isSubtypeOf(BoolType::get())) {
      return get_nontensor_argument_<bool>(
          argument.name(), argument.default_value());
    } else {
      // TODO Support more types
      AT_ERROR(
          "Error in caffe2->c10 wrapper: Unsupported argument type ",
          argument.type()->str(),
          " in c10 operator schema");
    }
  }

template <class T>
  IValue get_nontensor_argument_(
      const std::string& name,
      const c10::optional<IValue>& default_value) {
    if (default_value.has_value()) {
      return this->template GetSingleArgument<T>(name, default_value->to<T>());
    } else {
      TORCH_CHECK(
          this->template HasSingleArgumentOfType<T>(name),
          "Error in caffe2->c10 wrapper: Expected argument '",
          name,
          "' missing or wrong type.");
      return this->template GetSingleArgument<T>(name, 0);
    }
  }

c10::OperatorHandle op_;

// has_preallocated_outputs_ is true iff the operator schema has a last
  // argument that is a TensorList and has a name equal to with the name equal
  // to detail::PREALLOCATED_OUTPUT_ARGNAME. This argument is then used to pass
  // in preallocated output tensors to the caffe2 operator.
  bool has_preallocated_outputs_;

// this is stored as a member here to avoid having to re-allocate a stack
  // for each call. Between kernel calls, stack_.size() == 0, but capacity
  // should not need to be grown anymore after the first call.
  std::vector<IValue> stack_;
  std::mutex mutex_;
};

template <class Context>
inline std::function<
    std::unique_ptr<OperatorBase>(const OperatorDef&, Workspace*)>
createC10OperatorWrapper(const c10::OperatorName& op_name) {
  return [op_name](const OperatorDef& op_def, Workspace* ws) {
    auto op_handle =
        c10::Dispatcher::singleton().findSchema(op_name);
    AT_ASSERTM(
        op_handle.has_value(),
        "Tried to register c10 operator ",
        op_name.name,
        ".",
        op_name.overload_name,
        " with caffe2, but didn't find the c10 operator.");
    return std::make_unique<C10OperatorWrapper<Context>>(
        *op_handle, op_def, ws);
  };
}

} // namespace detail
} // namespace caffe2

#define C10_EXPORT_C10_OP_TO_CAFFE2_CPU(                       \
    OperatorName, Name)                                        \
  REGISTER_CPU_OPERATOR_CREATOR(                               \
      Name,                                                    \
      ::caffe2::detail::createC10OperatorWrapper<CPUContext>(  \
          {OperatorName, ""}))
#define C10_EXPORT_C10_OP_TO_CAFFE2_CUDA(                      \
    OperatorName, Name)                                        \
  REGISTER_CUDA_OPERATOR_CREATOR(                              \
      Name,                                                    \
      ::caffe2::detail::createC10OperatorWrapper<CUDAContext>( \
          {OperatorName, ""}))
#define C10_EXPORT_C10_OP_TO_CAFFE2_HIP(                       \
    OperatorName, Name)                                        \
  REGISTER_HIP_OPERATOR_CREATOR(                               \
      Name,                                                    \
      ::caffe2::detail::createC10OperatorWrapper<HIPContext>(  \
          {OperatorName, ""}))
#else
#define C10_EXPORT_C10_OP_TO_CAFFE2_CPU(                       \
    OperatorName, Name)
#define C10_EXPORT_C10_OP_TO_CAFFE2_CUDA(                      \
    OperatorName, Name)
#define C10_EXPORT_C10_OP_TO_CAFFE2_HIP(                       \
    OperatorName, Name)
#endif

cmd: