Introduce `tfjs-opt` test utility and remove TFJSDialect static registration

This is part of an effort to remove entirely the global constructor registration
in favor of explicit registry in order to get more manageable build dependencies.

PiperOrigin-RevId: 328176806
Change-Id: I2dd02eef77d439fcaf1f79f2fd537b8021aa2e2d

tensorflow/compiler/mlir/BUILD

@@ -132,7 +132,6 @@ tf_cc_binary(
         ":tf_mlir_opt_main",
         "//tensorflow/compiler/mlir/tensorflow:mlir_roundtrip_pass_registration",
         "//tensorflow/compiler/mlir/tensorflow:tf_graph_optimization_pass",
-        "//tensorflow/compiler/mlir/tfjs:tensorflow_js_dialect_registration",
         "//tensorflow/compiler/mlir/xla:all_xla_passes_for_testing",
     ],
 )
@@ -169,3 +168,5 @@ filegroup(
     name = "litfiles",
     srcs = glob(["runlit*py"]),
 )
+
+exports_files(["run_lit.sh"])

tensorflow/compiler/mlir/glob_lit_test.bzl

@@ -43,10 +43,10 @@ def _run_lit_test(name, data, size, tags, driver, features, exec_properties):
               and specifying a default driver will abort the tests.
       features: [str], list of extra features to enable.
     """
-    if driver != _default_driver:
-        fail("There is no present support for custom drivers. Please omit" +
-             " the driver parameter when running this test. If you require" +
-             " custom driver support, please file an issue to request it.")
+
+    # Remove the default_driver from the data: it does not exist as a file and is
+    # just a placeholder from the copybara rewrite.
+    data = [d for d in data if d != _default_driver]
 
     # Disable tests on windows for now, to enable testing rest of all xla and mlir.
     native.py_test(

tensorflow/compiler/mlir/runlit.cfg.py

@@ -74,7 +74,7 @@ tool_names = [
     'tf_tfjs_translate', 'flatbuffer_to_string', 'flatbuffer_translate',
     'tf-mlir-translate', 'mlir-tflite-runner', 'tfcompile',
     'json_to_flatbuffer', 'xla-gpu-opt', 'xla-opt', 'hlo_to_llvm_ir',
-    'kernel-gen-opt', 'xla-thunks-opt'
+    'kernel-gen-opt', 'xla-thunks-opt', 'tfjs-opt'
 ]
 tools = [ToolSubst(s, unresolved='ignore') for s in tool_names]
 llvm_config.add_tool_substitutions(tools, tool_dirs)

tensorflow/compiler/mlir/tfjs/BUILD

@@ -68,17 +68,6 @@ cc_library(
     alwayslink = 1,
 )
 
-cc_library(
-    name = "tensorflow_js_dialect_registration",
-    srcs = [
-        "ir/dialect_registration.cc",
-    ],
-    deps = [
-        ":tensorflow_js",
-    ],
-    alwayslink = 1,
-)
-
 gentbl(
     name = "tfjs_optimize_inc_gen",
     tbl_outs = [
@@ -107,7 +96,6 @@ cc_library(
     ],
     deps = [
         ":tensorflow_js",
-        ":tensorflow_js_dialect_registration",
         "//tensorflow/compiler/mlir/tensorflow",
         "@llvm-project//llvm:Support",
         "@llvm-project//mlir:Analysis",
@@ -149,7 +137,6 @@ cc_library(
     ],
     deps = [
         ":tensorflow_js",
-        ":tensorflow_js_dialect_registration",
         "//tensorflow/compiler/mlir/tensorflow",
         "//tensorflow/compiler/mlir/tensorflow:convert_graphdef",
         "//tensorflow/compiler/mlir/tensorflow:export_utils",
@@ -236,3 +223,20 @@ tf_cc_binary(
         "@llvm-project//mlir:Support",
     ],
 )
+
+tf_cc_binary(
+    name = "tfjs-opt",
+    srcs = [
+        "tfjs_opt.cc",
+    ],
+    deps = [
+        ":tensorflow_js",
+        ":tensorflow_js_passes",
+        "//tensorflow/compiler/mlir:init_mlir",
+        "//tensorflow/compiler/mlir/lite:tensorflow_lite_legalize_tf",
+        "//tensorflow/compiler/mlir/tensorflow",
+        "@llvm-project//mlir:AllPassesAndDialectsNoRegistration",
+        "@llvm-project//mlir:MlirOptLib",
+        "@llvm-project//mlir:StandardOps",
+    ],
+)

tensorflow/compiler/mlir/tfjs/tests/BUILD

@@ -3,8 +3,11 @@ load("//tensorflow/compiler/mlir:glob_lit_test.bzl", "glob_lit_tests")
 package(licenses = ["notice"])
 
 glob_lit_tests(
-    data = [":test_utilities"],
-    driver = "@llvm-project//mlir:run_lit.sh",
+    data = [
+        ":test_utilities",
+        "@llvm-project//mlir:run_lit.sh",
+    ],
+    driver = "//tensorflow/compiler/mlir:run_lit.sh",
     test_file_exts = ["mlir"],
 )
 
@@ -13,7 +16,7 @@ filegroup(
     name = "test_utilities",
     testonly = True,
     data = [
-        "//tensorflow/compiler/mlir:tf-opt",
+        "//tensorflow/compiler/mlir/tfjs:tfjs-opt",
         "@llvm-project//llvm:FileCheck",
         "@llvm-project//llvm:not",
     ],

tensorflow/compiler/mlir/tfjs/tests/ops.mlir

-// RUN: tf-opt -split-input-file -verify-diagnostics -tfl-runtime-verify %s | FileCheck %s
+// RUN: tfjs-opt -split-input-file -verify-diagnostics -tfl-runtime-verify %s | FileCheck %s
 
 // -----
 

tensorflow/compiler/mlir/tfjs/tests/optimize.mlir

 // Run optimize pass only and check the results.
-// RUN: tf-opt %s -tfjs-optimize | FileCheck %s
+// RUN: tfjs-opt %s -tfjs-optimize | FileCheck %s
 
 // CHECK-LABEL: prelu_fusion
 func @prelu_fusion(%arg0: tensor<2x3xf32>) -> tensor<2x3xf32> {

tensorflow/compiler/mlir/tfjs/ir/dialect_registration.cc → tensorflow/compiler/mlir/tfjs/tfjs_opt.cc

-/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+/* Copyright 2020 Google Inc. All Rights Reserved.
 
 Licensed under the Apache License, Version 2.0 (the "License");
 you may not use this file except in compliance with the License.
@@ -13,7 +13,21 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
+#include "mlir/Dialect/StandardOps/IR/Ops.h"  // from @llvm-project
+#include "mlir/InitAllPasses.h"  // from @llvm-project
+#include "mlir/Support/MlirOptMain.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/init_mlir.h"
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
 #include "tensorflow/compiler/mlir/tfjs/ir/tfjs_ops.h"
 
-// Static initialization for TensorFlow.js op registration.
-static mlir::DialectRegistration<mlir::tfjs::TFJSDialect> tfjs_ops;
+int main(int argc, char **argv) {
+  tensorflow::InitMlir y(&argc, &argv);
+
+  mlir::registerAllPasses();
+
+  mlir::DialectRegistry registry;
+  registry.insert<mlir::StandardOpsDialect>();
+  registry.insert<mlir::TF::TensorFlowDialect>();
+  registry.insert<mlir::tfjs::TFJSDialect>();
+  return failed(mlir::MlirOptMain(argc, argv, "TF JS pass driver\n", registry));
+}

tensorflow/compiler/mlir/tfjs/transforms/optimize.cc

@@ -37,6 +37,9 @@ namespace {
 
 // Optimize TFJS operations in functions.
 struct Optimize : public PassWrapper<Optimize, FunctionPass> {
+  void getDependentDialects(DialectRegistry &registry) const final {
+    registry.insert<TFJSDialect>();
+  }
   void runOnFunction() override;
 };
 

tensorflow/compiler/mlir/tfjs/translate/tf_tfjs_translate.cc

@@ -125,7 +125,6 @@ int main(int argc, char** argv) {
                                     "TF GraphDef to TFJS JSON converter\n");
 
   MLIRContext context;
-  context.loadAllGloballyRegisteredDialects();
   llvm::SourceMgr source_mgr;
   mlir::SourceMgrDiagnosticHandler sourceMgrHandler(source_mgr, &context);