<commit_before>/* Copyright 2019 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/mlir/tensorflow/transforms/bridge.h"

#include <memory>

#include "mlir/Pass/PassManager.h"  // TF:llvm-project
#include "mlir/Transforms/Passes.h"  // TF:llvm-project
#include "tensorflow/compiler/mlir/tensorflow/transforms/passes.h"
#include "tensorflow/compiler/mlir/tensorflow/utils/bridge_logger.h"
#include "tensorflow/compiler/mlir/tensorflow/utils/error_util.h"

namespace mlir {
namespace TFTPU {

void CreateTPUBridge(OpPassManager &pm) {
  // Run shape inference so that tf_executor/tf_device ops created later will
  // likely to inherit more concrete types.
  pm.addPass(TF::CreateTFShapeInferencePass());
  OpPassManager &func_pm = pm.nest<FuncOp>();
  func_pm.addPass(tf_executor::CreateTFExecutorIslandCoarseningPass());
  func_pm.addPass(CreateTPUClusterFormationPass());
  func_pm.addPass(createCanonicalizerPass());
  // Place DecomposeResourceOpsPass before TFExecutorConstantSinking pass
  // because DecomposeResourceOpsPass uses pattern rewriter which hoists
  // changed constants out of tf_device.Launch.
  func_pm.addPass(TFDevice::CreateDecomposeResourceOpsPass());

  // Run another shape inference pass because resource ecomposition might have
  // created new partial types.
  pm.addPass(TF::CreateTFShapeInferencePass());
  OpPassManager &func_pm2 = pm.nest<FuncOp>();
  func_pm2.addPass(tf_executor::CreateTFExecutorConstantSinkingPass());
  func_pm2.addPass(TFDevice::CreateResourceOpLiftingPass());

  pm.addPass(TF::CreateResourceDeviceInferencePass());<commit_msg>Reorder island coarsening and shape inference in MLIR TF/XLA Bridge.<commit_after>/* Copyright 2019 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/mlir/tensorflow/transforms/bridge.h"

#include <memory>

#include "mlir/Pass/PassManager.h"  // TF:llvm-project
#include "mlir/Transforms/Passes.h"  // TF:llvm-project
#include "tensorflow/compiler/mlir/tensorflow/transforms/passes.h"
#include "tensorflow/compiler/mlir/tensorflow/utils/bridge_logger.h"
#include "tensorflow/compiler/mlir/tensorflow/utils/error_util.h"

namespace mlir {
namespace TFTPU {

void CreateTPUBridge(OpPassManager &pm) {
  // Run island coarsening before shape inference to allow more exact shape
  // inference using constant folding within islands.
  OpPassManager &func_pm = pm.nest<FuncOp>();
  func_pm.addPass(tf_executor::CreateTFExecutorIslandCoarseningPass());
  // Run shape inference so that tf_executor/tf_device ops created later will
  // likely to inherit more concrete types.
  pm.addPass(TF::CreateTFShapeInferencePass());
  func_pm.addPass(CreateTPUClusterFormationPass());
  func_pm.addPass(createCanonicalizerPass());
  // Place DecomposeResourceOpsPass before TFExecutorConstantSinking pass
  // because DecomposeResourceOpsPass uses pattern rewriter which hoists
  // changed constants out of tf_device.Launch.
  func_pm.addPass(TFDevice::CreateDecomposeResourceOpsPass());

  // Run another shape inference pass because resource ecomposition might have
  // created new partial types.
  pm.addPass(TF::CreateTFShapeInferencePass());
  OpPassManager &func_pm2 = pm.nest<FuncOp>();
  func_pm2.addPass(tf_executor::CreateTFExecutorConstantSinkingPass());
  func_pm2.addPass(TFDevice::CreateResourceOpLiftingPass());

  pm.addPass(TF::CreateResourceDeviceInferencePass());
{'fraction_non_alphanumeric': 0.12242686890574214, 'fraction_numerical': 0.004767063921993499, 'mean_word_length': 5.788235294117647, 'pattern_counts': {'":': 0, '<': 9, '<?xml version=': 0, '>': 9, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1459273', 'n_tokens_mistral': 1321, 'n_tokens_neox': 1280, 'n_words': 369}