How to get mlir-opt to parse tf_executor dialect in mlir file?

Beginners
1

I used tf-mlir-translate to translate a tensorflow graphdef into mlir file - which contains tf graph embedded in tf_executor dialect as shown below:

module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 0 : i32}} {
  func @main() {
    tf_executor.graph {
       ...

How can I tell mlir-opt to interpret it using tf_executor and tf dialects ?

I am trying to write a simple pass by extending a module pass to walk over all functions and Operations to process each Operation. I use mlir-opt --<my_registered_pass> to test my code. But I am hitting a parse error

error: custom op 'tf_executor.graph' is unknown
    tf_executor.graph {
    ^

here’s the call stack

(anonymous namespace)::ModuleParser::parseModule((anonymous namespace)::ModuleParser * const this, mlir::ModuleOp module) (<path>//llvm-project/mlir/lib/Parser/Parser.cpp:4655)
mlir::parseSourceFile(const llvm::SourceMgr & sourceMgr, mlir::MLIRContext * context) (<path>/llvm-project/mlir/lib/Parser/Parser.cpp:4720)
performActions(llvm::raw_ostream & os, bool verifyDiagnostics, bool verifyPasses, llvm::SourceMgr & sourceMgr, mlir::MLIRContext * context, const mlir::PassPipelineCLParser & passPipeline) (<path>/llvm-project/mlir/lib/Support/MlirOptMain.cpp:44)
processBuffer(llvm::raw_ostream & os, std::unique_ptr<llvm::MemoryBuffer, std::default_delete<llvm::MemoryBuffer> > ownedBuffer, bool verifyDiagnostics, bool verifyPasses, const mlir::PassPipelineCLParser & passPipeline) (<path>/llvm-project/mlir/lib/Support/MlirOptMain.cpp:83)
mlir::MlirOptMain(llvm::raw_ostream & os, std::unique_ptr<llvm::MemoryBuffer, std::default_delete<llvm::MemoryBuffer> > buffer, const mlir::PassPipelineCLParser & passPipeline, bool splitInputFile, bool verifyDiagnostics, bool verifyPasses) (<path>/llvm-project/mlir/lib/Support/MlirOptMain.cpp:116)
main(int argc, char ** argv) (<path>/llvm-project/mlir/tools/mlir-opt/mlir-opt.cpp:74)
2

The TensorFlow dialect does not come with MLIR, it lives in the TensorFlow directory. There you can use tf-opt (just the same way you got tf-mlir-translate) that provides a superset of mlir-opt and includes the TensorFlow specific bits needed to process TensorFlow MLIR files.

3

Thanks for quick response. I will try it out with tf-opt.

4

Elaborating more, as I have been trying to do something similar and this may help someone else in the future.

The command line instructions you require after you have the TensorFlow source code are:

# cat add.pbtxt
node {
  name: "Add"
  op: "Add"
  input: "input0"
  input: "input1"
  attr {
    key: "T"
    value {
      type: DT_INT32
    }
  }
}
node {
  name: "input0"
  op: "Placeholder"
  attr {
    key: "dtype"
    value {
      type: DT_INT32
    }
  }
}
node {
  name: "input1"
  op: "Placeholder"
  attr {
    key: "dtype"
    value {
      type: DT_INT32
    }
  }
}
versions {
  producer: 27
}

# pb model to mlir tf dialect
bazel run //tensorflow/compiler/mlir:tf-mlir-translate -- \
    --graphdef-to-mlir --tf-input-arrays=input0,input1 \
    --tf-input-data-types=DT_INT32,DT_INT32 \
    --tf-input-shapes=10:10 --tf-output-arrays=Add \
    ${PWD}/add.pbtxt \
    -o ${PWD}/add.mlir

# Convert to mhlo
bazel run //tensorflow/compiler/mlir:tf-opt -- \
    --tf-executor-to-functional-conversion \
    --tf-region-control-flow-to-functional \
    --xla-legalize-tf \
    ${PWD}/add.mlir \
    -o ${PWD}/add-mhlo.mlir

# Convert to lhlo
bazel run //tensorflow/compiler/mlir:tf-opt -- \
    --xla-hlo-to-lhlo-with-xla \
    ${PWD}/add-mhlo.mlir \
    -o ${PWD}/add-lhlo.mlir

# Convert from lhlo to mlir.llvm.
#   There may be a need for other passes depending on the model
bazel run //tensorflow/compiler/mlir/hlo:mlir-hlo-opt -- \
    -mhlo-test-chlo-legalize-to-hlo \
    -hlo-legalize-to-lhlo=results-escape-function=true \
    -buffer-placement \
    -copy-removal \
    -canonicalize \
    -cse \
    -lhlo-legalize-to-linalg \
    -lhlo-fuse-linalg \
    -convert-linalg-to-loops \
    -convert-scf-to-std \
    -canonicalize \
    -cse \
    -test-lhlo-legalize-to-llvm  \
    ${PWD}/add-lhlo.mlir \
    -o ${PWD}/add-llvm.mlir
5

Newbie here, I am having difficulty lowering Constants using this pipeline.

If I have a frozen model with the weights encoded as constants, the constants are not lowered below lmhlo.const instructions. It feels like I am missing a pass, can anyone provide any hints to where I should direct my search queries.

6

no argumnet -lhlo-legalize-to-linalg of mlir-hlo-opt? changed?

7

Hi there, did you manage to parse tf/tf_executor dialect to HLO? I found the same issues as yours and even though I used tf-opt, the output file remained unchanged.