Define layer_norm schema in caffe2 (pytorch#16535)

Summary:
Pull Request resolved: pytorch#16535

There is now no need anymore to define the layer norm schema in a central location.
It can just be defined in caffe2 next to the kernel implementation.

Reviewed By: ezyang

Differential Revision: D13869503

fbshipit-source-id: c478153f8fd712ff6d507c794500286eb3583149

Author: smessmer

aten/src/ATen/core/dispatch/DispatchTable.h

@@ -99,7 +99,7 @@ class ThreadsafeOperatorTable_ final {
    }
 
    static std::string dispatch_key_to_string(TensorTypeId id) {
-     return std::string(toString(tensorTypeIdToBackend(id))) + "[" + guts::to_string(id) + "]";
+     return std::string(toString(tensorTypeIdToBackend(id))) + "[" + toString(id) + "]";
    }
 
    LeftRight<ska::flat_hash_map<TensorTypeId, DispatchTableEntry>> map_;

aten/src/ATen/core/opschema/layer_norm.cpp

@@ -7,4 +7,8 @@ std::ostream& operator<<(std::ostream& str, c10::TensorTypeId rhs) {
   return str << c10::to_string(rhs.underlyingId());
 }
 
+std::string toString(TensorTypeId id) {
+  return c10::to_string(id.underlyingId());
+}
+
 } // namespace c10

aten/src/ATen/core/opschema/layer_norm.h

@@ -32,8 +32,11 @@ class C10_API TensorTypeId final
 
   friend class TensorTypeIdCreator;
   friend C10_API std::ostream& operator<<(std::ostream&, TensorTypeId);
+  friend C10_API std::string toString(TensorTypeId);
 };
 
+std::string toString(TensorTypeId id);
+
 C10_API std::ostream& operator<<(std::ostream&, c10::TensorTypeId);
 
 } // namespace c10

c10/core/TensorTypeId.cpp

@@ -1,4 +1,4 @@
-#include <ATen/core/opschema/layer_norm.h>
+#include "caffe2/operators/layer_norm_op.h"
 #include "caffe2/core/operator_c10wrapper.h"
 
 namespace {
@@ -26,7 +26,7 @@ struct EpsilonParameter final {
 
 namespace caffe2 {
 REGISTER_C10_OPERATOR_FOR_CAFFE2_DISPATCH_WITH_PARAMETERS(
-    c10::core::opschema::LayerNorm,
+    caffe2::LayerNorm,
     C10LayerNorm_DontUseThisOpYet,
     3,
     ParameterHelper<AxisParameter>,

c10/core/TensorTypeId.h

@@ -1,8 +1,9 @@
 #include "caffe2/operators/layer_norm_op.h"
 #include "caffe2/utils/eigen_utils.h"
-#include <ATen/core/opschema/layer_norm.h>
+#include "caffe2/core/operator_c10wrapper.h"
 #include <ATen/core/dispatch/KernelRegistration.h>
 #include <c10/core/Tensor.h>
+#include <ATen/core/dispatch/OpSchemaRegistration.h>
 
 namespace caffe2 {
 
@@ -181,6 +182,22 @@ to the end.)
     .Output(1, "mean", "Mean values for each feature vector")
     .Output(2, "stddev", "Standard deviations for each feature vector");
 
+C10_DEFINE_OP_SCHEMA(LayerNorm, FunctionSchema(
+    "caffe2::layer_norm_dont_use_this_op_yet",
+    (std::vector<c10::Argument>{
+      c10::Argument("input"),
+      c10::Argument("axis", IntType::get()),
+      c10::Argument("epsilon", FloatType::get()),
+      c10::Argument("output", OptionalType::ofTensor(), nullopt, IValue()),
+      c10::Argument("output_mean", OptionalType::ofTensor(), nullopt, IValue()),
+      c10::Argument("output_stdev", OptionalType::ofTensor(), nullopt, IValue())
+    }), (std::vector<c10::Argument>{
+      c10::Argument("output"),
+      c10::Argument("mean"),
+      c10::Argument("stdev")
+    })
+));
+
 } // namespace caffe2
 
 
@@ -253,9 +270,10 @@ void layer_norm_c10(c10::Stack* stack, c10::KernelCache* cache_) { // TODO Pass
 
   return;
 }
+
 }
 namespace c10 {
-C10_REGISTER_KERNEL(c10::core::opschema::LayerNorm)
+C10_REGISTER_KERNEL(caffe2::LayerNorm)
     .withCache<Cache>()
     .kernel<&layer_norm_c10<float>>()
     .dispatchKey(CPUTensorId());

caffe2/operators/experimental/c10/schemas/layer_norm.cc

@@ -8,9 +8,12 @@
 #include "caffe2/core/operator.h"
 #include "caffe2/core/types.h"
 #include "caffe2/utils/math.h"
+#include <ATen/core/dispatch/OpSchemaRegistration.h>
 
 namespace caffe2 {
 
+C10_DECLARE_OP_SCHEMA(LayerNorm);
+
 template <class Context>
 class LayerNormOp final : public Operator<Context> {
  public: