iree: spirv: error: 'spv.Store' op mismatch in result type and pointer type

Describe the bug

-:17:12: error: 'spv.Store' op mismatch in result type and pointer type
-:17:12: note: see current operation: "spv.Store"(%9, %0) : (!spv.ptr<i64, StorageBuffer>, i32) -> ()

Note: this is the same reproducer as https://github.com/google/iree/issues/9155 and https://github.com/google/iree/issues/9126 cc @hanhanW for insights.

Full error log: https://gist.github.com/silvasean/4136eaa19f12910fc452ec76622bd157

Note: this blocks running BERT on IREE’s vulkan backend!

To Reproduce

iree-compile --iree-hal-target-backends=vulkan  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>
  }
}