[Mlir-commits] [mlir] [mlir][vector] Clarify the semantics of gather/scatter indexing (nfc) (PR #181357)

[Andrzej Warzyński]

banach-space wrote:

> @banach-space Re your question about Linalg, the code I'm thinking of - which is doing a bunch of multiplies - is
> 

Thanks! I am familiar with that implementation - I wrote it. I don't quite see where the consistency is? Why is the presence of multiplies significant here?

To make the discussion a bit more concrete, here is an example:
```mlir
func.func @example(%src:tensor<24x80x16xf32>, %offset: index, %output : tensor<2x1x4xf32>) -> tensor<2x1x4xf32> {
  %c13 = arith.constant 13 : index
  %1 = linalg.generic {
    indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1, d2)>],
    iterator_types = ["parallel", "parallel", "parallel"]
  } outs(%output : tensor<2x1x4xf32>) {
  ^bb0(%out: f32):

    %idx_0 = linalg.index 0 : index
    %idx_0_p_2 = affine.apply affine_map<(d0) -> (d0 + 2)>(%idx_0)

    %idx_2 = linalg.index 2 : index
    %idx_2_p_off = affine.apply affine_map<(d0, d1) -> (d0 + d1)>(%idx_2, %offset)

    %val = tensor.extract %src[%idx_0_p_2, %idx_2_p_off, %c13] : tensor<24x80x16xf32>

    linalg.yield %val : f32
  } -> tensor<2x1x4xf32>
  return %1 : tensor<2x1x4xf32>
}

module @transforms attributes { transform.with_named_sequence } {
  transform.named_sequence @__transform_main(%module: !transform.any_op {transform.readonly}) {

    %0 = transform.structured.match ops{["linalg.generic"]} in %module : (!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
   }
}
```

It vectorizes as:
```mlir
  func.func @vectorize_nd_tensor_extract_with_affine_apply_gather(%arg0: tensor<24x80x16xf32>, %arg1: index, %arg2: tensor<2x1x4xf32>) -> tensor<2x1x4xf32> {
    %cst = arith.constant dense<13> : vector<2x1x4xindex>
    %cst_0 = arith.constant dense<16> : vector<2x1x4xindex>
    %cst_1 = arith.constant dense<[160, 240]> : vector<2xindex>
    %c0 = arith.constant 0 : index
    %cst_2 = arith.constant dense<0.000000e+00> : vector<2x1x4xf32>
    %cst_3 = arith.constant dense<true> : vector<2x1x4xi1>
    %cst_4 = arith.constant dense<[0, 1, 2, 3]> : vector<4xindex>
    %0 = vector.broadcast %arg1 : index to vector<4xindex>
    %1 = arith.addi %0, %cst_4 : vector<4xindex>
    %2 = vector.broadcast %cst_1 : vector<2xindex> to vector<4x1x2xindex>
    %3 = vector.transpose %2, [2, 1, 0] : vector<4x1x2xindex> to vector<2x1x4xindex>
    %4 = vector.broadcast %1 : vector<4xindex> to vector<2x1x4xindex>
    %5 = arith.addi %4, %3 : vector<2x1x4xindex>
    %6 = arith.muli %5, %cst_0 : vector<2x1x4xindex>
    %7 = arith.addi %6, %cst : vector<2x1x4xindex>
    %8 = vector.gather %arg0[%c0, %c0, %c0] [%7], %cst_3, %cst_2 : tensor<24x80x16xf32>, vector<2x1x4xindex>, vector<2x1x4xi1>, vector<2x1x4xf32> into vector<2x1x4xf32>
    %9 = vector.transfer_write %8, %arg2[%c0, %c0, %c0] {in_bounds = [true, true, true]} : vector<2x1x4xf32>, tensor<2x1x4xf32>
    return %9 : tensor<2x1x4xf32>
  }
```

This looks OK to me - a vector of indices, one per output element in the output vector, is fed into `vector.gather`:
```mlir
%7 = arith.addi %6, %cst : vector<2x1x4xindex>
%8 = vector.gather %arg0[%c0, %c0, %c0] [%7], %cst_3, %cst_2 : tensor<24x80x16xf32>, vector<2x1x4xindex>, vector<2x1x4xi1>, vector<2x1x4xf32> into vector<2x1x4xf32>
```

Could you point me at the inconsistency that you are referring to? That would really help.

Let me also note, the documentation of `vector.gather` and `vector.scatter` pre-dates the implementation in the Linalg vectorizer. It is for the vectorizer to implement what is in the docs, rather than the docs reflect what the vectorizer does.  If that is not the case, I'd be happy to fix that.

https://github.com/llvm/llvm-project/pull/181357