[TF:XLA] Create Despecializing Pass Pipeline

When comparing backends, it is useful to take an HLO optimized for one backend and perform transformations in order to match numerics.  This can be thought of as finding a lowest common denominator.

Move this grouping of passes into its own HloPassPipeline that can be reused in a few different places.

PiperOrigin-RevId: 191914799

diff --git a/tensorflow/compiler/xla/service/BUILD b/tensorflow/compiler/xla/service/BUILD
index 3a99d84..db91e80 100644 (file)
--- a/tensorflow/compiler/xla/service/BUILD
+++ b/tensorflow/compiler/xla/service/BUILD
@@ -2640,6 +2640,21 @@ tf_cc_test(
 )
 
 cc_library(
+    name = "despecializer",
+    srcs = ["despecializer.cc"],
+    hdrs = ["despecializer.h"],
+    deps = [
+        ":bfloat16_normalization",
+        ":defuser",
+        ":hlo",
+        ":hlo_pass",
+        ":hlo_pass_pipeline",
+        ":implicit_broadcast_remover",
+        "//tensorflow/compiler/xla:statusor",
+    ],
+)
+
+cc_library(
     name = "source_map_util",
     srcs = ["source_map_util.cc"],
     hdrs = ["source_map_util.h"],

diff --git a/tensorflow/compiler/xla/service/despecializer.cc b/tensorflow/compiler/xla/service/despecializer.cc
new file mode 100644 (file)
index 0000000..d938f3a
--- /dev/null
+++ b/tensorflow/compiler/xla/service/despecializer.cc
@@ -0,0 +1,35 @@
+/* Copyright 2018 The TensorFlow Authors. 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.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/compiler/xla/service/despecializer.h"
+
+#include "tensorflow/compiler/xla/service/bfloat16_normalization.h"
+#include "tensorflow/compiler/xla/service/defuser.h"
+#include "tensorflow/compiler/xla/service/implicit_broadcast_remover.h"
+
+namespace xla {
+
+Despecializer::Despecializer() : pipeline_("despecializer") {
+  // TODO(b/70588125): Also deal with window reversal in a fast way.
+  pipeline_.AddPass<Defuser>();
+  pipeline_.AddPass<ImplicitBroadcastRemover>();
+  pipeline_.AddPass<BFloat16MixedPrecisionRemoval>();
+}
+
+StatusOr<bool> Despecializer::Run(HloModule* module) {
+  return pipeline_.Run(module);
+}
+
+}  // namespace xla

diff --git a/tensorflow/compiler/xla/service/despecializer.h b/tensorflow/compiler/xla/service/despecializer.h
new file mode 100644 (file)
index 0000000..af48f4a
--- /dev/null
+++ b/tensorflow/compiler/xla/service/despecializer.h
@@ -0,0 +1,45 @@
+/* Copyright 2018 The TensorFlow Authors. 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.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_COMPILER_XLA_SERVICE_DESPECIALIZER_H_
+#define TENSORFLOW_COMPILER_XLA_SERVICE_DESPECIALIZER_H_
+
+#include "tensorflow/compiler/xla/service/hlo_module.h"
+#include "tensorflow/compiler/xla/service/hlo_pass_interface.h"
+#include "tensorflow/compiler/xla/service/hlo_pass_pipeline.h"
+#include "tensorflow/compiler/xla/statusor.h"
+
+namespace xla {
+
+// Creates an HloPassPipeline containing multiple HloPasses that can
+// despecialize an optimized HloModule. This is useful to run an HloModule
+// optimized for one specfic platform on a different platform (undoing platform
+// specific passes) with matching numerics for comparison.
+//
+// Current despecialization passes are Defuser, ImplicitBroadcastRemover,
+// and BFloat16MixedPrecisionRemoval.
+class Despecializer : public HloPassInterface {
+ public:
+  Despecializer();
+  tensorflow::StringPiece name() const override { return "despecializer"; }
+  StatusOr<bool> Run(HloModule* module) override;
+
+ private:
+  HloPassPipeline pipeline_;
+};
+
+}  // namespace xla
+
+#endif  // TENSORFLOW_COMPILER_XLA_SERVICE_DESPECIALIZER_H_