From 1d6e818f2c49a4149893f1c316c2b49c30b3e4bf Mon Sep 17 00:00:00 2001
From: Sebastian Messmer <messmer@fb.com>
Date: Tue, 8 Jan 2019 20:22:42 -0800
Subject: [PATCH] Move LayerNorm op schema to c10 (#15199)

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

In order to call it from PyTorch, this op schema can't live in caffe2 but must be included from PyTorch.
Moving it to c10. This is not where it should be in the end (that's why there is a large TODO here),
but an intermediate hack to enable this use case and proof-of-concept.

Reviewed By: ezyang

Differential Revision: D13462124

fbshipit-source-id: 1e187b9def8ef049c91e6de947ea4a85758d711b
---
 c10/core/opschema/layer_norm.cpp                    |  4 ++++
 .../c10/schemas => c10/core/opschema}/layer_norm.h  | 21 ++++++++++++++-------
 .../experimental/c10/schemas/layer_norm.cc          | 12 ++++--------
 caffe2/operators/layer_norm_op.cc                   |  6 +++---
 4 files changed, 25 insertions(+), 18 deletions(-)
 create mode 100644 c10/core/opschema/layer_norm.cpp
 rename {caffe2/operators/experimental/c10/schemas => c10/core/opschema}/layer_norm.h (66%)

diff --git a/c10/core/opschema/layer_norm.cpp b/c10/core/opschema/layer_norm.cpp
new file mode 100644
index 0000000..c8f71ba
--- /dev/null
+++ b/c10/core/opschema/layer_norm.cpp
@@ -0,0 +1,4 @@
+#include <c10/core/opschema/layer_norm.h>
+#include <c10/core/dispatch/OpSchemaRegistration.h>
+
+C10_DEFINE_OP_SCHEMA(c10::core::opschema::LayerNorm);
diff --git a/caffe2/operators/experimental/c10/schemas/layer_norm.h b/c10/core/opschema/layer_norm.h
similarity index 66%
rename from caffe2/operators/experimental/c10/schemas/layer_norm.h
rename to c10/core/opschema/layer_norm.h
index 11df08b..a2fe639 100644
--- a/caffe2/operators/experimental/c10/schemas/layer_norm.h
+++ b/c10/core/opschema/layer_norm.h
@@ -2,12 +2,18 @@
 
 #include <c10/core/Tensor.h>
 #include <c10/util/Array.h>
-#include "caffe2/core/context_base.h"
-#include "caffe2/core/tensor.h"
 
-namespace caffe2 {
-namespace ops {
+namespace at {
+class BaseContext;
+}
 
+namespace c10 {
+namespace core {
+namespace opschema {
+
+// TODO This op schema should probably not live in c10 since it's not a method
+// on Tensor. It's only here as a proof-of-concept op and for LATTE team
+// to be able to call caffe2 layer norm from PyTorch.
 struct LayerNorm final {
   static constexpr const char* name = "LayerNorm";
 
@@ -24,7 +30,7 @@ struct LayerNorm final {
       int axis,
       float epsilon,
       Cache* cache,
-      BaseContext* context);
+      at::BaseContext* context);
 
   static constexpr size_t num_dispatch_args() {return 1;}
 
@@ -34,5 +40,6 @@ struct LayerNorm final {
       {"input", "output", "output_mean", "output_stddev", "axis", "epsilon", "cache", "context"}};
 };
 
-} // namespace ops
-} // namespace caffe2
+} // namespace opschema
+} // namespace core
+} // namespace c10
diff --git a/caffe2/operators/experimental/c10/schemas/layer_norm.cc b/caffe2/operators/experimental/c10/schemas/layer_norm.cc
index 149ef0c..e5be324 100644
--- a/caffe2/operators/experimental/c10/schemas/layer_norm.cc
+++ b/caffe2/operators/experimental/c10/schemas/layer_norm.cc
@@ -1,11 +1,6 @@
-#include "caffe2/operators/experimental/c10/schemas/layer_norm.h"
-#include <c10/core/dispatch/OpSchemaRegistration.h>
+#include <c10/core/opschema/layer_norm.h>
 #include "caffe2/core/operator_c10wrapper.h"
 
-using caffe2::CPUContext;
-
-C10_DEFINE_OP_SCHEMA(caffe2::ops::LayerNorm);
-
 namespace {
 
 struct AxisParameter final {
@@ -28,10 +23,11 @@ struct EpsilonParameter final {
 };
 } // namespace
 
+
 namespace caffe2 {
 REGISTER_C10_OPERATOR_FOR_CAFFE2_DISPATCH_WITH_PARAMETERS(
-    ops::LayerNorm,
-    ops::LayerNorm::Cache,
+    c10::core::opschema::LayerNorm,
+    c10::core::opschema::LayerNorm::Cache,
     C10LayerNorm_DontUseThisOpYet,
     ParameterHelper<AxisParameter>,
     ParameterHelper<EpsilonParameter>)
diff --git a/caffe2/operators/layer_norm_op.cc b/caffe2/operators/layer_norm_op.cc
index 5e63dd4..4806844 100644
--- a/caffe2/operators/layer_norm_op.cc
+++ b/caffe2/operators/layer_norm_op.cc
@@ -1,6 +1,6 @@
 #include "caffe2/operators/layer_norm_op.h"
-#include "caffe2/operators/experimental/c10/schemas/layer_norm.h"
 #include "caffe2/utils/eigen_utils.h"
+#include <c10/core/opschema/layer_norm.h>
 #include <c10/core/dispatch/KernelRegistration.h>
 
 namespace caffe2 {
@@ -193,7 +193,7 @@ void layer_norm_c10(
     const c10::C10Tensor& sig_,
     int axis,
     float epsilon,
-    caffe2::ops::LayerNorm::Cache* cache,
+    c10::core::opschema::LayerNorm::Cache* cache,
     caffe2::BaseContext* context) {
   caffe2::Tensor X(X_);
   caffe2::Tensor Y(Y_);
@@ -220,7 +220,7 @@ void layer_norm_c10(
 }
 }
 namespace c10 {
-C10_REGISTER_KERNEL(caffe2::ops::LayerNorm)
+C10_REGISTER_KERNEL(c10::core::opschema::LayerNorm)
     .kernel(&layer_norm_c10<float>)
     .dispatchKey(c10::DispatchKey<1>{
         c10::details::TensorParameterDispatchKey{DeviceTypeId::CPU,
-- 
2.7.4