[MLIR][linalg][vectorization] failure for an nd extract case

[nirvedhmeshram]

For the following IR

#map = affine_map<(d0, d1) -> (d0, d1)>
#map1 = affine_map<(d0, d1, d2) -> (d0 + d1 + d2)>
func.func @gather_failure(%arg0: tensor<8x128x768xf32>, %arg2: tensor<8x1xf32>, %arg3 : index) -> tensor<8x1xf32> {
  %c0 = arith.constant 0 : index
  %1 = linalg.generic {
    indexing_maps = [#map], 
    iterator_types = ["parallel", "parallel"]
  } outs(%arg2 : tensor<8x1xf32>) {
  ^bb0(%arg5: f32):
      %2 = linalg.index 0 : index
      %3 = linalg.index 1 : index
      %4 = affine.apply #map1(%arg3, %3, %arg3)
    %extracted = tensor.extract %arg0[%2, %c0, %4] : tensor<8x128x768xf32>
    linalg.yield %extracted : f32
  } -> tensor<8x1xf32>
  return %1 : tensor<8x1xf32>
}

module attributes {transform.with_named_sequence} {
  transform.named_sequence @__transform_main(%arg2: !transform.any_op {transform.readonly}) {
    %0 = transform.structured.match ops{["linalg.generic"]} in %arg2 : (!transform.any_op) -> !transform.any_op
    %1 = transform.get_parent_op %0 {isolated_from_above} : (!transform.any_op) -> !transform.any_op
    %2 = transform.structured.vectorize_children_and_apply_patterns %1 {vectorize_nd_extract} : (!transform.any_op) -> !transform.any_op
    transform.yield
  }
}

run

mlir-opt -transform-interpreter -split-input-file test.mlir 

It gives the error

within split at test_gather_core.mlir:1 offset :13:18: error: 'vector.shape_cast' op source/result number of elements must match
    %extracted = tensor.extract %arg0[%2, %c0, %4] : tensor<8x128x768xf32>
                 ^
within split at test_gather_core.mlir:1 offset :13:18: note: see current operation: %7 = "vector.shape_cast"(%5) : (vector<8x1xindex>) -> vector<1xindex>

with verify-each=0 you can see the following output IR

#map = affine_map<(d0, d1, d2) -> (d1, d2)>
#map1 = affine_map<(d0, d1) -> (d0, d1)>
"builtin.module"() ({
  "func.func"() <{function_type = (tensor<8x128x768xf32>, tensor<8x1xf32>, index) -> tensor<8x1xf32>, sym_name = "gather_failure"}> ({
  ^bb0(%arg1: tensor<8x128x768xf32>, %arg2: tensor<8x1xf32>, %arg3: index):
    %3 = "arith.constant"() <{value = 0.000000e+00 : f32}> : () -> f32
    %4 = "arith.constant"() <{value = 0 : i32}> : () -> i32
    %5 = "arith.constant"() <{value = 0 : index}> : () -> index
    %6 = "arith.constant"() <{value = dense<[0, 1, 2, 3, 4, 5, 6, 7]> : vector<8xindex>}> : () -> vector<8xindex>
    %7 = "vector.broadcast"(%6) : (vector<8xindex>) -> vector<1x8xindex>
    %8 = "vector.transpose"(%7) <{permutation = array<i64: 1, 0>}> : (vector<1x8xindex>) -> vector<8x1xindex>
    %9 = "arith.addi"(%arg3, %arg3) <{overflowFlags = #arith.overflow<none>}> : (index, index) -> index
    %10 = "vector.shape_cast"(%8) : (vector<8x1xindex>) -> vector<1xindex>
    %11 = "vector.extractelement"(%10, %4) : (vector<1xindex>, i32) -> index
    %12 = "vector.transfer_read"(%arg1, %11, %5, %9, %3) <{in_bounds = [true, true], operandSegmentSizes = array<i32: 1, 3, 1, 0>, permutation_map = #map}> : (tensor<8x128x768xf32>, index, index, index, f32) -> vector<8x1xf32>
    %13 = "vector.transfer_write"(%12, %arg2, %5, %5) <{in_bounds = [true, true], operandSegmentSizes = array<i32: 1, 1, 2, 0>, permutation_map = #map1}> : (vector<8x1xf32>, tensor<8x1xf32>, index, index) -> tensor<8x1xf32>
    "func.return"(%13) : (tensor<8x1xf32>) -> ()
  }) : () -> ()

[nirvedhmeshram]

cc @MaheshRavishankar @qedawkins

[MaheshRavishankar]

@banach-space can you provide some guidance here?

[banach-space]

Sorry that you are hitting this - looks like a vectorisation bug and most definitely my fault. I am OoO this and the coming week, so apologies if replying with a delay.

This shape_cast is clearly wrong:

%10 = "vector.shape_cast"(%8) : (vector<8x1xindex>) -> vector<1xindex>

It should be this instead:

%10 = "vector.shape_cast"(%8) : (vector<8x1xindex>) -> vector<8xindex>

These Ops are created when generating indices for vector.gather. It happens here:

llvm-project/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp

Line 1082 in ddf40e0

auto indexAs1dVector = rewriter.create<vector::ShapeCastOp>(

Hope this helps :)

[nirvedhmeshram]

Thanks for the pointers @banach-space , looking at the code around here and the other test cases, I think the problem is that we have a

"vector.transpose"(%7) <{permutation = array<i64: 1, 0>}> : (vector<1x8xindex>) -> vector<8x1xindex>

between the vector.broadcast and transferread indices
so the assumption of

  // Collect indices for `vector.transfer_read`. At this point, the indices will
  // either be scalars or would have been broadcast to vectors matching the
  // result type. 

is untrue.
The other thing I want to ask is do we really want the provided IR classified as a VectorMemoryAccessKind::Contiguous
if it were classified as a VectorMemoryAccessKind::Gather things would workout?

[nirvedhmeshram]

I am able to support the contiguous load lowering by checking for transpose case and generate this IR

module {
  func.func @gather_failure(%arg0: tensor<8x128x768xf32>, %arg1: tensor<1x?x8x1xf32>, %arg2: index) -> tensor<8x1xf32> {
    %cst = arith.constant dense<[0, 1, 2, 3, 4, 5, 6, 7]> : vector<8xindex>
    %c0_i32 = arith.constant 0 : i32
    %cst_0 = arith.constant 0.000000e+00 : f32
    %c0 = arith.constant 0 : index
    %0 = tensor.empty() : tensor<8x1xf32>
    %1 = vector.broadcast %cst : vector<8xindex> to vector<1x8xindex>
    %2 = vector.transpose %1, [1, 0] : vector<1x8xindex> to vector<8x1xindex>
    %3 = arith.addi %arg2, %arg2 : index
    %4 = vector.shape_cast %2 : vector<8x1xindex> to vector<8xindex>
    %5 = vector.extractelement %4[%c0_i32 : i32] : vector<8xindex>
    %6 = vector.transfer_read %arg0[%5, %c0, %3], %cst_0 {in_bounds = [true, true]} : tensor<8x128x768xf32>, vector<8x1xf32>
    %7 = vector.transfer_write %6, %0[%c0, %c0] {in_bounds = [true, true]} : vector<8x1xf32>, tensor<8x1xf32>
    return %7 : tensor<8x1xf32>
  }
}

Edit: After looking at the above "correct" IR. I am seeing the issue that %2 = vector.transpose %1, [1, 0] : vector<1x8xindex> to vector<8x1xindex> seems unnecessary to begin with. Now I am investigating why its added in the first place.

[nirvedhmeshram]

@banach-space let me know if #107922 is an okay way to handle this.

[llvmbot]

@llvm/issue-subscribers-mlir-linalg

Author: Nirvedh Meshram (nirvedhmeshram)

For the following IR ``` #map = affine_map<(d0, d1) -> (d0, d1)> #map1 = affine_map<(d0, d1, d2) -> (d0 + d1 + d2)> func.func @gather_failure(%arg0: tensor<8x128x768xf32>, %arg2: tensor<8x1xf32>, %arg3 : index) -> tensor<8x1xf32> { %c0 = arith.constant 0 : index %1 = linalg.generic { indexing_maps = [#map], iterator_types = ["parallel", "parallel"] } outs(%arg2 : tensor<8x1xf32>) { ^bb0(%arg5: f32): %2 = linalg.index 0 : index %3 = linalg.index 1 : index %4 = affine.apply #map1(%arg3, %3, %arg3) %extracted = tensor.extract %arg0[%2, %c0, %4] : tensor<8x128x768xf32> linalg.yield %extracted : f32 } -> tensor<8x1xf32> return %1 : tensor<8x1xf32> }

module attributes {transform.with_named_sequence} {
transform.named_sequence @__transform_main(%arg2: !transform.any_op {transform.readonly}) {
%0 = transform.structured.match ops{["linalg.generic"]} in %arg2 : (!transform.any_op) -> !transform.any_op
%1 = transform.get_parent_op %0 {isolated_from_above} : (!transform.any_op) -> !transform.any_op
%2 = transform.structured.vectorize_children_and_apply_patterns %1 {vectorize_nd_extract} : (!transform.any_op) -> !transform.any_op
transform.yield
}
}

run

mlir-opt -transform-interpreter -split-input-file test.mlir

It gives the error

within split at test_gather_core.mlir:1 offset :13:18: error: 'vector.shape_cast' op source/result number of elements must match
%extracted = tensor.extract %arg0[%2, %c0, %4] : tensor<8x128x768xf32>
^
within split at test_gather_core.mlir:1 offset :13:18: note: see current operation: %7 = "vector.shape_cast"(%5) : (vector<8x1xindex>) -> vector<1xindex>

with `verify-each=0` you can see the following output IR

#map = affine_map<(d0, d1, d2) -> (d1, d2)>
#map1 = affine_map<(d0, d1) -> (d0, d1)>
"builtin.module"() ({
"func.func"() <{function_type = (tensor<8x128x768xf32>, tensor<8x1xf32>, index) -> tensor<8x1xf32>, sym_name = "gather_failure"}> ({
^bb0(%arg1: tensor<8x128x768xf32>, %arg2: tensor<8x1xf32>, %arg3: index):
%3 = "arith.constant"() <{value = 0.000000e+00 : f32}> : () -> f32
%4 = "arith.constant"() <{value = 0 : i32}> : () -> i32
%5 = "arith.constant"() <{value = 0 : index}> : () -> index
%6 = "arith.constant"() <{value = dense<[0, 1, 2, 3, 4, 5, 6, 7]> : vector<8xindex>}> : () -> vector<8xindex>
%7 = "vector.broadcast"(%6) : (vector<8xindex>) -> vector<1x8xindex>
%8 = "vector.transpose"(%7) <{permutation = array<i64: 1, 0>}> : (vector<1x8xindex>) -> vector<8x1xindex>
%9 = "arith.addi"(%arg3, %arg3) <{overflowFlags = #arith.overflow<none>}> : (index, index) -> index
%10 = "vector.shape_cast"(%8) : (vector<8x1xindex>) -> vector<1xindex>
%11 = "vector.extractelement"(%10, %4) : (vector<1xindex>, i32) -> index
%12 = "vector.transfer_read"(%arg1, %11, %5, %9, %3) <{in_bounds = [true, true], operandSegmentSizes = array<i32: 1, 3, 1, 0>, permutation_map = #map}> : (tensor<8x128x768xf32>, index, index, index, f32) -> vector<8x1xf32>
%13 = "vector.transfer_write"(%12, %arg2, %5, %5) <{in_bounds = [true, true], operandSegmentSizes = array<i32: 1, 1, 2, 0>, permutation_map = #map1}> : (vector<8x1xf32>, tensor<8x1xf32>, index, index) -> tensor<8x1xf32>
"func.return"(%13) : (tensor<8x1xf32>) -> ()
}) : () -> ()

</details>

[nirvedhmeshram]

Fixed by #107922