How to get multiple output in mlir?

MLIR
1

Hello, everyone! In my mlir function, I have two value to return.

return %4171, %4166 : tensor<1x80x32000xf32>, tensor<1x80x4096xf32>

I lower it to llvm ir.

llvm.func @_mlir_ciface_forward(%arg0: !llvm.ptr, %arg1: !llvm.ptr, %arg2: !llvm.ptr) attributes {llvm.emit_c_interface} {
    %0 = llvm.load %arg1 : !llvm.ptr -> !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
    %1 = llvm.extractvalue %0[0] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> 
    %2 = llvm.extractvalue %0[1] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> 
    %3 = llvm.extractvalue %0[2] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> 
    %4 = llvm.extractvalue %0[3, 0] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> 
    %5 = llvm.extractvalue %0[4, 0] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> 
    %6 = llvm.load %arg2 : !llvm.ptr -> !llvm.struct<(ptr, ptr, i64, array<2 x i64>, array<2 x i64>)>
    %7 = llvm.extractvalue %6[0] : !llvm.struct<(ptr, ptr, i64, array<2 x i64>, array<2 x i64>)> 
    %8 = llvm.extractvalue %6[1] : !llvm.struct<(ptr, ptr, i64, array<2 x i64>, array<2 x i64>)> 
    %9 = llvm.extractvalue %6[2] : !llvm.struct<(ptr, ptr, i64, array<2 x i64>, array<2 x i64>)> 
    %10 = llvm.extractvalue %6[3, 0] : !llvm.struct<(ptr, ptr, i64, array<2 x i64>, array<2 x i64>)> 
    %11 = llvm.extractvalue %6[3, 1] : !llvm.struct<(ptr, ptr, i64, array<2 x i64>, array<2 x i64>)> 
    %12 = llvm.extractvalue %6[4, 0] : !llvm.struct<(ptr, ptr, i64, array<2 x i64>, array<2 x i64>)> 
    %13 = llvm.extractvalue %6[4, 1] : !llvm.struct<(ptr, ptr, i64, array<2 x i64>, array<2 x i64>)> 
    %14 = llvm.call @forward(%1, %2, %3, %4, %5, %7, %8, %9, %10, %11, %12, %13) : (!llvm.ptr, !llvm.ptr, i64, i64, i64, !llvm.ptr, !llvm.ptr, i64, i64, i64, i64, i64) -> !llvm.struct<(struct<(ptr, ptr, i64, array<3 x i64>, array<3 x i64>)>, struct<(ptr, ptr, i64, array<3 x i64>, array<3 x i64>)>)>
    llvm.store %14, %arg0 : !llvm.struct<(struct<(ptr, ptr, i64, array<3 x i64>, array<3 x i64>)>, struct<(ptr, ptr, i64, array<3 x i64>, array<3 x i64>)>)>, !llvm.ptr
    llvm.return
  }

It will put two return values to !llvm.struct<(struct<(ptr, ptr, i64, array<3 x i64>, array<3 x i64>)>, struct<(ptr, ptr, i64, array<3 x i64>, array<3 x i64>)>)>. Two return values both store in one llvm.struct.
Can it split to two llvm.struct? Thank you for your help.

2

This is a matter of calling convention. You could implement your own but MLIR provides this: https://mlir.llvm.org/docs/TargetLLVMIR/#calling-conventions

3

Thank you for your answer!

4

Also note that LLVM IR, and therefore LLVM dialect, do not support returning multiple results from a function. MLIR does, so you should be able to do this in your own dialect.

1 Like