Implement calling of c10 ops from c2 (pytorch#9369)

Summary:
Pull Request resolved: pytorch#9369

This adds the capability for caffe2 to call c10 operators and adds a dummy c10 sigmoid op as a proof of concept.

I used this test script to make sure it works:

    from caffe2.python import workspace, model_helper
    import numpy as np

    data1 = np.random.rand(16, 100).astype(np.float32)
    workspace.FeedBlob("data1", data1)
    m = model_helper.ModelHelper(name="my net")
    sigmoid1 = m.net.C10Sigmoid_DontUseThisOpYet("data1", "sigmoid1")
    sigmoid2 = m.net.Sigmoid("data1", "sigmoid2")

    workspace.RunNetOnce(m.param_init_net)
    workspace.CreateNet(m.net)
    data1 = np.random.rand(16, 100).astype(np.float32)
    workspace.FeedBlob("data1", data1)
    workspace.RunNet(m.name, 1)

    print(workspace.FetchBlob("data1"))
    print(workspace.FetchBlob("sigmoid1"))
    print(workspace.FetchBlob("sigmoid2"))

(and check that both sigmoid outputs are the same)

Reviewed By: ezyang

Differential Revision: D8814669

fbshipit-source-id: eeb0e7a854727f1617a3c592a662a7e5ae226f40

Author: smessmer

caffe2/core/operator.cc

@@ -14,6 +14,8 @@
 #include "caffe2/utils/proto_utils.h"
 #include "caffe2/utils/string_utils.h"
 
+#include "caffe2/core/operator_c10wrapper.h"
+
 CAFFE2_DEFINE_int(
     caffe2_operator_max_engine_name_length,
     10,
@@ -75,8 +77,10 @@ GlobalEnginePrefType& g_global_engine_pref() {
   return *g_global_engine_pref_;
 }
 
-unique_ptr<OperatorBase> TryCreateOperator(
-    const string& key, const OperatorDef& operator_def, Workspace* ws) {
+unique_ptr<OperatorBase> TryCreateC2Operator(
+    const string& key,
+    const OperatorDef& operator_def,
+    Workspace* ws) {
   const auto& type = operator_def.device_option().device_type();
   CAFFE_ENFORCE(
       gDeviceTypeRegistry()->count(type),
@@ -96,6 +100,24 @@ unique_ptr<OperatorBase> TryCreateOperator(
   }
 }
 
+unique_ptr<OperatorBase> TryCreateC10Operator(
+    const string& key,
+    const OperatorDef& operator_def,
+    Workspace* ws) {
+  return C10OperatorRegistry()->Create(key, operator_def, ws);
+}
+
+unique_ptr<OperatorBase> TryCreateOperator(
+    const string& key,
+    const OperatorDef& operator_def,
+    Workspace* ws) {
+  if (auto op = TryCreateC10Operator(key, operator_def, ws)) {
+    return op;
+  } else {
+    return TryCreateC2Operator(key, operator_def, ws);
+  }
+}
+
 unique_ptr<OperatorBase> _CreateOperator(
     const OperatorDef& operator_def,
     Workspace* ws) {
@@ -656,11 +678,17 @@ std::set<std::string> GetRegisteredOperators() {
   for (const auto& name : CUDAOperatorRegistry()->Keys()) {
     all_keys.emplace(name);
   }
+
   // HIP operators
   for (const auto& name : HIPOperatorRegistry()->Keys()) {
     all_keys.emplace(name);
   }
 
+  // C10 operators
+  for (const auto& name : C10OperatorRegistry()->Keys()) {
+    all_keys.emplace(name);
+  }
+
   return all_keys;
 }
 

caffe2/core/operator_c10wrapper.cc

@@ -0,0 +1,10 @@
+#include "caffe2/core/operator_c10wrapper.h"
+
+namespace caffe2 {
+
+CAFFE_DEFINE_REGISTRY(
+    C10OperatorRegistry,
+    OperatorBase,
+    const OperatorDef&,
+    Workspace*);
+}

caffe2/core/operator_c10wrapper.h

@@ -0,0 +1,54 @@
+#pragma once
+
+#include "caffe2/core/dispatch/Dispatcher.h"
+#include "caffe2/core/operator.h"
+
+namespace caffe2 {
+
+/**
+ * To make a c10 operator "C10Add" callable from caffe2 as "C2MyAddOpName", just
+ * write
+ *
+ *     REGISTER_C10_OPERATOR_FOR_CAFFE2_DISPATCH(C10Add, C2MyAddOpName)
+ *
+ * Note: This wrapper currently only supports C10 ops that have exactly one
+ * output and take that in the last parameter as "Tensor* output".
+ * TODO: Figure out a better way to handle output parameters
+ */
+
+template <class OpSchemaDef, class Context>
+class C10OperatorWrapper final : public Operator<Context> {
+  using Schema = c10::OpSchema<OpSchemaDef>;
+
+ public:
+  C10OperatorWrapper(const OperatorDef& operator_def, Workspace* ws)
+      : Operator<Context>(operator_def, ws) {}
+
+  USE_OPERATOR_CONTEXT_FUNCTIONS;
+
+  bool RunOnDevice() override {
+    RunOnDevice_(
+        c10::guts::make_index_sequence<Schema::signature::num_args - 1>());
+    return true;
+  }
+
+ private:
+  template <size_t... InputIndex>
+  void RunOnDevice_(c10::guts::index_sequence<InputIndex...>) {
+    c10::Dispatcher<OpSchemaDef>::call(Input(InputIndex)..., Output(0));
+  }
+};
+
+CAFFE_DECLARE_REGISTRY(
+    C10OperatorRegistry,
+    OperatorBase,
+    const OperatorDef&,
+    Workspace*);
+
+// TODO Currently we only register the CPU variant. This is going to be fixed
+//      once the tensor detemplatization lands.
+#define REGISTER_C10_OPERATOR_FOR_CAFFE2_DISPATCH(OpSchemaDef, Name) \
+  CAFFE_REGISTER_CLASS(                                              \
+      C10OperatorRegistry, Name, C10OperatorWrapper<OpSchemaDef, CPUContext>)
+
+} // namespace caffe2

caffe2/operators/c10_sigmoid_op.cc

@@ -0,0 +1,14 @@
+#include "c10_sigmoid_op.h"
+#include "caffe2/core/dispatch/OpSchemaRegistration.h"
+#include "caffe2/core/operator_c10wrapper.h"
+
+using caffe2::CPUContext;
+using caffe2::Tensor;
+
+C10_DEFINE_OP_SCHEMA(caffe2::SigmoidOp);
+
+namespace caffe2 {
+REGISTER_C10_OPERATOR_FOR_CAFFE2_DISPATCH(
+    SigmoidOp,
+    C10Sigmoid_DontUseThisOpYet)
+}

caffe2/operators/c10_sigmoid_op.h

@@ -0,0 +1,17 @@
+#pragma once
+
+#include "caffe2/core/tensor.h"
+#include "caffe2/utils/Array.h"
+
+namespace caffe2 {
+
+struct SigmoidOp final {
+  static constexpr const char* name = "sigmoid";
+
+  using Signature = void(const Tensor& input, Tensor* output);
+
+  static constexpr c10::guts::array<const char*, 2> parameter_names = {
+      {"input", "output"}};
+};
+
+} // namespace caffe2

caffe2/operators/c10_sigmoid_op_cpu.cc

@@ -0,0 +1,29 @@
+#include "c10_sigmoid_op.h"
+#include "caffe2/core/dispatch/KernelRegistration.h"
+#include "caffe2/utils/eigen_utils.h"
+#include "caffe2/utils/math.h"
+
+using caffe2::CPUContext;
+using caffe2::Tensor;
+
+namespace {
+template <class DataType>
+void sigmoid_op_cpu_impl(
+    const Tensor& input,
+    Tensor* output) {
+  output->ResizeLike(input);
+
+  caffe2::ConstEigenVectorArrayMap<DataType> xM(
+      input.data<DataType>(), input.size());
+  caffe2::EigenVectorArrayMap<DataType>(
+      output->mutable_data<DataType>(), input.size()) = 1. / (1. + (-xM).exp());
+}
+} // namespace
+
+namespace c10 {
+C10_REGISTER_KERNEL(caffe2::SigmoidOp)
+    .kernel(&sigmoid_op_cpu_impl<float>)
+    .dispatchKey({DeviceTypeId::CPU,
+                  LayoutId(0),
+                  caffe2::TypeMeta::Id<float>()});
+} // namespace c10