Prevent duplicate operator registrations

Summary: - This is a simple hack to prevent duplicator operator registrations

Differential Revision: D76469487

Author: q10

fbgemm_gpu/codegen/training/pt2/embedding_split_host_pt2_cpu_wrapper_template.cpp

@@ -32,6 +32,7 @@
 {%- if has_vbe_support %}
 #include "fbgemm_gpu/utils/pt2_autograd_utils.h"
 {%- endif %}
+#include "fbgemm_gpu/utils/torch_schema.h"
 
 using Tensor = at::Tensor;
 using namespace fbgemm_gpu;
@@ -352,7 +353,9 @@ TORCH_LIBRARY_FRAGMENT(fbgemm, m) {
         ndesc, wdesc, vdesc
         )
     %}
-    {%- if not nobag %} {#-/*nobag schema is registered in cuda template*/#}
+
+    
+    if (!utils::torch::schemaExists("fbgemm::{{ embedding_codegen_forward_op }}_wrapper")) {
     m.def("{{ embedding_codegen_forward_op }}_wrapper("
         "    Tensor host_weights, "
         "    Tensor dev_weights, "
@@ -396,7 +399,7 @@ TORCH_LIBRARY_FRAGMENT(fbgemm, m) {
         , {PT2_COMPLIANT_TAG}
         {%- endif %}
         );
-    {%- endif %} {#-/*if not nobag*/#}    
+    }
     DISPATCH_TO_CPU("{{ embedding_codegen_forward_op }}_wrapper", {{ embedding_codegen_forward_op }}_cpu_wrapper);
 
     {%- else %} {#-/* backward */#}

fbgemm_gpu/include/fbgemm_gpu/utils/torch_schema.h

@@ -0,0 +1,19 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#include <ATen/core/dispatch/Dispatcher.h>
+
+namespace fbgemm_gpu::utils::torch {
+
+inline bool schemaExists(const std::string& qualified_name) {
+  return c10::Dispatcher::realSingleton()
+      .findSchema({qualified_name, ""})
+      .has_value();
+}
+
+} // namespace fbgemm_gpu::utils::torch