How to lower linalg.softmax?

MLIR
1

Hi ! I have a mlir func with linalg.softmax.

module {
  func.func @forward(%arg0: tensor<13x13xf32>) -> tensor<13x13xf32> {
    %0 = tensor.empty() : tensor<13x13xf32>
    %1 = linalg.softmax dimension(0) ins(%arg0 : tensor<13x13xf32>) outs(%0 : tensor<13x13xf32>) -> tensor<13x13xf32>
    return %1 : tensor<13x13xf32>
  }
}

Here is my passpipeline.

        -arith-expand \
        -eliminate-empty-tensors \
        -empty-tensor-to-alloc-tensor \
        -convert-elementwise-to-linalg \
        -linalg-bufferize \
        -convert-linalg-to-affine-loops \
        -affine-loop-fusion \
        -affine-parallelize \
        -lower-affine \
        -convert-scf-to-openmp \
        -func-bufferize \
        -arith-bufferize \
        -tensor-bufferize \
        -buffer-deallocation \
        -finalizing-bufferize \
        -convert-vector-to-scf \
        -expand-strided-metadata \
        -convert-vector-to-llvm \
        -memref-expand \
        -arith-expand \
        -convert-arith-to-llvm \
        -finalize-memref-to-llvm \
        -convert-scf-to-cf \
        -llvm-request-c-wrappers \
        -convert-openmp-to-llvm \
        -convert-math-to-llvm \
        -convert-math-to-libm  \
        -convert-func-to-llvm \
        -reconcile-unrealized-casts \
	-mlir-print-ir-before-all \

It will finally result this.

// -----// IR Dump Before FinalizingBufferize (finalizing-bufferize) //----- //
func.func @forward(%arg0: memref<13x13xf32>) -> memref<13x13xf32> {
  %0 = bufferization.to_tensor %arg0 : memref<13x13xf32>
  %alloc = memref.alloc() {alignment = 64 : i64} : memref<13x13xf32>
  %1 = bufferization.to_tensor %alloc : memref<13x13xf32>
  memref.dealloc %alloc : memref<13x13xf32>
  %2 = linalg.softmax dimension(0) ins(%0 : tensor<13x13xf32>) outs(%1 : tensor<13x13xf32>) -> tensor<13x13xf32>
  %3 = bufferization.to_memref %2 : memref<13x13xf32>
  return %3 : memref<13x13xf32>
}
error: failed to legalize operation 'linalg.softmax'

I think it is because linalg.softmax didn’t lower. Are there some wrongs in my pipeline? Thank you for your help.

2

Softmax implements the AggregatedOpInterface and you can decompose it into a sequence of simpler operations. You can use this method: llvm-project/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp at 81cedac8f60cf5a0c2c3724a8260d46792b9d637 · llvm/llvm-project · GitHub.

Have a look a the following commit: [mlir][Linalg] Add an interface to decompose complex ops · llvm/llvm-project@9be8219 · GitHub there are some examples on how to decompose the operation using the transform dialect, if you want to use a pass you can have a look at: tpp-mlir/lib/TPP/Transforms/DecomposeAggregatedOps.cpp at e18be6965896cef85f4fa3d233e40df7a27e1bee · plaidml/tpp-mlir · GitHub

1 Like
3

Thank you for your help! I will try it.