integrations/tensorflow/iree_tf_compiler/TFL/VerifyFullyConverted.cpp

// Copyright 2021 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/Utils/ConversionUtils.h"
#include "iree_tf_compiler/TFL/Passes.h"
#include "llvm/ADT/StringExtras.h"
#include "llvm/Support/FormatVariadic.h"
#include "mlir/Pass/Pass.h"
#include "mlir/Support/LLVM.h"
#include "mlir/Transforms/DialectConversion.h"
#include "tensorflow/compiler/mlir/lite/ir/tfl_ops.h"

namespace mlir {
namespace iree_integrations {
namespace TFL {

class VerifyFullyConvertedPass
    : public PassWrapper<VerifyFullyConvertedPass, FunctionPass> {
 public:
  StringRef getArgument() const override {
    return "iree-tflite-verify-fully-converted";
  }

  StringRef getDescription() const override {
    return "Verifies that all TFLite frontend ops were converted and none "
           "remain";
  }

  // Validates that no TFLite frontends ops are in the function.
  void runOnFunction() override {
    ConversionTarget target(getContext());
    target.markUnknownOpDynamicallyLegal([](Operation *) { return true; });
    target.addIllegalDialect<mlir::TFL::TensorFlowLiteDialect>();
    if (failed(
            iree_compiler::verifyAllOperationsAreLegal(getOperation(), target)))
      return signalPassFailure();
  }
};

static PassRegistration<VerifyFullyConvertedPass> pass;

std::unique_ptr<OperationPass<FuncOp>> createVerifyFullyConvertedPass() {
  return std::make_unique<VerifyFullyConvertedPass>();
}

}  // namespace TFL
}  // namespace iree_integrations
}  // namespace mlir