chore: import upstream snapshot with attribution

tensorflow/compiler/mlir/tfr/passes/canonicalize.cc

@@ -0,0 +1,177 @@
+/* Copyright 2020 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 <cstdint>
+
+#include "llvm/Support/raw_ostream.h"
+#include "mlir/Conversion/SCFToControlFlow/SCFToControlFlow.h"  // from @llvm-project
+#include "mlir/Dialect/Affine/LoopUtils.h"  // from @llvm-project
+#include "mlir/Dialect/Arith/IR/Arith.h"  // from @llvm-project
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/Dialect/SCF/IR/SCF.h"  // from @llvm-project
+#include "mlir/IR/Attributes.h"  // from @llvm-project
+#include "mlir/IR/IRMapping.h"  // from @llvm-project
+#include "mlir/IR/MLIRContext.h"  // from @llvm-project
+#include "mlir/IR/Matchers.h"  // from @llvm-project
+#include "mlir/IR/PatternMatch.h"  // from @llvm-project
+#include "mlir/IR/Region.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "mlir/Support/LogicalResult.h"  // from @llvm-project
+#include "mlir/Transforms/Inliner.h"  // from @llvm-project
+#include "mlir/Transforms/InliningUtils.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
+#include "tensorflow/compiler/mlir/tfr/ir/tfr_ops.h"
+#include "tensorflow/compiler/mlir/tfr/passes/passes.h"
+
+//===----------------------------------------------------------------------===//
+// Canonicalization patterns for the scf.for and scf.if ops. They are used to
+// optimize the control flow in the tfr function. Technically, both patterns
+// should be upstreamed to be part of the op definition.
+// TODO(fengliuai): sync with the llvm upstream for both patterns.
+//
+namespace mlir {
+namespace TFR {
+
+namespace {
+
+class UnrollSCFForOp : public OpRewritePattern<scf::ForOp> {
+  using OpRewritePattern<scf::ForOp>::OpRewritePattern;
+
+ public:
+  LogicalResult matchAndRewrite(scf::ForOp for_op,
+                                PatternRewriter &rewriter) const override {
+    Location loc = for_op.getLoc();
+    APInt lower_bound, upper_bound, step;
+    if (!matchPattern(for_op.getLowerBound(), m_ConstantInt(&lower_bound)) ||
+        !matchPattern(for_op.getUpperBound(), m_ConstantInt(&upper_bound)) ||
+        !matchPattern(for_op.getStep(), m_ConstantInt(&step))) {
+      return failure();
+    }
+    uint64_t trip_count = (upper_bound - lower_bound).sdiv(step).getZExtValue();
+    if (trip_count <= 0) return failure();
+
+    // TODO(fengliuai): use loopUnrollByFactor once the iter_arg is supported
+
+    Block *single_block = for_op.getBody();
+    IRMapping mapping;
+    Value iv = for_op.getInductionVar();
+    for (auto iter_op :
+         llvm::zip(for_op.getRegionIterArgs(), for_op.getInitArgs())) {
+      mapping.map(std::get<0>(iter_op), std::get<1>(iter_op));
+    }
+    mapping.map(iv, for_op.getLowerBound());
+    for (auto i = 0; i < trip_count; ++i) {
+      if (!iv.use_empty()) {
+        // iv' = iv + step * i;
+        Value iter = arith::ConstantIndexOp::create(rewriter, loc, i);
+        Value step_cst =
+            arith::ConstantIndexOp::create(rewriter, loc, step.getSExtValue());
+        Value stride = arith::MulIOp::create(rewriter, loc, step_cst, iter);
+        Value iv_unroll =
+            arith::AddIOp::create(rewriter, loc, mapping.lookup(iv), stride);
+        mapping.map(iv, iv_unroll);
+      }
+
+      Operation *terminator_op;
+      for (auto it = single_block->begin(); it != single_block->end(); ++it) {
+        terminator_op = rewriter.clone(*it, mapping);
+      }
+      // Map the block arguments to the yield results.
+      for (auto iter_op : llvm::zip(for_op.getRegionIterArgs(),
+                                    terminator_op->getOperands())) {
+        mapping.map(std::get<0>(iter_op), std::get<1>(iter_op));
+      }
+      rewriter.eraseOp(terminator_op);
+    }
+    SmallVector<Value, 4> returned;
+    for (Value arg : for_op.getRegionIterArgs()) {
+      returned.push_back(mapping.lookup(arg));
+    }
+    rewriter.replaceOp(for_op, returned);
+    return success();
+  }
+};
+
+// TODO(fengliuai): up stream this pattern.
+class SimplifySCFIfOp : public OpRewritePattern<scf::IfOp> {
+  using OpRewritePattern<scf::IfOp>::OpRewritePattern;
+
+ public:
+  LogicalResult matchAndRewrite(scf::IfOp if_op,
+                                PatternRewriter &rewriter) const override {
+    // Then branch
+    if (matchPattern(if_op.getCondition(), m_NonZero())) {
+      return InlineRegion(if_op.getLoc(), rewriter, if_op,
+                          &if_op.getThenRegion());
+    }
+
+    // Else branch
+    if (matchPattern(if_op.getCondition(), m_Zero())) {
+      if (if_op.getElseRegion().empty()) {
+        // Remove the op
+        rewriter.eraseOp(if_op);
+        return success();
+      } else {
+        return InlineRegion(if_op.getLoc(), rewriter, if_op,
+                            &if_op.getElseRegion());
+      }
+    }
+
+    // Not a constant condition
+    return failure();
+  }
+
+ private:
+  LogicalResult InlineRegion(Location loc, PatternRewriter &rewriter,
+                             Operation *inline_point, Region *region) const;
+};
+
+LogicalResult SimplifySCFIfOp::InlineRegion(Location loc,
+                                            PatternRewriter &rewriter,
+                                            Operation *inline_point,
+                                            Region *region) const {
+  InlinerInterface interface(loc.getContext());
+  InlinerConfig config;
+  if (failed(inlineRegion(interface, config.getCloneCallback(), region,
+                          inline_point, {}, inline_point->getResults(), loc,
+                          /*shouldCloneInlinedRegion=*/true))) {
+    return failure();
+  }
+
+  // If the inlining was successful then erase the scf.if op.
+  rewriter.eraseOp(inline_point);
+  return success();
+}
+
+}  // namespace
+
+void populateCanonicalizationPatterns(func::FuncOp func,
+                                      RewritePatternSet &patterns) {
+  MLIRContext *context = func.getContext();
+  mlir::Dialect *tf = context->getLoadedDialect<mlir::TF::TensorFlowDialect>();
+  // Load all official canonicalization patterns. Here we skip the
+  // canonicalization of the ops in the tf dialect, because they couldn't
+  // propagate the attributes correctly. These optimization will be played by
+  // bridge.
+  func->walk([&](Operation *op) {
+    if (op->getDialect() != tf) {
+      op->getRegisteredInfo()->getCanonicalizationPatterns(patterns, context);
+    }
+  });
+  patterns.add<UnrollSCFForOp, SimplifySCFIfOp>(context);
+}
+
+}  // namespace TFR
+}  // namespace mlir

