Running Triton in IREE

[parsifal-47]

Hi All,
I am trying to run Triton kernels in IREE, for that, I get arith+linalg IR from https://github.com/microsoft/triton-shared , bufferize it with:

mlir-opt --one-shot-bufferize=allow-return-allocs-from-loops=true --reconcile-unrealized-casts ttshared.mlir -o ttshared_bufferized.mlir

but when I run:

./tools/iree-compile ttshared_bufferized.mlir --output-format=vm-bytecode -o tt.vmfb
ttshared_bufferized.mlir:3:3: error: failed to legalize unresolved materialization from ('!vm.ref<memref<*xf32>>') to ('memref<*xf32>') that remained live after conversion
...

the original code contains no !vm.ref<memref<*xf32>> that is likely something internal to IREE, do you have an advice on how to prepare the IR to be compilable with IREE?

Complete code below, thank you!

#map = affine_map<(d0) -> (d0)>
module {
  func.func @add_kernel(%arg0: memref<*xf32> {tt.divisibility = 16 : i32}, %arg1: memref<*xf32> {tt.divisibility = 16 : i32}, %arg2: memref<*xf32> {tt.divisibility = 16 : i32}, %arg3: i32 {tt.divisibility = 16 : i32}, %arg4: i32, %arg5: i32, %arg6: i32, %arg7: i32, %arg8: i32, %arg9: i32) {
    %c1024_i32 = arith.constant 1024 : i32
    %c1024 = arith.constant 1024 : index
    %0 = arith.muli %arg7, %c1024_i32 : i32
    %1 = arith.index_cast %0 : i32 to index
    %reinterpret_cast = memref.reinterpret_cast %arg0 to offset: [%1], sizes: [1024], strides: [1] : memref<*xf32> to memref<1024xf32, strided<[1], offset: ?>>
    %2 = arith.addi %1, %c1024 : index
    %3 = arith.index_cast %arg3 : i32 to index
    %4 = arith.minsi %2, %3 : index
    %5 = arith.maxsi %4, %1 : index
    %6 = arith.subi %5, %1 : index
    %alloc = memref.alloc() : memref<1024xf32>
    %subview = memref.subview %reinterpret_cast[0] [%6] [1] : memref<1024xf32, strided<[1], offset: ?>> to memref<?xf32, strided<[1], offset: ?>>
    %subview_0 = memref.subview %alloc[0] [%6] [1] : memref<1024xf32> to memref<?xf32, strided<[1]>>
    memref.copy %subview, %subview_0 : memref<?xf32, strided<[1], offset: ?>> to memref<?xf32, strided<[1]>>
    %reinterpret_cast_1 = memref.reinterpret_cast %arg1 to offset: [%1], sizes: [1024], strides: [1] : memref<*xf32> to memref<1024xf32, strided<[1], offset: ?>>
    %alloc_2 = memref.alloc() : memref<1024xf32>
    %subview_3 = memref.subview %reinterpret_cast_1[0] [%6] [1] : memref<1024xf32, strided<[1], offset: ?>> to memref<?xf32, strided<[1], offset: ?>>
    %subview_4 = memref.subview %alloc_2[0] [%6] [1] : memref<1024xf32> to memref<?xf32, strided<[1]>>
    memref.copy %subview_3, %subview_4 : memref<?xf32, strided<[1], offset: ?>> to memref<?xf32, strided<[1]>>
    linalg.generic {indexing_maps = [#map, #map, #map], iterator_types = ["parallel"]} ins(%alloc, %alloc_2 : memref<1024xf32>, memref<1024xf32>) outs(%alloc : memref<1024xf32>) {
    ^bb0(%in: f32, %in_8: f32, %out: f32):
      %7 = arith.addf %in, %in_8 : f32
      linalg.yield %7 : f32
    }
    %reinterpret_cast_5 = memref.reinterpret_cast %arg2 to offset: [%1], sizes: [1024], strides: [1] : memref<*xf32> to memref<1024xf32, strided<[1], offset: ?>>
    %subview_6 = memref.subview %alloc[0] [%6] [1] : memref<1024xf32> to memref<?xf32, strided<[1]>>
    %subview_7 = memref.subview %reinterpret_cast_5[0] [%6] [1] : memref<1024xf32, strided<[1], offset: ?>> to memref<?xf32, strided<[1], offset: ?>>
    memref.copy %subview_6, %subview_7 : memref<?xf32, strided<[1]>> to memref<?xf32, strided<[1], offset: ?>>
    return
  }
}

[ScottTodd]

The * in memref<*xf32> indicates "unranked": https://mlir.llvm.org/docs/Dialects/Builtin/#unrankedmemreftype. IREE only supports ranked tensors (and I suspect memrefs), though ranked tensors can have dynamic shapes.

For example,

• tensor<?xf32> has one (unbounded) dynamic dimension and could represent for example tensor<8xf32> or tensor<64xf32> but not tensor<4x4xf32>
• tensor<?x?xf32> has two dynamic dimensions and could be tensor<4x4xf32> but not tensor<4xf32> or tensor<4x4x4xf32>
• tensor<*xf32> is unranked and could be tensor<4x4x4xf32>, tensor<8xf32> or mostly anything else

Can you produce programs with known ranks? I see some notes on UnkrankedMemRefType that could be helpful: https://mlir.llvm.org/docs/Dialects/Builtin/#unrankedmemreftype

The purpose of unranked memrefs is to allow external library functions to receive memref arguments of any rank without versioning the functions based on the rank. Other uses of this type are disallowed or will have undefined behavior.

It looks like the ranks and shapes are actually known in your program there, based on the memref.reinterpret_cast calls, but IREE might still be trying to respect the unranked types in the function arguments.

[parsifal-47]

Thank you for detailed answer, yes, I would likely be able to produce ranked tensor.

[parsifal-47]

Follow-up question, I tried to go from unranked to unbounded *xf32 -> ?xf32, but I still see the same error:

ttshared_bufferized.mlir:3:3: error: failed to legalize unresolved materialization from ('!vm.ref<memref<?xf32>>') to ('memref<?xf32>') that remained live after conversion
  func.func @add_kernel(%arg0: memref<?xf32> {tt.divisibility = 16 : i32}, %arg1: memref<?xf32> {tt.divisibility = 16 : i32}, %arg2: memref<?xf32> {tt.divisibility = 16 : i32}, %arg3: i32 {tt.divisibility = 16 : i32}, %arg4: i32, %arg5: i32, %arg6: i32, %arg7: i32, %arg8: i32, %arg9: i32) {

do I need to have static shapes or now I miss something else?

Thank you!

[ScottTodd]

Ah, that might be tricky. Probably want someone more familiar with codegen (@MaheshRavishankar ?) or Triton (@antiagainst ?) to weigh in. I suspect that IREE is expecting "linalg on tensors" input and this input is "linalg on memrefs". Is there a way to convert between those formats or have your input program constructed on tensors? IREE treats memrefs as a low level construct that each backend can use as needed, while backends prefer to start with abstract representations instead of concrete memory locations.