Define softmax semantics and implement a lowering to the correct sequence of named ops

[kurapov-peter]

For this, we need to lower it from multiple models and multiple frameworks and see "what they do". At least two potential options: exp(arg) and exp(arg-mean), might be more. We need to pick one, provide reasoning, and then implement a lowering upstream.

[rengolin]

I'd start by looking what TensorFlow, PyTorch, ONNX, OpenVino, oneDNN graph do currently. Look at their definitions, if possible lower a model on their graph description to MLIR and see what comes, etc.

Once we know what people will use it for, we encode that as "semantics" and lower it to that sequence of named ops.

We may need to create a new pass in addition to the legalizer. Something like linalg-decomposition.

[kurapov-peter]

Sounds good, let's start by collecting data from frameworks.

[chudur-budur]

I'd start by looking what TensorFlow, PyTorch, ONNX, OpenVino, oneDNN graph do currently. Look at their definitions, if possible lower a model on their graph description to MLIR and see what comes, etc.

Once we know what people will use it for, we encode that as "semantics" and lower it to that sequence of named ops.

We may need to create a new pass in addition to the legalizer. Something like linalg-decomposition.

I tried with TensorFlow, they have a function that spits out MLIR from a function:

import tensorflow as tf
from tensorflow.python.ops.gen_math_ops import exp as expf

expf = tf.function(expf)

cfunc = expf.get_concrete_function(tf.TensorSpec(None, tf.dtypes.float32))

mlir_str = tf.mlir.experimental.convert_function(cfunc, pass_pipeline="builtin.module(any(cse,canonicalize),convert-func-to-llvm{use-bare-ptr-memref-call-conv=1})")

print(mlir_str)

and I am getting this:

module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 1766 : i32}} {
  func.func @__inference_exp_5(%arg0: tensor<*xf32> {tf._user_specified_name = "x"}) -> tensor<*xf32> attributes {allow_soft_placement = false, tf.entry_function = {control_outputs = "", inputs = "x", outputs = "identity_RetVal"}} {
    %0 = tf_executor.graph {
      %outputs, %control = tf_executor.island wraps "tf.Exp"(%arg0) {device = ""} : (tensor<*xf32>) -> tensor<*xf32>
      %outputs_0, %control_1 = tf_executor.island wraps "tf.Identity"(%outputs) {device = ""} : (tensor<*xf32>) -> tensor<*xf32>
      tf_executor.fetch %outputs_0 : tensor<*xf32>
    }
    return %0 : tensor<*xf32>
  }
}

I am looking into how to get MLIR from tf.Exp.

[kurapov-peter]

Does their tf-standard-pipeline produce the same output? I wonder what is it exactly, too.

[chudur-budur]

Does their tf-standard-pipeline produce the same output? I wonder what is it exactly, too.

I made a mistake, you were right, the output above was with tf-standard-pipeline, now it's fixed.

The tf.Exp is an exp() based on tf dialect. https://www.tensorflow.org/mlir/tf_ops#tfexp_tfexpop

[chudur-budur]

Most likely tf.Exp is a "legacy call" OP defined here: https://cs.opensource.google/tensorflow/tensorflow/+/master:tensorflow/compiler/mlir/tensorflow/ir/tf_ops.td;drc=ff9a2058aae80dd43b722caad7db88d631ffa7ea;l=443

[rengolin]

You need a conversion program to lower tf to linalg. We used IREE last year (it was painful), not sure it still works the same way, but it's a way to get linalg at the end.

[kurapov-peter]

Since TF created mlir as a part of the framework I'd assume they have some lowering in mind? Is it IREE?

[chudur-budur]

For the tensforlow part, I need tools like tf-opt and tf-mlir-translate. To get those, I need to compile tensorflow from source.

But I am having these issues:

tensorflow/tensorflow#68247
tensorflow/tensorflow#68246

Any idea?

[rengolin]

When we did the same thing over a year ago, we also had to build TF from source, which was painful. I created a branch almost 2 years ago (so YMMV) to automate that.

However, I do recommend you use IREE directly (from pip install) and go from there. IREE is much easier to use.

[chudur-budur]

When we did the same thing over a year ago, we also had to build TF from source, which was painful. I created a branch almost 2 years ago (so YMMV) to automate that.

Thanks. But the script in that branch doesn't compile tf-apis. I managed to compile TF from source, but couldn't compile the tf-api (tf-opt, tf-mlir-translate) to a functional binary. I can compile but getting linker problems.

tensorflow/tensorflow#69367

However, I do recommend you use IREE directly (from pip install) and go from there. IREE is much easier to use.

IREE doesn't understand tf_executor dialect. Do you have any idea how do I register it to IREE?

[rengolin]

IIRC, you need to use iree-import-tf (https://iree.dev/guides/ml-frameworks/tensorflow/) and pull the IR from a saved model.

If trying TF directly, there is a tf.mlir.experimental (https://www.tensorflow.org/api_docs/python/tf/mlir/experimental/convert_graph_def) section that may help.

Once you have HLO, you should be able to use iree-opt to lower to Linalg (https://groups.google.com/g/iree-discuss/c/u2zqHKVqPTc/m/UCGeYoauGwAJ).

[rengolin]

Note, on specific help for IREE, please use their mailing list directly or Discord. It's the quickest way to get the answer you need.