vmvx: error: failed to legalize operation 'memref.alloca' that was explicitly marked illegal

[silvasean]

Describe the bug

error: failed to legalize operation 'memref.alloca' that was explicitly marked illegal

FYI, the reproducer is the same as in #9126 cc @hanhanW for insights.

Full log: https://gist.github.com/silvasean/72957031fea1be083f0172508f94358b

Note: this blocks running BERT on IREE's vmvx backend!

To Reproduce

iree-compile --iree-hal-target-backends=vmvx  /tmp/iree-temps/core-input.mlir

#map0 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>
#map1 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, 0)>
module attributes {torch.debug_module_name = "OnlyLogitsHuggingFaceModel"} {
  func.func @forward(%arg0: tensor<1x12xi64>, %arg1: tensor<1x12x12x12xf32>) -> tensor<1x12x12x1xf32> {
    %c0_i64 = arith.constant 0 : i64
    %cst = arith.constant -3.40282347E+38 : f32
    %0 = linalg.init_tensor [1, 12, 12, 1] : tensor<1x12x12x1xi64>
    %1 = linalg.fill ins(%c0_i64 : i64) outs(%0 : tensor<1x12x12x1xi64>) -> tensor<1x12x12x1xi64>
    %2 = linalg.init_tensor [1, 12, 12, 1] : tensor<1x12x12x1xf32>
    %3 = linalg.fill ins(%cst : f32) outs(%2 : tensor<1x12x12x1xf32>) -> tensor<1x12x12x1xf32>
    %4:2 = linalg.generic {indexing_maps = [#map0, #map1, #map1], iterator_types = ["parallel", "parallel", "parallel", "reduction"]} ins(%arg1 : tensor<1x12x12x12xf32>) outs(%3, %1 : tensor<1x12x12x1xf32>, tensor<1x12x12x1xi64>) {
    ^bb0(%arg2: f32, %arg3: f32, %arg4: i64):
      %5 = linalg.index 3 : index
      %6 = arith.index_cast %5 : index to i64
      %7 = arith.cmpf ogt, %arg2, %arg3 : f32
      %8 = arith.select %7, %arg2, %arg3 : f32
      %9 = arith.select %7, %6, %arg4 : i64
      linalg.yield %8, %9 : f32, i64
    } -> (tensor<1x12x12x1xf32>, tensor<1x12x12x1xi64>)
    return %4#0 : tensor<1x12x12x1xf32>
  }
}

[hanhanW]

This is an argmin/max op. It's a dup of #8899

We will need a stack allocation, which is bounded. Other backends allocate bounded static buffer to work it around. For vmvx, I think we're not going to make vmvx accept stack allocation. For now there is no way around it.

[silvasean]

Hmm... that kind of defeats the purpose of VMVX being a reference backend... will bring this up.

[benvanik]

vmvx can easily handle stack allocs like this

[benvanik]

https://github.com/google/iree/blob/41a2cebe9560fda8a3ef171bb2dd5f51434897f2/compiler/src/iree/compiler/Dialect/VM/Conversion/MemRefToVM/ConvertMemRefToVM.cpp
just needs a pattern from memref.alloca to IREE::VM::BufferAllocOp

[hanhanW]

cool, if we accept bounded stack allocs for vmvx backend, I can take a look at it.

[benvanik]

Workgroup local memory is always going to be better (on all implementations), but that's a bigger change - having memref.alloc/alloca -> vm.buffer.alloc is always useful though and I think that'd unblock all this.