Allow exposing caffe2 operators with variable number of input tensors to c10 (#17491)

Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/17491

Before, there was no way to expose a caffe2 operator that had a variable number of inputs.
Now, this is allowed by giving the operator one tensor list input.
Note that the tensor list must be the first input, and that any other tensor inputs will be ignored and inaccessible in this case.

Reviewed By: ezyang

Differential Revision: D14220705

fbshipit-source-id: 7f921bfb581caf46b229888c409bbcc40f7dda80

diff --git a/caffe2/core/c10_operator.h b/caffe2/core/c10_operator.h
index e05d5d7..41b8058 100644 (file)
--- a/caffe2/core/c10_operator.h
+++ b/caffe2/core/c10_operator.h
@@ -103,7 +103,7 @@ inline c10::FunctionSchema make_function_schema_for_c10(const char* OperatorName
  * In caffe2/operators/MyOperator.h:
  *
  * > C10_DECLARE_CAFFE2_OPERATOR(C10MyOperator) // C10MyOperator is the name
- * used by c10 for this operator
+ *                                              // used by c10 for this operator
  *
  * In caffe2/operators/MyOperator.cc
  *
@@ -111,27 +111,34 @@ inline c10::FunctionSchema make_function_schema_for_c10(const char* OperatorName
  * >    C10MyOperator,
  * >    (std::vector<c10::Argument>{
  * >      c10::Argument("input1"),
- * >      c10::Argument("input2", c10::IntType::get()),
- * >      c10::Argument("input3", c10::FloatType::get())
+ * >      c10::Argument("argument2", c10::IntType::get()),
+ * >      c10::Argument("argument3", c10::FloatType::get())
  * >    }), (std::vector<c10::Argument>{
  * >      c10::Argument("output1"),
  * >      c10::Argument("output2")
  * >    }),
  * >    caffe2::MyOperator<caffe2::CPUContext> // This is the caffe2 operator
- * class template > )
+ * >                                           // class template
+ * > )
  *
  * In caffe2/operators/MyOperator.cu
  *
  * > C10_REGISTER_CAFFE2_OPERATOR_CUDA(C10MyOperator,
- * caffe2::MyOperator<caffe2::CUDAContext>)
+ *   caffe2::MyOperator<caffe2::CUDAContext>)
  *
  * Notes:
  * - all macros must be defined in the top level namespace, not in namespace
- * caffe2.
+ *   caffe2.
  * - all operators must call C10_DECLARE_CAFFE2_OPERATOR and
- * C10_REGISTER_CAFFE2_OPERATOR_CPU.
+ *   C10_REGISTER_CAFFE2_OPERATOR_CPU.
  * - calling C10_REGISTER_CAFFE2_OPERATOR_CUDA is optional and can be omitted if
- * you don't want to expose the operator for CUDA operations.
+ *   you don't want to expose the operator for CUDA operations.
+ * - caffe2 arguments must come after caffe2 inputs, in other words, any tensor
+ *   inputs must precede any non-tensor inputs.
+ *
+ * More complex use cases:
+ * - If your operator has a variable number of input tensors, make the first (!)
+ *   input an input of type TensorList. There must be no other tensor inputs.
  */
 #ifndef C10_MOBILE
 #define C10_DECLARE_CAFFE2_OPERATOR(OperatorName) \

diff --git a/caffe2/core/operator.h b/caffe2/core/operator.h
index e062305..75c73c6 100644 (file)
--- a/caffe2/core/operator.h
+++ b/caffe2/core/operator.h
@@ -150,11 +150,23 @@ class CAFFE2_API OperatorBase : public Observable<OperatorBase> {
         throw enf;
       }
     }
-    DCHECK_LT(idx, ivalue_inputs_.size());
-    auto ival = ivalue_inputs_[idx];
+    DCHECK_LT(0, ivalue_inputs_.size());
+    IValue ival;
+    if (ivalue_inputs_[0].isTensorList()) {
+      // if the first input is a tensor list, we get input tensors by indexing into that list.
+      // currently, this means that only tensors from that list are accessible as inputs.
+      // any hypothetical input tensors that come after the list are not accessible.
+      const auto& tensorList = ivalue_inputs_[0].toTensorListRef();
+      DCHECK_LT(idx, tensorList.size());
+      ival = tensorList[idx];
+    } else {
+      // if the first input is not a tensor list, we get input tensors by indexing into the inputs.
+      DCHECK_LT(idx, ivalue_inputs_.size());
+      ival = ivalue_inputs_[idx];
+    }
     CAFFE_ENFORCE(
         ival.isTensor(),
-        "Inpput(int, DeviceType) is only available for IValues that store Tensors");
+        "Input(int, DeviceType) is only available for IValues that store Tensors");
     Tensor tensor = caffe2::Tensor(ival.toTensor());
     CAFFE_ENFORCE_EQ(tensor.GetDeviceType(), type);
     input_tensors_[idx] = std::move(tensor);

diff --git a/torch/csrc/jit/register_c10_ops.cpp b/torch/csrc/jit/register_c10_ops.cpp
index 0ed473c..6849e4c 100644 (file)
--- a/torch/csrc/jit/register_c10_ops.cpp
+++ b/torch/csrc/jit/register_c10_ops.cpp
@@ -5,6 +5,13 @@ namespace torch {
 namespace jit {
 namespace {
 
+at::Tensor unwrap(at::Tensor&& tensor) {
+  if (tensor.requires_grad()) {
+    throw std::runtime_error("Autograd not yet supported for c10 ops.");
+  }
+  return torch::autograd::Variable(std::move(tensor)).data();
+}
+
 // TODO This currently only handles tensors with requires_grad==False correctly.
 //      It should also handle autograd.
 Operator createOperatorFromC10(const c10::OperatorHandle& op) {
@@ -16,11 +23,11 @@ Operator createOperatorFromC10(const c10::OperatorHandle& op) {
       for (auto iter = stack.end() - input_size; iter != stack.end(); ++iter) {
         // TODO Remove the .defined() check once we don't have undefined tensors on the stack anymore (@wanchaol is working on this)
         if (iter->isTensor() && iter->toTensor().defined()) {
-          at::Tensor tensor = std::move(*iter).toTensor();
-          if (tensor.requires_grad()) {
-            throw std::runtime_error("Autograd not yet supported for c10 ops.");
+          *iter = unwrap(std::move(*iter).toTensor());
+        } else if (iter->isTensorList()) {
+          for (auto& item : iter->toTensorList()->elements()) {
+            item = unwrap(std::move(item));
           }
-          *iter = torch::autograd::Variable(std::move(tensor)).data();
         }
       }