tensorflow/compiler/mlir

606 阅读1分钟

tensorflow compiler mlir提供了各种“方言”,针对TensorFlow IR和XLA HLO IR做一个介绍。TensorFlow IR,代表 TensorFlow 计算图中的所有可能事物,XLA HLO IR则对齐XLA。

从其他的IR表示转化为mlir方言的过程一般被称为tanslate,不同的方言之间的转换一般被称为conversion或者legalization,同一种方言内的转换(optimizer pass)一般称为transforms

  • tensorflow mlir dialect

当前有两种不同的方言对 MLIR 中的 TensorFlow 图进行建模,tf dialect和tf_executor dialect,前者代表TensorFlow graph中常规的operations,后者主要用来表示一些executor相关的原语比如control dependencies,没有深入了解

下列是一个tf dialect

func @ValidConv2D(%arg0: tensor<256x32x32x3xf32>, %arg1: tensor<3x3x3x16xf32>) -> tensor<256x32x32x16xf32> {
  %0 = "tf.Conv2D"(%arg0, %arg1) {padding = "SAME", strides = [1, 1, 1, 1]} : (tensor<256x32x32x3xf32>, tensor<3x3x3x16xf32>) -> tensor<256x32x32x16xf32>
  return %0 : tensor<256x32x32x16xf32>
}