tensorflow/compiler/mlir/tfr/passes/decompose.cc

@@ -0,0 +1,370 @@
+/* Copyright 2020 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 <algorithm>
+#include <cmath>
+#include <cstdint>
+#include <iterator>
+#include <limits>
+#include <memory>
+#include <optional>
+#include <string>
+#include <utility>
+
+#include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/ADT/StringRef.h"
+#include "llvm/Support/Casting.h"
+#include "llvm/Support/CommandLine.h"
+#include "llvm/Support/raw_ostream.h"
+#include "mlir/Dialect/Arith/IR/Arith.h"  // from @llvm-project
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/Dialect/SCF/IR/SCF.h"  // from @llvm-project
+#include "mlir/IR/Attributes.h"  // from @llvm-project
+#include "mlir/IR/Builders.h"  // from @llvm-project
+#include "mlir/IR/BuiltinAttributes.h"  // from @llvm-project
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
+#include "mlir/IR/MLIRContext.h"  // from @llvm-project
+#include "mlir/IR/SymbolTable.h"  // from @llvm-project
+#include "mlir/IR/Value.h"  // from @llvm-project
+#include "mlir/IR/Visitors.h"  // from @llvm-project
+#include "mlir/Pass/Pass.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "mlir/Support/LogicalResult.h"  // from @llvm-project
+#include "mlir/Transforms/DialectConversion.h"  // from @llvm-project
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"  // from @llvm-project
+#include "mlir/Transforms/Inliner.h"  // from @llvm-project
+#include "mlir/Transforms/InliningUtils.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
+#include "tensorflow/compiler/mlir/tfr/ir/tfr_ops.h"
+#include "tensorflow/compiler/mlir/tfr/ir/tfr_types.h"
+#include "tensorflow/compiler/mlir/tfr/passes/passes.h"
+#include "tensorflow/compiler/mlir/tfr/utils/utils.h"
+#include "tensorflow/core/lib/monitoring/counter.h"
+
+namespace tensorflow {
+namespace {
+
+auto* tf_core_op_expansion_op_counter =
+    monitoring::Counter<1>::New("/tensorflow/core/op_expansion/op_counter",
+                                "The number of composite op expanded.", "name");
+}
+
+void IncreaseOpExpansionExecuteCounterByOne(const std::string& op_name) {
+  tf_core_op_expansion_op_counter->GetCell(op_name)->IncrementBy(1);
+}
+
+}  // namespace tensorflow
+
+//===----------------------------------------------------------------------===//
+// The pass to decompose unregistered TF ops with the TFR compose function.
+//
+namespace mlir {
+namespace TFR {
+
+namespace {
+
+// Quantize the float value based on given scale and zero point attributes.
+IntegerAttr Quantize(float value, Attribute scale_attr, Attribute zp_attr,
+                     OpBuilder builder) {
+  double scale = mlir::cast<FloatAttr>(scale_attr).getValueAsDouble();
+  int64_t zp = mlir::cast<IntegerAttr>(zp_attr).getInt();
+
+  int quantized = static_cast<int>(std::round(value / scale) + zp);
+  quantized =
+      std::min(quantized, static_cast<int>(std::numeric_limits<int8_t>::max()));
+  quantized =
+      std::max(quantized, static_cast<int>(std::numeric_limits<int8_t>::min()));
+  return builder.getI32IntegerAttr(quantized);
+}
+
+// Decompose the TF ops with the registered composition library.
+class DecomposeTFOpsPass
+    : public PassWrapper<DecomposeTFOpsPass, OperationPass<func::FuncOp>> {
+ public:
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(DecomposeTFOpsPass)
+
+  explicit DecomposeTFOpsPass(std::optional<ModuleOp> external_tfr_module)
+      : external_tfr_module_(external_tfr_module) {}
+
+  StringRef getArgument() const final { return "tfr-decompose"; }
+
+  StringRef getDescription() const final {
+    return "Decompose TF ops with the registered composition library.";
+  }
+
+  void runOnOperation() override;
+
+ private:
+  // Apply canonicalization, mainly constant folding, on the function.
+  void ApplyCanonicalization();
+
+  // Rewrite unregistered TF ops to TFR func call ops. Return failure if all the
+  // ops are registered or the compose function doesn't exist.
+  LogicalResult RewriteUnregisteredTFOps();
+
+  // Inline the TFR func call ops.
+  LogicalResult InlineTFRFuncCalls();
+
+  // Optional external symbol table to look up the TFR function.
+  std::optional<ModuleOp> external_tfr_module_;
+};
+
+#include "tensorflow/compiler/mlir/tfr/passes/generated_decompose.inc"
+
+void DecomposeTFOpsPass::ApplyCanonicalization() {
+  func::FuncOp func = getOperation();
+  RewritePatternSet patterns(&getContext());
+
+  populateWithGenerated(patterns);
+  populateCanonicalizationPatterns(func, patterns);
+
+  (void)applyPatternsGreedily(func, std::move(patterns));
+}
+
+LogicalResult DecomposeTFOpsPass::RewriteUnregisteredTFOps() {
+  func::FuncOp func = getOperation();
+  SymbolTable table(external_tfr_module_.has_value()
+                        ? *external_tfr_module_
+                        : func->getParentOfType<ModuleOp>());
+  OpBuilder builder(func);
+  bool changed = false;
+  func.walk([&table, &builder, &changed](Operation* op) {
+    // Only the un-registered ops requires decomposition. The remaining ones
+    // either will be constant folded or lowered by the rules defined in the
+    // bridge.
+    if (op->isRegistered()) {
+      return WalkResult::advance();
+    }
+
+    // Find out the compose function
+    auto compose_func_name = GetComposeFuncName(op->getName().getStringRef());
+    auto compose_func = table.lookup<TFRFuncOp>(compose_func_name);
+    if (!compose_func || compose_func.isExternal()) {
+      // There are no decomposition methods defined for this op, skip.
+      return WalkResult::advance();
+    }
+
+    // Make sure all the attributes are valid. An attribute is valid when it is
+    // in the signature or it is allowed explicitly.
+    auto compose_func_signature =
+        table.lookup<TFRFuncOp>(compose_func_name + "_");
+    if (!compose_func_signature) compose_func_signature = compose_func;
+    auto defined_attrs = compose_func_signature.getDefinedAttributeNames();
+    if (failed(ValidateAttrs(op, defined_attrs))) {
+      return WalkResult::interrupt();
+    }
+
+    tensorflow::IncreaseOpExpansionExecuteCounterByOne(
+        op->getName().getStringRef().str());
+
+    auto compose_func_type = compose_func.getFunctionType();
+    builder.setInsertionPoint(op);
+    TFRTensorType unconstrainted_tensor_type = builder.getType<TFRTensorType>();
+
+    // Create the new operands. This is mapping the operands from the target
+    // TF ops to the TFR function arguments. If the TFR function argument is
+    // a tensor_list, a "tfr.build_list" op is used to concat the available
+    // TF op operands. If the TFR function argument isn't a tensor/tensor_list,
+    // a constant is created by using the attribute stored in the TF op or the
+    // default value in the argument attribute.
+    llvm::SmallVector<Value, 4> new_operands;
+    for (auto arg : llvm::enumerate(compose_func_type.getInputs())) {
+      if (auto tensor_type = mlir::dyn_cast<TFRTensorType>(arg.value())) {
+        auto casted = CastOp::create(builder, op->getLoc(), tensor_type,
+                                     op->getOperand(arg.index()));
+        new_operands.push_back(casted);
+      } else if (auto list_type =
+                     mlir::dyn_cast<TFRTensorListType>(arg.value())) {
+        llvm::SmallVector<Value, 4> variadic_operands;
+        for (int i = arg.index(); i < op->getNumOperands(); i++) {
+          auto casted =
+              CastOp::create(builder, op->getLoc(), unconstrainted_tensor_type,
+                             op->getOperand(i));
+          variadic_operands.push_back(casted);
+        }
+        auto build_list_op = BuildListOp::create(builder, op->getLoc(),
+                                                 list_type, variadic_operands);
+        new_operands.push_back(build_list_op.getOut());
+      } else {
+        auto attr_name = compose_func.getArgAttrOfType<StringAttr>(
+            arg.index(), kAttrArgumentNameAttr);
+        auto attribute = op->getAttr(attr_name.getValue());
+        if (!attribute) {
+          attribute =
+              compose_func.getArgAttr(arg.index(), kAttrArgumentDefaultAttr);
+        }
+        if (!attribute && attr_name.getValue() == "out_type") {
+          auto type = op->getResult(0).getType();
+          if (mlir::isa<TensorType>(type)) {
+            type = mlir::cast<TensorType>(type).getElementType();
+          }
+          attribute = TypeAttr::get(type);
+        }
+        Value attr_cst;
+        // Wrap these special attributes as a special TFR constant, so the SSA
+        // value has a valid type to be used as TFR function argument. These
+        // attributes are not expected to be manipulated by the lowering passes.
+        if (mlir::isa<TypeAttr>(attribute) || mlir::isa<ArrayAttr>(attribute) ||
+            mlir::isa<StringAttr>(attribute) ||
+            mlir::isa<FlatSymbolRefAttr>(attribute)) {
+          TFRAttrType output_type = TFRAttrType::get(builder.getContext());
+          attr_cst =
+              ConstOp::create(builder, op->getLoc(), output_type, attribute);
+        } else {
+          attr_cst = mlir::arith::ConstantOp::create(
+              builder, op->getLoc(), cast<TypedAttr>(attribute));
+        }
+        new_operands.push_back(attr_cst);
+      }
+    }
+
+    // Create the TFR call op
+    auto new_op = CallOp::create(
+        builder, op->getLoc(), compose_func_type.getResults(),
+        SymbolRefAttr::get(builder.getContext(), compose_func.getName()),
+        new_operands, /*args_attrs=*/nullptr, /*res_attrs=*/nullptr);
+
+    // Replace the use of the old op. This is mapping the results from the
+    // target TF ops to the TFR function returns. If the TFR function return is
+    // a tensor_list, "tfr.get_element" op is used to extract the required TF
+    // op result.
+    llvm::SmallVector<Value, 4> new_results;
+    for (auto res : llvm::enumerate(compose_func_type.getResults())) {
+      if (mlir::dyn_cast<TFRTensorType>(res.value())) {
+        new_results.push_back(new_op.getResult(res.index()));
+      } else if (auto list_type =
+                     mlir::dyn_cast<TFRTensorListType>(res.value())) {
+        for (int i = res.index(), j = 0; i < op->getNumResults(); i++, j++) {
+          auto index = mlir::arith::ConstantOp::create(builder, op->getLoc(),
+                                                       builder.getIndexAttr(j));
+          auto element_op = GetElementOp::create(
+              builder, op->getLoc(), unconstrainted_tensor_type,
+              new_op.getResult(res.index()), index.getResult());
+          new_results.push_back(element_op.getOut());
+        }
+      }
+    }
+    for (auto res : llvm::zip(op->getResults(), new_results)) {
+      auto casted = CastOp::create(
+          builder, op->getLoc(), std::get<0>(res).getType(), std::get<1>(res));
+      std::get<0>(res).replaceAllUsesWith(casted.getOut());
+    }
+
+    // Copy all the unregisted attributes to the new op.
+    if (failed(CopyAllowedUnregisteredAttrs(op, new_op, defined_attrs))) {
+      return WalkResult::interrupt();
+    }
+
+    op->erase();
+    changed |= true;
+    return WalkResult::advance();
+  });
+
+  // If `changed` is false, it is considered as a failure, so the recursive
+  // rewrite will stop.
+  return success(changed);
+}
+
+LogicalResult DecomposeTFOpsPass::InlineTFRFuncCalls() {
+  // The Inliner will automatically use the registered dialect inliner.
+  InlinerInterface inliner(&getContext());
+  InlinerConfig config;
+  func::FuncOp func = getOperation();
+  SymbolTable table(external_tfr_module_.has_value()
+                        ? *external_tfr_module_
+                        : func->getParentOfType<ModuleOp>());
+
+  // The inliner only inlines the TFR call op.
+  bool changed = false;
+  auto walk_result = func.walk([&](CallOp call_op) {
+    auto callee = table.lookup<TFRFuncOp>(call_op.getCallee());
+    if (!callee || callee.isExternal()) return WalkResult::advance();
+
+    // Record the boundary of the inlined operations. The inlined operation will
+    // be inserted between these two operations.
+    Operation* inlined_point = call_op.getOperation();
+    Operation* after_inlined_point =
+        &*std::next(Block::iterator(call_op.getOperation()));
+
+    // Use the inliner to replace all the uses of the call_op by its
+    // composition.
+    if (failed(inlineCall(inliner, config.getCloneCallback(),
+                          cast<CallOpInterface>(call_op.getOperation()),
+                          cast<CallableOpInterface>(callee.getOperation()),
+                          callee.getCallableRegion(),
+                          /**shouldCloneInLinedRegion=*/true))) {
+      // This failure is usually because the decompose function is not defined.
+      // This call will be raised to TF ops.
+      return WalkResult::interrupt();
+    }
+
+    // Propagate all the attributes to the inlined operations, which are defined
+    // by the two boundary operations.
+    PropagateAttrsToOperations(call_op, Block::iterator(inlined_point),
+                               Block::iterator(after_inlined_point));
+
+    // Remove the call_op to finish the op expansion.
+    call_op.erase();
+    changed |= true;
+    return WalkResult::advance();
+  });
+
+  if (walk_result.wasInterrupted()) {
+    signalPassFailure();
+    return failure();
+  }
+
+  // If `changed` is false, it is considered as a failure, so the recursive
+  // rewrite will stop.
+  return success(changed);
+}
+
+void DecomposeTFOpsPass::runOnOperation() {
+  // Set a maximum iteration threshold in case there are infinite loops in the
+  // call stack.
+  int max_iterators = 10;
+  do {
+    // canonicalization
+    ApplyCanonicalization();
+
+    // rewrite unregistered tf ops. Failed either because no ops can be
+    // decomposed or the compose function isn't defined.
+    auto rewrite_status = RewriteUnregisteredTFOps();
+    // inline the tfr call op until there are no tfr.call op can be inlined.
+    auto inline_status = InlineTFRFuncCalls();
+
+    if (failed(rewrite_status) && failed(inline_status)) {
+      break;
+    }
+  } while (max_iterators-- >= 0);
+}
+
+}  // namespace
+
+// Creates an instance of the pass to decompose the TF ops.
+std::unique_ptr<OperationPass<func::FuncOp>> CreateDecomposeTFOpsPass(
+    std::optional<ModuleOp> tfr_module) {
+  return std::make_unique<DecomposeTFOpsPass>(tfr_module);
+}
+
+static PassRegistration<DecomposeTFOpsPass> pass([] {
+  return CreateDecomposeTFOpsPass();
+});
+
+}  // namespace TFR
+}  // namespace mlir

