move AliasWithNameOp to caffe2/operators

Summary: Pull Request resolved: pytorch#31281

Reviewed By: houseroad

Differential Revision: D19053453

fbshipit-source-id: 350bfd5c001db9c17916dcae7ade8f56db1e9841

Author: wat3rBro

caffe2/operators/alias_with_name.cc

@@ -0,0 +1,25 @@
+#include "caffe2/operators/alias_with_name.h"
+
+namespace caffe2 {
+
+REGISTER_CPU_OPERATOR(AliasWithName, AliasWithNameOp<CPUContext>);
+
+OPERATOR_SCHEMA(AliasWithName)
+    .NumInputs(1)
+    .NumOutputs(1)
+    .AllowInplace({{0, 0}})
+    .IdenticalTypeAndShape()
+    .SetDoc(R"DOC(
+Similar with AliasOp, storing the alias name as operator argument.
+)DOC")
+    .Arg("name", "name of the aliasing")
+    .Arg("is_backward", "weather or not to alias forward or backward")
+    .Input(0, "input", "Input tensor whose storage will be shared.")
+    .Output(0, "output", "Tensor of same shape as input, sharing its storage.");
+
+} // namespace caffe2
+
+C10_EXPORT_CAFFE2_OP_TO_C10_CPU(
+    AliasWithName,
+    "_caffe2::AliasWithName(Tensor input, str name, bool is_backward = False) -> (Tensor output)",
+    caffe2::AliasWithNameOp<caffe2::CPUContext>);

caffe2/operators/alias_with_name.cu

@@ -0,0 +1,12 @@
+#include "caffe2/core/context_gpu.h"
+#include "caffe2/operators/alias_with_name.h"
+
+namespace caffe2 {
+
+REGISTER_CUDA_OPERATOR(AliasWithName, AliasWithNameOp<CUDAContext>);
+
+} // namespace caffe2
+
+C10_EXPORT_CAFFE2_OP_TO_C10_CUDA(
+    AliasWithName,
+    caffe2::AliasWithNameOp<caffe2::CUDAContext>);

caffe2/operators/alias_with_name.h

@@ -0,0 +1,46 @@
+#ifndef ALIAS_WITH_NAME_OP_H_
+#define ALIAS_WITH_NAME_OP_H_
+
+#include "caffe2/core/context.h"
+#include "caffe2/core/export_caffe2_op_to_c10.h"
+#include "caffe2/core/operator.h"
+
+C10_DECLARE_EXPORT_CAFFE2_OP_TO_C10(AliasWithName)
+
+namespace caffe2 {
+
+template <class Context>
+class AliasWithNameOp final : public Operator<Context> {
+ public:
+  USE_OPERATOR_CONTEXT_FUNCTIONS;
+  template <class... Args>
+  explicit AliasWithNameOp(Args&&... args)
+      : Operator<Context>(std::forward<Args>(args)...),
+        name_(this->template GetSingleArgument<std::string>(
+            "name",
+            "invalid_name")),
+        is_backward_(
+            this->template GetSingleArgument<bool>("is_backward", false)) {
+    CAFFE_ENFORCE(
+        OperatorBase::HasArgument("name"), "You have to specify argument name");
+  }
+
+  bool RunOnDevice() override {
+    auto& input = Input(0);
+    CAFFE_ENFORCE_GE(input.numel(), 0, "Tensor is not initialized");
+
+    // This doesn't work anymore as this is "newstyle" operator
+    // OutputTensorAlias(0, input);
+
+    OperatorBase::SetOutputTensor(0, input.Alias());
+    return true;
+  }
+
+ protected:
+  std::string name_;
+  bool is_backward_;
+};
+
+} // namespace caffe2
+
+#endif // ALIAS_WITH_NAME_OP_H_

caffe2/python/operator_test/alias_with_name_test.py

@@ -0,0 +1,31 @@
+#!/usr/bin/env python3
+
+import caffe2.python.hypothesis_test_util as hu
+import hypothesis.strategies as st
+import numpy as np
+from caffe2.python import core, utils
+from hypothesis import given
+
+
+class TestAliasWithNameOp(hu.HypothesisTestCase):
+    @given(
+        shape=st.lists(st.integers(0, 5), min_size=1, max_size=3),
+        dtype=st.sampled_from([np.float32, np.int64]),
+        **hu.gcs
+    )
+    def test_alias_with_name_op(self, shape, dtype, dc, gc):
+        test_input = (100 * np.random.random(shape)).astype(dtype)
+        test_inputs = [test_input]
+
+        alias_op = core.CreateOperator(
+            "AliasWithName",
+            ["input"],
+            ["output"],
+            device_option=gc,
+        )
+        alias_op.arg.add().CopyFrom(utils.MakeArgument("name", "whatever_name"))
+
+        def reference_func(x):
+            return (x,)
+
+        self.assertReferenceChecks(gc, alias_op, test_inputs, reference_func)

caffe2/python/operator_test/torch_integration_test.py

@@ -710,6 +710,15 @@ def _piecewise_linear_ref(X):
 
         torch.testing.assert_allclose(torch.tensor(expected_output), actual_output)
 
+    def test_alias_with_name_is_in_place(self):
+        device = "cuda" if workspace.has_cuda_support else "cpu"
+        x = torch.Tensor([3, 42]).to(device)
+        y = torch.ops._caffe2.AliasWithName(x, "new_name")
+        x[1] = 6
+        torch.testing.assert_allclose(x, torch.Tensor([3, 6]).to(device))
+        # y should also change because y is alias of x
+        torch.testing.assert_allclose(y, torch.Tensor([3, 6]).to(device))
+
 
 if __name__ == '__main__':
     unittest.main()