Repository: alibaba/BladeDISC (language: python)

Issue:
Wont' build intrenal decoupled XLA from BladeDISC any more.



Discussion / hints:

Fix patch:
diff --git a/tao/tao_bridge/common.cc b/tao/tao_bridge/common.cc
index 2ce161f3ec9..831a06d602b 100644
--- a/tao/tao_bridge/common.cc
+++ b/tao/tao_bridge/common.cc
@@ -49,7 +49,7 @@ static void AllocateTaoBridgeFlags() {
                                 &opts->disc_cache_path));
   CHECK_OK(ReadBoolFromEnvVar("TAO_ENABLE_MLIR", false, &opts->enable_mlir));
 
-  CHECK_OK(ReadBoolFromEnvVar("TAO_MLIR_BRANCH_ONLY", false,
+  CHECK_OK(ReadBoolFromEnvVar("TAO_MLIR_BRANCH_ONLY", true,
                               &opts->tao_mlir_branch_only));
 
   CHECK_OK(ReadBoolFromEnvVar("TAO_DUMP_PASS_OUTPUT", false,
@@ -102,9 +102,8 @@ static void AllocateTaoBridgeFlags() {
   CHECK_OK(ReadBoolFromEnvVar("TAO_WHALE_SEPERATE_MICRO_BATCH", false,
                               &opts->seperate_whale_micro_batch));
 
-  // Experimental features
   CHECK_OK(ReadBoolFromEnvVar(
-      "TAO_EXPERIMENTAL_ENABLE_MLIR_WHOLE_GRAPH_COMPILATION", false,
+      "TAO_EXPERIMENTAL_ENABLE_MLIR_WHOLE_GRAPH_COMPILATION", true,
       &opts->experimental_enable_mlir_whole_graph_compilation));
 
   CHECK_OK(ReadBoolFromEnvVar("TAO_SKIP_TRAINING_GRAPH", false,


Test patch:
diff --git a/tao/tao_bridge/test/gpu/mlir/base.py b/tao/tao_bridge/test/gpu/mlir/base.py
index e64d053167b..0be4d3c2468 100644
--- a/tao/tao_bridge/test/gpu/mlir/base.py
+++ b/tao/tao_bridge/test/gpu/mlir/base.py
@@ -33,6 +33,7 @@ def setUpWithMLIR(testcase=''):
         os.environ["TF_XLA_FLAGS"]="--tf_xla_min_cluster_size=1"
         os.environ["TAO_COMPILATION_MODE_ASYNC"] = "false"
         os.environ["TAO_ENABLE_FALLBACK"] = "false"
+        os.environ["TAO_MLIR_BRANCH_ONLY"] = "true"
         TaoTestCase._setup_dump_graph(testcase)
         TaoTestCase._setup_tf_logging()
         TaoTestCase._locate_tao_compiler()
diff --git a/tao/tao_bridge/test/gpu/mlir/test_mlir.py b/tao/tao_bridge/test/gpu/mlir/test_mlir.py
index 47721e36437..9810dbd0dca 100644
--- a/tao/tao_bridge/test/gpu/mlir/test_mlir.py
+++ b/tao/tao_bridge/test/gpu/mlir/test_mlir.py
@@ -28,41 +28,18 @@ class TestTaoMlir(MlirTestCase):
     def setUp(self):
         MlirTestCase.setUpWithMLIR(self._testMethodName)
 
-    def _check_disc(self, fn):
+    def _check_launch_op(self, fn, launch_op_name):
         with open(fn) as f:
             graph_def = tf.GraphDef()
             text_format.Merge(f.read(), graph_def)
 
-            disc_op = self.get_node(graph_def, "DiscLaunch")
-
+            disc_op = self.get_node(graph_def, launch_op_name)
             mlir_func_name = disc_op.attr["mlir_function"].func.name
             self.assertNotEqual(mlir_func_name, "")
 
             mlir_func = self.get_func(graph_def, mlir_func_name)
             self.get_node(mlir_func, "MatMul")
 
-    def _check_tao(self, fn):
-        with open(fn) as f:
-            graph_def = tf.GraphDef()
-            text_format.Merge(f.read(), graph_def)
-
-            tao_op = self.get_node(graph_def, "TaoLaunch")
-
-            self.assertNotEqual(tao_op.attr["function"].func.name, "")
-            mlir_func_name = tao_op.attr["mlir_function"].func.name
-            self.assertEqual(mlir_func_name,
-                tao_op.attr["function"].func.name + "_mlir")
-
-            mlir_func = self.get_func(graph_def, mlir_func_name)
-            tao_op_inner = self.get_node(mlir_func, "TaoMlirLaunch")
-            self.assertEqual(
-                tao_op_inner.attr["function"].func.name, "")
-            inner_mlir_func_name = tao_op_inner.attr["mlir_function"].func.name
-            self.assertNotEqual(inner_mlir_func_name, "")
-
-            inner_mlir_func = self.get_func(graph_def, inner_mlir_func_name)
-            self.get_node(inner_mlir_func, "MatMul")
-
     def test_with_mlir(self):
         np_dtype = np.float32
         dtype = tf.float32
@@ -98,9 +75,9 @@ def _expected_value():
 
         fn = super(TestTaoMlir, self).dumped_file('after_tao_pass.pbtxt')
         if self.is_platform_alibaba():
-            self._check_tao(fn)
+            self._check_launch_op(fn, "TaoMlirLaunch")
         else:
-            self._check_disc(fn)
+            self._check_launch_op(fn, "DiscLaunch")
 
 
 if __name__ == "__main__":
diff --git a/tao/tao_bridge/test/gpu/mlir/test_mlir_static_rank_cache.py b/tao/tao_bridge/test/gpu/mlir/test_mlir_static_rank_cache.py
index 649fa167477..0bac1d4b74c 100644
--- a/tao/tao_bridge/test/gpu/mlir/test_mlir_static_rank_cache.py
+++ b/tao/tao_bridge/test/gpu/mlir/test_mlir_static_rank_cache.py
@@ -33,28 +33,10 @@ def _check_tao_time_const(self, fn):
             graph_def = tf.GraphDef()
             text_format.Merge(f.read(), graph_def)
 
-            tao_op = self.get_node(graph_def, "TaoLaunch")
-
-            self.assertNotEqual(tao_op.attr["function"].func.name, "")
-            mlir_func_name = tao_op.attr["mlir_function"].func.name
-            self.assertEqual(mlir_func_name,
-                tao_op.attr["function"].func.name + "_mlir")
-
-            mlir_func = self.get_func(graph_def, mlir_func_name)
-            tao_op_inner = self.get_node(mlir_func, "TaoMlirLaunch")
-            self.assertEqual(
-                tao_op_inner.attr["function"].func.name, "")
-            inner_mlir_func_name = tao_op_inner.attr["mlir_function"].func.name
-            self.assertNotEqual(inner_mlir_func_name, "")
-
-            self.assertEqual(
-                tao_op_inner.attr["mlir_function"].func.attr["_TaoXlaNumConstantArgs"].i, 1)
-            self.assertEqual(
-                tao_op_inner.attr["mlir_function"].func.attr["_TaoMlirNumFixedShapeArgs"].i, 0)
-            self.assertEqual(
-                tao_op.attr["function"].func.attr["_TaoXlaNumConstantArgs"].i, 1)
-            self.assertEqual(
-                tao_op.attr["function"].func.attr["_TaoMlirNumFixedShapeArgs"].i, 0)
+            disc_op = self.get_node(graph_def, "TaoMlirLaunch")
+
+            self.assertEqual(disc_op.attr["mlir_function"].func.attr["_TaoXlaNumConstantArgs"].i, 1)
+            self.assertEqual(disc_op.attr["mlir_function"].func.attr["_TaoMlirNumFixedShapeArgs"].i, 0)
 
     def _check_disc_time_const(self, fn):
         with open(fn) as f:
@@ -111,29 +93,10 @@ def _check_tao_fixed_shape(self, fn):
         with open(fn) as f:
             graph_def = tf.GraphDef()
             text_format.Merge(f.read(), graph_def)
+            disc_op = self.get_node(graph_def, "TaoMlirLaunch")
 
-            tao_op = self.get_node(graph_def, "TaoLaunch")
-
-            self.assertNotEqual(tao_op.attr["function"].func.name, "")
-            mlir_func_name = tao_op.attr["mlir_function"].func.name
-            self.assertEqual(mlir_func_name,
-                tao_op.attr["function"].func.name + "_mlir")
-
-            mlir_func = self.get_func(graph_def, mlir_func_name)
-            tao_op_inner = self.get_node(mlir_func, "TaoMlirLaunch")
-            self.assertEqual(
-                tao_op_inner.attr["function"].func.name, "")
-            inner_mlir_func_name = tao_op_inner.attr["mlir_function"].func.name
-            self.assertNotEqual(inner_mlir_func_name, "")
-
-            self.assertEqual(
-                tao_op_inner.attr["mlir_function"].func.attr["_TaoXlaNumConstantArgs"].i, 0)
-            self.assertEqual(
-                tao_op_inner.attr["mlir_function"].func.attr["_TaoMlirNumFixedShapeArgs"].i, 2)
-            self.assertEqual(
-                tao_op.attr["function"].func.attr["_TaoXlaNumConstantArgs"].i, 2)
-            self.assertEqual(
-                tao_op.attr["function"].func.attr["_TaoMlirNumFixedShapeArgs"].i, 0)
+            self.assertEqual(disc_op.attr["mlir_function"].func.attr["_TaoXlaNumConstantArgs"].i, 0)
+            self.assertEqual(disc_op.attr["mlir_function"].func.attr["_TaoMlirNumFixedShapeArgs"].i, 2)
 
     def _check_disc_fixed_shape(self, fn):
         with open(fn) as f:
@@ -142,8 +105,8 @@ def _check_disc_fixed_shape(self, fn):
 
             disc_op = self.get_node(graph_def, "DiscLaunch")
 
-            self.assertEqual(disc_op.attr["mlir_function"].func.attr["_TaoXlaNumConstantArgs"].i, 2)
-            self.assertEqual(disc_op.attr["mlir_function"].func.attr["_TaoMlirNumFixedShapeArgs"].i, 0)
+            self.assertEqual(disc_op.attr["mlir_function"].func.attr["_TaoXlaNumConstantArgs"].i, 0)
+            self.assertEqual(disc_op.attr["mlir_function"].func.attr["_TaoMlirNumFixedShapeArgs"].i, 2)
 
     # test for must_be_fixed_shape inputs for mlir
     def test_compile_time_fixed_shape(self):