text: // Copyright 2019 The IREE Authors
  //
  // Licensed under the Apache License v2.0 with LLVM Exceptions.
  // See https://llvm.org/LICENSE.txt for license information.
  // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
  
  #include "iree/compiler/Dialect/Flow/IR/FlowDialect.h"
  #include "iree/compiler/Dialect/Flow/IR/FlowOps.h"
  #include "iree/compiler/Dialect/IREE/IR/IREEDialect.h"
  #include "iree/compiler/Dialect/IREE/IR/IREETypes.h"
  #include "iree_tf_compiler/TF/Passes.h"
  #include "llvm/ADT/PostOrderIterator.h"
  #include "llvm/ADT/STLExtras.h"
  #include "mlir/IR/Attributes.h"
  #include "mlir/IR/Dialect.h"
  #include "mlir/IR/MLIRContext.h"
  #include "mlir/IR/RegionGraphTraits.h"
  #include "mlir/IR/SymbolTable.h"
  #include "mlir/Pass/Pass.h"
  #include "mlir/Pass/PassRegistry.h"
  #include "mlir/Support/LLVM.h"
  #include "mlir/Support/LogicalResult.h"
  #include "mlir/Transforms/Utils.h"
  #include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
  #include "tensorflow/compiler/mlir/tensorflow/ir/tf_saved_model.h"
  #include "tensorflow/compiler/mlir/tensorflow/ir/tf_types.h"
  
  namespace mlir {
  namespace iree_integrations {
  namespace TF {
  
  static LogicalResult rewriteTFVariableOpToFlowVariableOp(Operation &op,
                                                           Value flowPtr) {
    OpBuilder builder(&op);
    if (auto readVariable = dyn_cast<mlir::TF::ReadVariableOp>(op)) {
      auto load =
          builder.create<iree_compiler::IREE::Flow::VariableLoadIndirectOp>(
              readVariable.getLoc(), readVariable.value().getType(), flowPtr);
      readVariable.value().replaceAllUsesWith(load.result());
      readVariable.erase();
    } else if (auto assignVariable = dyn_cast<mlir::TF::AssignVariableOp>(op)) {
      builder.create<iree_compiler::IREE::Flow::VariableStoreIndirectOp>(
          assignVariable.getLoc(), assignVariable.value(), flowPtr);
      assignVariable.erase();
    } else {
      return op.emitError() << "could not lower resource op to flow: "
                            << op.getName();
    }
    return success();
  }
  
  static LogicalResult convertTFGlobalTensorsToFlowVariables(ModuleOp module) {
    OpBuilder globalBuilder(module.getBodyRegion());
    SymbolTable symbolTable(module);
  
    if (auto sessionInitializer =
            tf_saved_model::GetSessionInitializerOp(module)) {
      return sessionInitializer.emitError()
             << "Session initializer is not supported yet";
    }
  
    DenseMap<StringRef, std::string> symNameToFlowSymName;
    for (auto globalTensor : module.getOps<tf_saved_model::GlobalTensorOp>()) {
      auto exportedNames = tf_saved_model::GetExportedNames(globalTensor);
      std::string flowSymName;
      if (exportedNames.empty()) {
        flowSymName = "__iree_flow_" + globalTensor.sym_name().str();
      } else if (exportedNames.size() == 1) {
        flowSymName = exportedNames[0].str();
      } else {
        return globalTensor.emitError()
               << "Multiple exported names for global tensor not supported yet";
      }
      symNameToFlowSymName[globalTensor.sym_name()] = flowSymName;
      auto variableOp =
          globalBuilder.create<iree_compiler::IREE::Flow::VariableOp>(
              globalTensor.getLoc(), flowSymName, globalTensor.is_mutable(),
              globalTensor.type(), globalTensor.value());
      variableOp.setPrivate();
    }
  
    // TODO(silvasean): Make this conversion interprocedural.
    for (auto func : module.getOps<FuncOp>()) {
      if (!tf_saved_model::IsExported(func)) {
        continue;
      }
      SmallVector<unsigned, 4> argsToErase;
      OpBuilder builder(func.getBody());
      SmallVector<Value, 8> typeConversionWorklist;
      for (int i = 0, e = func.getNumArguments(); i < e; i++) {
        auto globalTensor = tf_saved_model::LookupBoundInputOfType<
            tf_saved_model::GlobalTensorOp>(func, i, symbolTable);
        if (!globalTensor) {
          continue;
        }
        auto variableAddressOp =
            builder.create<iree_compiler::IREE::Flow::VariableAddressOp>(
                globalTensor.getLoc(),
                iree_compiler::IREE::PtrType::get(globalTensor.type()),
                builder.getSymbolRefAttr(
                    symNameToFlowSymName[globalTensor.sym_name()]));
        typeConversionWorklist.push_back(variableAddressOp.getResult());
        func.getArgument(i).replaceAllUsesWith(variableAddressOp.getResult());
        argsToErase.push_back(i);
      }
      func.eraseArguments(argsToErase);
  
      while (!typeConversionWorklist.empty()) {
        Value v = typeConversionWorklist.pop_back_val();
        Type desiredType = v.getType();
        for (OpOperand &use : llvm::make_early_inc_range(v.getUses())) {
          Operation *owner = use.getOwner();
  
          // If the user is already in the flow dialect, then everything is ok.
          if (iree_compiler::IREE::Flow::FlowDialect::isDialectOp(owner)) {
            continue;
          }
  
          // If a user is just a terminator passing the value through a successor
          // operand, propagate through the successor operand.
          // TODO(silvasean): Handle case of different types in preds.
          // This would require calculating a common type.
          // This won't be a problem unless we see IR that effectively phi's
          // together different resources, which I don't think tensorflow does.
          if (BranchOpInterface branchOp = dyn_cast<BranchOpInterface>(owner)) {
            if (auto arg =
                    branchOp.getSuccessorBlockArgument(use.getOperandNumber())) {
              if (arg->getType() != desiredType) {
                arg->setType(desiredType);
                typeConversionWorklist.push_back(*arg);
              }
              continue;
            }
          }
  
          // Resource types can have subtypes (or lack thereof) and casting
          // between them is allowed. Here we just pass through.
          if (auto castOp = dyn_cast<mlir::TF::CastOp>(owner)) {
            assert(v == castOp.x());
            castOp.y().replaceAllUsesWith(castOp.x());
            castOp.erase();
            // The RAUW could have added more uses of `v`, so put it back on the
            // worklist and process it again.
            typeConversionWorklist.push_back(v);
            break;
          }
  
          if (failed(rewriteTFVariableOpToFlowVariableOp(*owner, v))) {
            return failure();
          }
        }
      }
    }
  
    // Erase all the global tensors.
    for (auto globalTensor : llvm::make_early_inc_range(
             module.getOps<tf_saved_model::GlobalTensorOp>())) {
      globalTensor.erase();
    }
    return success();
  }
  
  class LowerGlobalTensors
      : public PassWrapper<LowerGlobalTensors, OperationPass<ModuleOp>> {
   public:
    void getDependentDialects(DialectRegistry &registry) const override {
      registry.insert<mlir::tf_saved_model::TensorFlowSavedModelDialect,
                      iree_compiler::IREE::Flow::FlowDialect,
                      iree_compiler::IREEDialect>();
    }
  
    void runOnOperation() override {
      if (failed(convertTFGlobalTensorsToFlowVariables(getOperation()))) {
        signalPassFailure();
      }
    }
  };
  
  std::unique_ptr<OperationPass<ModuleOp>> createLowerGlobalTensorsPass() {
    return std::make_unique<LowerGlobalTensors>();
  }
  
  static PassRegistration<LowerGlobalTensors> pass(
      "iree-tf-saved-model-lower-global-tensors",
      "Lowers tf_saved_model global tensors to IREE flow dialect variables");
  
  }  // namespace TF
  }  // namespace iree_integrations
  }  // namespace mlir
  