tensorflow/compiler/mlir/tfr/passes/decompose_patterns.td

@@ -0,0 +1,59 @@
+/* Copyright 2021 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 "mlir/IR/OpBase.td"
+include "mlir/IR/PatternBase.td"
+include "mlir/Dialect/Arith/IR/ArithOps.td"
+include "mlir/Dialect/Func/IR/FuncOps.td"
+include "tensorflow/compiler/mlir/tfr/ir/tfr_ops.td"
+
+class Quantize<string value> : NativeCodeCall<"TFR::Quantize(" # value # ", $0, $1, $_builder)">;
+
+class HasStringAttr<string value> : AttrConstraint<
+    CPred<"llvm::cast<StringAttr>($_self).getValue() == \"" # value # "\"">>;
+
+def QuantActRangeNonePattern :
+  Pattern<
+    (TFR_TFRQuantActRangeOp
+     (TFR_ConstOp HasStringAttr<"NONE">:$act), $scale, $zp),
+    [(TFR_ConstantTensorOp (Arith_ConstantOp ConstantAttr<I32Attr, "-128">)),
+     (TFR_ConstantTensorOp (Arith_ConstantOp ConstantAttr<I32Attr, "127">))]>;
+
+def QuantActRangeReluPattern :
+  Pattern<
+    (TFR_TFRQuantActRangeOp
+     (TFR_ConstOp HasStringAttr<"RELU">:$act),
+     (ConstantLikeMatcher F32Attr:$scale),
+     (ConstantLikeMatcher I64Attr:$zp)),
+    [(TFR_ConstantTensorOp (Arith_ConstantOp (Quantize<"0.0f"> $scale, $zp))),
+     (TFR_ConstantTensorOp (Arith_ConstantOp ConstantAttr<I32Attr, "127">))]>;
+
+def QuantActRangeRelu6Pattern :
+  Pattern<
+    (TFR_TFRQuantActRangeOp
+     (TFR_ConstOp HasStringAttr<"RELU6">:$act),
+     (ConstantLikeMatcher F32Attr:$scale),
+     (ConstantLikeMatcher I64Attr:$zp)),
+    [(TFR_ConstantTensorOp (Arith_ConstantOp (Quantize<"0.0f"> $scale, $zp))),
+     (TFR_ConstantTensorOp (Arith_ConstantOp (Quantize<"6.0f"> $scale, $zp)))]>;
+
+
+def QuantActRangeReluN1To1Pattern :
+  Pattern<
+    (TFR_TFRQuantActRangeOp
+     (TFR_ConstOp HasStringAttr<"RELU_N1_TO_1">:$act),
+     (ConstantLikeMatcher F32Attr:$scale),
+     (ConstantLikeMatcher I64Attr:$zp)),
+    [(TFR_ConstantTensorOp (Arith_ConstantOp (Quantize<"-1.0f"> $scale, $zp))),
+     (TFR_ConstantTensorOp (Arith_ConstantOp (Quantize<"1.0f"> $scale, $zp)))]>;

tensorflow/compiler/mlir/tfr/passes/passes.h

@@ -0,0 +1,52 @@
+/* Copyright 2020 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_MLIR_TFR_PASSES_PASSES_H_
+#define TENSORFLOW_COMPILER_MLIR_TFR_PASSES_PASSES_H_
+
+#include <memory>
+#include <optional>
+
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/Pass/Pass.h"  // from @llvm-project
+#include "mlir/Support/LogicalResult.h"  // from @llvm-project
+
+namespace mlir {
+namespace TFR {
+
+// Scans the func op and adds all the canonicalization patterns of the ops
+// except the tf ops, inside the function.
+void populateCanonicalizationPatterns(func::FuncOp func,
+                                      RewritePatternSet &patterns);
+
+// Decompose ops.
+std::unique_ptr<OperationPass<func::FuncOp>> CreateDecomposeTFOpsPass(
+    std::optional<ModuleOp> tfr_module = std::nullopt);
+
+// Rewrites quantized operands and results with their storage types.
+// This pass should be run at module level after decomposition, if there are
+// quantized operands or results.
+std::unique_ptr<OperationPass<ModuleOp>> CreateRewriteQuantizedIOPass();
+
+// Raise to TF ops.
+std::unique_ptr<OperationPass<func::FuncOp>> CreateRaiseToTFOpsPass(
+    std::optional<ModuleOp> tfr_module = std::nullopt,
+    bool materialize_derived_attrs = false);
+
+}  // namespace TFR
+}  // namespace mlir
+
+#endif  // TENSORFLOW_COMPILER_MLIR_TFR_PASSES_PASSES_H_

tensorflow/compiler/mlir/tfr/passes/raise_to_tf.cc

@@ -0,0 +1,522 @@
+/* Copyright 2020 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 <memory>
+#include <optional>
+#include <string>
+#include <utility>
+
+#include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/DenseSet.h"
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/ADT/StringMap.h"
+#include "llvm/ADT/StringRef.h"
+#include "llvm/Support/Casting.h"
+#include "llvm/Support/CommandLine.h"
+#include "llvm/Support/raw_ostream.h"
+#include "mlir/Dialect/Arith/IR/Arith.h"  // from @llvm-project
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/Dialect/SCF/IR/SCF.h"  // from @llvm-project
+#include "mlir/IR/Attributes.h"  // from @llvm-project
+#include "mlir/IR/Builders.h"  // from @llvm-project
+#include "mlir/IR/BuiltinAttributes.h"  // from @llvm-project
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
+#include "mlir/IR/Diagnostics.h"  // from @llvm-project
+#include "mlir/IR/Location.h"  // from @llvm-project
+#include "mlir/IR/MLIRContext.h"  // from @llvm-project
+#include "mlir/IR/Matchers.h"  // from @llvm-project
+#include "mlir/IR/OperationSupport.h"  // from @llvm-project
+#include "mlir/IR/SymbolTable.h"  // from @llvm-project
+#include "mlir/IR/Types.h"  // from @llvm-project
+#include "mlir/IR/Value.h"  // from @llvm-project
+#include "mlir/IR/Visitors.h"  // from @llvm-project
+#include "mlir/Pass/Pass.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "mlir/Support/LogicalResult.h"  // from @llvm-project
+#include "mlir/Transforms/DialectConversion.h"  // from @llvm-project
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"  // from @llvm-project
+#include "mlir/Transforms/InliningUtils.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
+#include "tensorflow/compiler/mlir/tfr/ir/tfr_ops.h"
+#include "tensorflow/compiler/mlir/tfr/ir/tfr_types.h"
+#include "tensorflow/compiler/mlir/tfr/passes/passes.h"
+#include "tensorflow/compiler/mlir/tfr/utils/utils.h"
+
+//===----------------------------------------------------------------------===//
+// The pass to rewrite the TFR function call ops by TF ops. The callee of the
+// TFR function call defines the signatures of the TF ops.
+//
+namespace mlir {
+namespace TFR {
+
+namespace {
+
+// This pattern is to rewrite the "tfr.call" op and the "tfr.cast" ops on the
+// operands by a TF op with "tfr.cast" ops on the results. The result type of
+// the new TF op is an unranked tensor with element type derived.
+class RewriteTFRCallOp : public OpRewritePattern<CallOp> {
+  using OpRewritePattern<CallOp>::OpRewritePattern;
+
+ public:
+  explicit RewriteTFRCallOp(MLIRContext* context, const SymbolTable& table,
+                            bool materialize_derived_attrs)
+      : OpRewritePattern<CallOp>(context),
+        symbol_table_(table),
+        materialize_derived_attrs_(materialize_derived_attrs) {}
+
+  LogicalResult matchAndRewrite(CallOp call_op,
+                                PatternRewriter& rewriter) const override;
+
+ private:
+  // Derives the attribute values for the attributes attached to the
+  // `input_tfr_type`. These attributes are only for the element type of the
+  // inputs, and these type information has been collected in the `input_types`.
+  // The result is stored in `derived_attrs` as the named attributes. Returns
+  // failure if the attributes stored in the `input_tfr_type` violates the
+  // assumptions.
+  LogicalResult AddDerivedAttrs(
+      PatternRewriter& rewriter, Type input_tfr_type,
+      ArrayRef<Attribute> input_types,
+      llvm::StringMap<Attribute>* derived_attrs) const;
+
+  // Collects the operands and attributes for the TF op. At the same time, it
+  // collects all the derived attribute values to derive the output types of the
+  // TF op.
+  LogicalResult CollectInputsAndAttributes(
+      PatternRewriter& rewriter, TFRFuncOp signature, CallOp call_op,
+      SmallVectorImpl<Value>* inputs, NamedAttrList* arg_attrs,
+      llvm::StringMap<Attribute>* derived_attrs) const;
+
+  // Uses the collected attribute values to derive all the output types.
+  LogicalResult DeriveOutputTypes(Location loc, FunctionType signature,
+                                  const llvm::StringMap<Attribute>& attrs,
+                                  SmallVectorImpl<Type>* output_types) const;
+
+  // Creates the TF op and also the necessary tfr.cast ops to replace the
+  // original TFR call op.
+  LogicalResult CreateAndReplaceOp(
+      PatternRewriter& rewriter, CallOp call_op,
+      const SmallVectorImpl<Type>& output_types,
+      const SmallVectorImpl<Value>& inputs, const NamedAttrList& attr_list,
+      const llvm::StringMap<Attribute>& derived_attrs) const;
+
+  // Converts the attribute to the specific type.
+  Attribute ProcessAttributeValue(Attribute attr, StringAttr attr_type) const;
+
+  Type GetFixedElementType(StringRef element_type, Builder& builder) const {
+    if (element_type == "i32_") return builder.getI32Type();
+    if (element_type == "i64_") return builder.getI64Type();
+    if (element_type == "f32_") return builder.getF32Type();
+    if (element_type == "i1_") return builder.getI1Type();
+    return {};
+  }
+
+  // Adds a tf.Cast op if the tfr.tensor attribute indicated a fixed element
+  // type.
+  // TODO(fengliuai): This method is required when the operand types are not set
+  // by the frontend correctly.
+  Value CastToNonDerivedType(PatternRewriter& rewriter, Location loc,
+                             CastOp cast_op, Type input_tfr_type) const {
+    auto tensor_type = mlir::dyn_cast<TFRTensorType>(input_tfr_type);
+    if (!tensor_type) return cast_op.getArg();
+
+    auto attr_names = tensor_type.getAttrKeys();
+    if (attr_names.empty() || attr_names.size() > 1) return cast_op.getArg();
+    StringRef tfr_type_attr = attr_names[0].getValue();
+    if (!fixed_elt_type_attrs_.contains(tfr_type_attr)) return cast_op.getArg();
+
+    Type result_elt_type = GetFixedElementType(tfr_type_attr, rewriter);
+    if (!result_elt_type) {
+      return cast_op.getArg();
+    }
+
+    Type original_input_type =
+        mlir::cast<TypeAttr>(cast_op.getInputElementType()).getValue();
+    if (result_elt_type != original_input_type) {
+      UnrankedTensorType result_type = UnrankedTensorType::get(result_elt_type);
+      return TF::CastOp::create(rewriter, loc, result_type, cast_op.getArg());
+    }
+    return cast_op.getArg();
+  }
+
+  // For variadic operands, we have to enforce them to use the same types.
+  // TODO(fengliuai): This method is required when the operand types are not set
+  // by the frontend correctly.
+  void CastValuesToSameType(PatternRewriter& rewriter, Location loc,
+                            const llvm::SmallVectorImpl<Attribute>& input_types,
+                            llvm::SmallVectorImpl<Value>& input_values) const {
+    if (input_types.size() <= 1) return;
+
+    Type target_input_type = mlir::cast<TypeAttr>(input_types[0]).getValue();
+    auto result_type = UnrankedTensorType::get(target_input_type);
+    for (auto i = 1; i < input_types.size(); ++i) {
+      Type current_input_type = mlir::cast<TypeAttr>(input_types[i]).getValue();
+      if (current_input_type != target_input_type) {
+        input_values[i] =
+            TF::CastOp::create(rewriter, loc, result_type, input_values[i]);
+      }
+    }
+  }
+
+  const SymbolTable& symbol_table_;
+  const bool materialize_derived_attrs_;
+  const llvm::SmallDenseSet<StringRef, 4> fixed_elt_type_attrs_{"i32_", "i64_",
+                                                                "f32_", "i1_"};
+};
+
+LogicalResult RewriteTFRCallOp::AddDerivedAttrs(
+    PatternRewriter& rewriter, Type input_tfr_type,
+    ArrayRef<Attribute> input_types,
+    llvm::StringMap<Attribute>* derived_attrs) const {
+  // If there is an attribute associated to the input in the signature, we
+  // store it as an derived attribute.
+  if (auto tensor_type = mlir::dyn_cast<TFRTensorType>(input_tfr_type)) {
+    auto attr_names = tensor_type.getAttrKeys();
+    if (attr_names.empty()) return success();
+
+    if (attr_names.size() == 1) {
+      derived_attrs->insert({attr_names[0].getValue(), input_types[0]});
+      return success();
+    }
+  }
+
+  // If there is an attribute associated to the input in the signature,
+  // we store it as an derived attribute.
+  if (auto list_type = mlir::dyn_cast<TFRTensorListType>(input_tfr_type)) {
+    auto attr_names = list_type.getAttrKeys();
+    if (attr_names.empty()) return success();
+
+    // N*T case
+    if (attr_names.size() == 2) {
+      derived_attrs->insert({attr_names[0].getValue(),
+                             rewriter.getI32IntegerAttr(input_types.size())});
+      // Note that this uses the first element of the list to infer the T value.
+      // A tf.Cast is required to cast the other inputs to the same type.
+      derived_attrs->insert({attr_names[1].getValue(), input_types[0]});
+      return success();
+    }
+
+    // list(dtype) case
+    if (attr_names.size() == 1) {
+      derived_attrs->insert(
+          {attr_names[0].getValue(), rewriter.getArrayAttr(input_types)});
+      return success();
+    }
+  }
+
+  return failure();
+}
+
+LogicalResult RewriteTFRCallOp::CollectInputsAndAttributes(
+    PatternRewriter& rewriter, TFRFuncOp signature, CallOp call_op,
+    SmallVectorImpl<Value>* inputs, NamedAttrList* arg_attrs,
+    llvm::StringMap<Attribute>* derived_attrs) const {
+  for (const auto& operand :
+       llvm::enumerate(signature.getFunctionType().getInputs())) {
+    // If the index is larger than the operand number of the call_op, the
+    // default value of the operand needs to be used.
+    if (operand.index() >= call_op.getNumOperands()) {
+      auto attr_name = signature.getArgAttrOfType<StringAttr>(
+          operand.index(), kAttrArgumentNameAttr);
+      auto attr_value =
+          signature.getArgAttr(operand.index(), kAttrArgumentDefaultAttr);
+      arg_attrs->push_back(
+          rewriter.getNamedAttr(attr_name.getValue(), attr_value));
+      continue;
+    }
+
+    // The index is valid for the call_op.
+    Value input = call_op.getOperand(operand.index());
+    Operation* input_op = input.getDefiningOp();
+    auto input_tfr_type =
+        signature.getFunctionType().getInputs()[operand.index()];
+
+    // There are three cases for the preceding input_op:
+
+    // 1. The preceding op can be a tfr.cast op, which will be fused to the
+    // current op, so the result op has input with tensor type.
+    if (auto cast_op = dyn_cast_or_null<CastOp>(input_op)) {
+      Value input_to_cast = CastToNonDerivedType(rewriter, call_op.getLoc(),
+                                                 cast_op, input_tfr_type);
+      inputs->push_back(input_to_cast);
+      if (failed(AddDerivedAttrs(rewriter, input_tfr_type,
+                                 {cast_op.getInputElementType()},
+                                 derived_attrs))) {
+        return failure();
+      }
+      continue;
+    }
+
+    // 2. The preceding op is a tfr.build_list op, which collects multiple
+    // values with tensor types via the tfr.cast ops. These ops will be fused
+    // to the current op as well, so all the tfr.cast op inputs will be inputs
+    // to the result op.
+    if (auto list_op = dyn_cast_or_null<BuildListOp>(input_op)) {
+      // Find out all the inputs to the build list op
+      // TODO(fengliuai): make build_list op only take tensor argument
+      llvm::SmallVector<Attribute, 4> list_input_types;
+      llvm::SmallVector<Value, 4> list_inputs;
+      for (auto list_input : list_op.getOperands()) {
+        auto cast_op = dyn_cast_or_null<CastOp>(list_input.getDefiningOp());
+        if (!cast_op) return failure();
+        list_inputs.push_back(cast_op.getArg());
+        list_input_types.push_back(cast_op.getInputElementType());
+      }
+      CastValuesToSameType(rewriter, call_op.getLoc(), list_input_types,
+                           list_inputs);
+      inputs->append(list_inputs.begin(), list_inputs.end());
+      if (failed(AddDerivedAttrs(rewriter, input_tfr_type, list_input_types,
+                                 derived_attrs))) {
+        return failure();
+      }
+      continue;
+    }
+
+    // 3. The preceding op is a constant, thus the value of this constant is
+    // used to create an attribute of the result op, according to the signature.
+    Attribute arg_value;
+    // A failure indicates the argument isn't a constant value, so we should
+    // not use it as an attribute.
+    if (!matchPattern(input, m_Constant(&arg_value))) {
+      return failure();
+    }
+    auto attr_name = signature.getArgAttrOfType<StringAttr>(
+        operand.index(), kAttrArgumentNameAttr);
+    auto attr_type = signature.getArgAttrOfType<StringAttr>(
+        operand.index(), kAttrArgumentTypeAttr);
+    auto value = ProcessAttributeValue(arg_value, attr_type);
+    arg_attrs->push_back(rewriter.getNamedAttr(attr_name.getValue(), value));
+  }
+  return success();
+}
+
+Attribute RewriteTFRCallOp::ProcessAttributeValue(Attribute attr,
+                                                  StringAttr attr_type) const {
+  if (!attr_type) return attr;
+
+  if (attr_type.getValue() == "tensor") {
+    if (auto f = mlir::dyn_cast<FloatAttr>(attr)) {
+      RankedTensorType type = RankedTensorType::get({}, f.getType());
+      return DenseFPElementsAttr::get(type, attr);
+    }
+    // TODO(fengliuai): handles ArrayAttr. Note that it can be nested ArrayAttr.
+  }
+
+  return attr;
+}
+
+// For each output, uses the attribute name associated to the tfr types to find
+// out the attribute value from the collected `attrs` and create the output type
+// of the result op by using the attribute value as the element type.
+LogicalResult RewriteTFRCallOp::DeriveOutputTypes(
+    Location loc, FunctionType signature,
+    const llvm::StringMap<Attribute>& attrs,
+    SmallVectorImpl<Type>* output_types) const {
+  for (auto res : llvm::enumerate(signature.getResults())) {
+    if (auto tensor_type = mlir::dyn_cast<TFRTensorType>(res.value())) {
+      // tfr.tensor should only have one attribute attached.
+      auto attr_key = tensor_type.getAttrKeys().front();
+      Builder builder(signature.getContext());
+      if (auto attr = attrs.lookup(attr_key.getValue())) {
+        output_types->push_back(
+            UnrankedTensorType::get(mlir::cast<TypeAttr>(attr).getValue()));
+      } else if (Type element_type =
+                     GetFixedElementType(attr_key.getValue(), builder)) {
+        output_types->push_back(UnrankedTensorType::get(element_type));
+      } else {
+        emitError(loc) << "type " << attr_key.getValue()
+                       << " can't be resolved for the signature of the op";
+        return failure();
+      }
+      continue;
+    }
+
+    if (auto list_type = mlir::dyn_cast<TFRTensorListType>(res.value())) {
+      // There are two cases: N*T or list(dtype)
+      auto attr_keys = list_type.getAttrKeys();
+      // N*T case
+      if (attr_keys.size() == 2) {
+        // The first one is N, and the second one is T
+        int list_size =
+            mlir::cast<IntegerAttr>(attrs.lookup(attr_keys[0].getValue()))
+                .getInt();
+        Type list_type =
+            mlir::cast<TypeAttr>(attrs.lookup(attr_keys[1].getValue()))
+                .getValue();
+        for (int i = 0; i < list_size; ++i) {
+          output_types->push_back(UnrankedTensorType::get(list_type));
+        }
+        continue;
+      }
+      // TODO(fengliuai): list(dtype) case
+    }
+    return failure();
+  }
+  return success();
+}
+
+LogicalResult RewriteTFRCallOp::CreateAndReplaceOp(
+    PatternRewriter& rewriter, CallOp call_op,
+    const SmallVectorImpl<Type>& output_types,
+    const SmallVectorImpl<Value>& inputs, const NamedAttrList& attr_list,
+    const llvm::StringMap<Attribute>& derived_attrs) const {
+  // Create the new op
+  Location loc = call_op.getLoc();
+  rewriter.setInsertionPointAfter(call_op);
+  std::string tf_op_name = GetTFOpName(call_op.getCallee());
+  OperationState new_state(loc, tf_op_name, inputs, output_types, attr_list);
+  Operation* new_op = rewriter.create(new_state);
+  if (materialize_derived_attrs_) {
+    for (const auto& attr : derived_attrs) {
+      // Add or update the derived attribute with the value. Skip the fixed
+      // element type attributes, in case they are present in the NodeDef.
+      if (!fixed_elt_type_attrs_.contains(attr.first())) {
+        new_op->setAttr(attr.first(), attr.second);
+      }
+    }
+  }
+  // Create the tfr.cast ops on the results and replace the uses of the
+  // original call op.
+  TFRTensorType unconstrainted_type = rewriter.getType<TFRTensorType>();
+  SmallVector<Value, 4> new_results;
+  for (auto res : llvm::enumerate(call_op.getResultTypes())) {
+    Type res_type = res.value();
+    if (mlir::dyn_cast<TFRTensorType>(res_type)) {
+      Value new_res = new_op->getResult(res.index());
+      auto casted = CastOp::create(rewriter, loc, res_type, new_res);
+      new_results.push_back(casted.getOut());
+    } else if (auto list_type =
+                   mlir::dyn_cast<TFRTensorListType>(res.value())) {
+      SmallVector<Value, 4> tensor_list;
+      for (int i = res.index(); i < new_op->getNumResults(); i++) {
+        Value new_res = new_op->getResult(i);
+        auto casted =
+            CastOp::create(rewriter, loc, unconstrainted_type, new_res);
+        tensor_list.push_back(casted.getOut());
+      }
+      auto list_op = BuildListOp::create(rewriter, loc, res_type, tensor_list);
+      new_results.push_back(list_op.getOut());
+    }
+  }
+
+  // Copy all the allowed attributes to the new op.
+  if (failed(CopyNonSymbolRefAttrs(call_op, new_op))) return failure();
+
+  rewriter.replaceOp(call_op, new_results);
+  return success();
+}
+
+LogicalResult RewriteTFRCallOp::matchAndRewrite(
+    CallOp call_op, PatternRewriter& rewriter) const {
+  // Get the func op and verify that it is external. The type of this external
+  // func op is used as the signature of the corresponding TF ops. All the
+  // external func ops have the trailing underscore.
+  std::string external_callee_name = call_op.getCallee().str().append("_");
+  TFRFuncOp func = symbol_table_.lookup<TFRFuncOp>(external_callee_name);
+  if (!func || !func.isExternal()) return failure();
+  // Get the inputs and attributes. The attributes include these from the
+  // argument list and also these derived from the inputs.
+  SmallVector<Value, 4> inputs;
+  NamedAttrList argument_attrs;
+  llvm::StringMap<Attribute> derived_attrs;
+  if (failed(CollectInputsAndAttributes(rewriter, func, call_op, &inputs,
+                                        &argument_attrs, &derived_attrs))) {
+    return failure();
+  }
+
+  // Derive the output types. The result type is derived by using the
+  // attributes attched to the result type of the signature. The attribute
+  // value should be either in the attribute argument list or the derived
+  // attribute from the input tensors. All the result type
+  // are unranked, and shape inference should be applied afterwards.
+  SmallVector<Type, 4> output_types;
+
+  // Merge the attributes from the argument list to the derived ones.
+  for (auto& attr : argument_attrs) {
+    derived_attrs.insert({attr.getName(), attr.getValue()});
+  }
+
+  // Derive the output types by using the attributes attached to the tfr
+  // types.
+  if (failed(DeriveOutputTypes(call_op->getLoc(), func.getFunctionType(),
+                               derived_attrs, &output_types))) {
+    return failure();
+  }
+
+  // Create the new op and replace the old TFR call op.
+  return CreateAndReplaceOp(rewriter, call_op, output_types, inputs,
+                            argument_attrs, derived_attrs);
+}
+
+// Raise TFR call ops to the TF ops.
+class RaiseToTFOpsPass
+    : public PassWrapper<RaiseToTFOpsPass, OperationPass<func::FuncOp>> {
+ public:
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(RaiseToTFOpsPass)
+
+  void getDependentDialects(DialectRegistry& registry) const override {
+    registry.insert<TFRDialect, TF::TensorFlowDialect, scf::SCFDialect,
+                    arith::ArithDialect, func::FuncDialect>();
+  }
+
+  explicit RaiseToTFOpsPass(std::optional<ModuleOp> tfr_module,
+                            bool materialize_derived_attrs)
+      : external_tfr_module_(tfr_module),
+        materialize_derived_attrs_(materialize_derived_attrs) {}
+
+  StringRef getArgument() const final { return "tfr-raise-to-tf"; }
+
+  StringRef getDescription() const final {
+    return "Raise all the TFR call ops to TF ops.";
+  }
+
+  void runOnOperation() override;
+
+ private:
+  std::optional<ModuleOp> external_tfr_module_;
+  const bool materialize_derived_attrs_;
+};
+
+void RaiseToTFOpsPass::runOnOperation() {
+  func::FuncOp func = getOperation();
+  MLIRContext* ctx = &getContext();
+  SymbolTable table(external_tfr_module_.has_value()
+                        ? *external_tfr_module_
+                        : func->getParentOfType<ModuleOp>());
+
+  RewritePatternSet patterns(&getContext());
+  patterns.add<RewriteTFRCallOp>(ctx, table, materialize_derived_attrs_);
+
+  populateCanonicalizationPatterns(func, patterns);
+
+  (void)applyPatternsGreedily(func, std::move(patterns));
+}
+}  // namespace
+
+// Creates an instance of the pass to raise TFR call ops to the TF ops.
+std::unique_ptr<OperationPass<func::FuncOp>> CreateRaiseToTFOpsPass(
+    std::optional<ModuleOp> tfr_module, bool materialize_derived_attrs) {
+  return std::make_unique<RaiseToTFOpsPass>(tfr_module,
+                                            materialize_derived_attrs);
+}
+
+static PassRegistration<RaiseToTFOpsPass> pass([] {
+  return CreateRaiseToTFOpsPass();
+});
+
+}  // namespace TFR
+}  // namespace mlir

tensorflow/compiler/mlir/tfr/passes/rewrite_quantized_io.cc

@@ -0,0 +1,120 @@
+/* Copyright 2020 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 <memory>
+#include <string>
+
+#include "llvm/ADT/StringRef.h"
+#include "llvm/Support/Casting.h"
+#include "mlir/Dialect/Quant/IR/QuantTypes.h"  // from @llvm-project
+#include "mlir/IR/Block.h"  // from @llvm-project
+#include "mlir/IR/Builders.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "mlir/IR/Value.h"  // from @llvm-project
+#include "mlir/Pass/Pass.h"  // from @llvm-project
+#include "mlir/Pass/PassRegistry.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tfr/ir/tfr_ops.h"
+#include "tensorflow/compiler/mlir/tfr/passes/passes.h"
+
+namespace mlir {
+namespace TFR {
+
+class RewriteQuantizedIOPass
+    : public PassWrapper<RewriteQuantizedIOPass, OperationPass<ModuleOp>> {
+ public:
+  StringRef getArgument() const final { return "tfr-rewrite-quantized-io"; }
+
+  StringRef getDescription() const final {
+    return "Replaces operands and results that has quantized type with their "
+           "storage types.";
+  }
+  void runOnOperation() override;
+};
+
+void RewriteQuantizedIOPass::runOnOperation() {
+  ModuleOp module = getOperation();
+  OpBuilder builder(module);
+  module.walk([&](func::FuncOp func) {
+    Block& block = func.front();
+    Operation* terminator = block.getTerminator();
+
+    // Replace input_arg(tensor<quant_type>) -> tfr.cast
+    // with input_arg(tensor<storage_type>) -> tfr.cast
+    for (BlockArgument arg : block.getArguments()) {
+      Type arg_type = arg.getType();
+      if (auto quant_type = llvm::dyn_cast<quant::QuantizedType>(
+              llvm::cast<TensorType>(arg_type).getElementType())) {
+        if (arg.hasOneUse() && llvm::isa<TFR::CastOp>(*arg.user_begin())) {
+          arg.setType(llvm::cast<TensorType>(arg_type).clone(
+              quant_type.getStorageType()));
+        } else {
+          std::string error_message;
+          llvm::raw_string_ostream os{error_message};
+          os << "The argument with type ";
+          arg.getType().print(os);
+          os << " should have one user, which should be tfr.cast.";
+          func->emitError(error_message);
+          return;
+        }
+      }
+    }
+
+    builder.setInsertionPoint(terminator);
+    // Replace tfr.cast(tensor<quant_type>) -> output
+    // with tfr.cast(tensor<storage_type>) -> output
+    for (OpOperand& returned_value : terminator->getOpOperands()) {
+      auto returned_type =
+          llvm::dyn_cast<TensorType>(returned_value.get().getType());
+      if (!returned_type ||
+          !llvm::isa<quant::QuantizedType>(returned_type.getElementType())) {
+        continue;
+      }
+
+      if (auto returned_op =
+              returned_value.get().getDefiningOp<TFR::CastOp>()) {
+        auto new_type = returned_type.clone(
+            llvm::cast<quant::QuantizedType>(returned_type.getElementType())
+                .getStorageType());
+        auto new_op = TFR::CastOp::create(builder, returned_op->getLoc(),
+                                          new_type, returned_op.getArg());
+        returned_value.set(new_op.getResult());
+        if (returned_op.use_empty()) {
+          returned_op.erase();
+        }
+      } else {
+        returned_value.get().getDefiningOp()->emitError(
+            "The producer of quantized type result should be a tfr.cast op.");
+        return;
+      }
+    }
+
+    auto new_func_type = builder.getFunctionType(block.getArgumentTypes(),
+                                                 terminator->getOperandTypes());
+    func.setType(new_func_type);
+  });
+}
+
+// Creates an instance of the pass to decompose the TF ops.
+std::unique_ptr<OperationPass<ModuleOp>> CreateRewriteQuantizedIOPass() {
+  return std::make_unique<RewriteQuantizedIOPass>();
+}
+
+static PassRegistration<RewriteQuantizedIOPass> pass([] {
+  return CreateRewriteQuantizedIOPass();
+});
+
+}  // namespace TFR
+}  // namespace mlir

tensorflow/compiler/mlir/tfr/passes/tfr_opt.cc

@@ -0,0 +1,44 @@
+/* Copyright 2020 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 "mlir/Dialect/Arith/IR/Arith.h"  // from @llvm-project
+#include "mlir/Dialect/Func/Extensions/AllExtensions.h"  // from @llvm-project
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/Dialect/Quant/IR/Quant.h"  // from @llvm-project
+#include "mlir/Dialect/Quant/IR/QuantTypes.h"  // from @llvm-project
+#include "mlir/Dialect/SCF/IR/SCF.h"  // from @llvm-project
+#include "mlir/Dialect/Shape/IR/Shape.h"  // from @llvm-project
+#include "mlir/InitAllDialects.h"  // from @llvm-project
+#include "mlir/InitAllPasses.h"  // from @llvm-project
+#include "mlir/Tools/mlir-opt/MlirOptMain.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/init_mlir.h"
+#include "tensorflow/compiler/mlir/quantization/common/ir/QuantOps.h"
+#include "tensorflow/compiler/mlir/tensorflow/dialect_registration.h"
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
+#include "tensorflow/compiler/mlir/tfr/ir/tfr_ops.h"
+
+int main(int argc, char **argv) {
+  tensorflow::InitMlir y(&argc, &argv);
+
+  mlir::registerAllPasses();
+
+  mlir::DialectRegistry registry;
+  registry.insert<mlir::scf::SCFDialect, mlir::TF::TensorFlowDialect,
+                  mlir::arith::ArithDialect, mlir::func::FuncDialect,
+                  mlir::shape::ShapeDialect, mlir::quant::QuantDialect,
+                  mlir::quant::ir::TFQuantDialect, mlir::TFR::TFRDialect>();
+  mlir::func::registerAllExtensions(registry);
+  return failed(mlir::MlirOptMain(argc, argv, "TFR Pass Driver\n", registry));
+}