llvm / llvm/llvm-project

[mlir][vector] `vector.gather` on a strided memref currently lowers inconsistently

Open
#187,215 6 comments 0 reactions 0 assignees View on GitHub
mlir:vector
Dominant language
LLVM
Stars
40.5k
Forks
18.7k
PR merge metrics
PR metrics pending

Description

**PROBLEM**

There are two ways to lower `vector.gather` to LLVM:
* directly to LLVM in `--test-lower-to-llvm`,
* indirectly via `--test-vector-gather-lowering` ( `vector.gather` -> `vector.load`) , then `--test-lower-to-llvm`.

For the same source program, these two lowering paths produce different runtime results on a **non-identity-layout** memref. Since both are lowerings of the same op, they should agree. This suggests that the semantics of `vector.gather` on strided memrefs are currently underspecified, or that one lowering is inconsistent with the intended semantics.

This affects `vector.scatter` as well.

Discussion based on LLVM SHA: bc54aeff7445028f2b23713a2ba43392bae3d4e.

**ACTUAL DISCREPANCY**

For the reproducer below, the two lowering paths produce:
* direct lowering to LLVM: (0, 1, 2, 3, 4, 10, 11, 12)
* lowering via `--test-vector-gather-lowering`: (0, 1, 2, 10, 11, 12, 20, 21)

**CONTEXT**

Example operation:
```mlir
// %base is a 3x3 view into a 10x5 MemRef
vector.gather %base[%c0, %c0][%indices], %mask, %pass_thru : memref>, vector<8xindex>, vector<8xi1>, vector<8xf32> into vector<8xf32>
```

The key question is:
* _What are the semantics of vector.gather on a strided memref?_

The discrepancy seems to come from two possible interpretations:
* gather from the logical element sequence of the memref view;
* gather from a contiguous sequence in the underlying storage, starting from the computed base address.

Today, the two lowering paths appear to implement different interpretations.

**REPRODUCER**

Based on https://github.com/llvm/llvm-project/blob/main/mlir/test/Integration/Dialect/Vector/CPU/gather.mlir.

> NOTE: The current upstream version of the test will only check MemRef(s) with **identity** layouts. @krzysz00 kindly provided a snippet (in this [comment](https://github.com/llvm/llvm-project/pull/187071#issuecomment-4077846358)) with non-identity layout. I incorporated that into the original test.

```mlir
// DEFINE: %{entry_point} = main
// DEFINE: %{run} = mlir-runner -e entry -entry-point-result=void \
// DEFINE: -shared-libs=%native_mlir_runner_utils,%native_mlir_c_runner_utils

/// TEST 1. Verify default compilation (direct lowering of `vector.gather` to LLVM)
// DEFINE: %{compile} = mlir-opt %s -test-lower-to-llvm
// RUN: %{compile} | %{run} | FileCheck %s

/// TEST 2. Verify compilation via `test-vector-gather-lowering` (`vector.gather`
/// lowerd to LLVM via `vector.load`)
// REDEFINE: %{compile} = mlir-opt %s --test-vector-gather-lowering | mlir-opt -test-lower-to-llvm
// RUN: %{compile} | %{run} | FileCheck %s

//===----------------------------------------------------------------------===//
// @gather8_strided
//
// Convenience wrapper
//===----------------------------------------------------------------------===//
func.func @gather8_strided(%base: memref>, %indices: vector<8xindex>,
%mask: vector<8xi1>, %pass_thru: vector<8xf32>) -> vector<8xf32> {
%c0 = arith.constant 0: index
%g = vector.gather %base[%c0, %c0][%indices], %mask, %pass_thru
: memref>, vector<8xindex>, vector<8xi1>, vector<8xf32> into vector<8xf32>
return %g : vector<8xf32>
}

//===----------------------------------------------------------------------===//
// @entry
//
// Main entry point
//===----------------------------------------------------------------------===//
func.func @entry() {
// 1. Set up memory (10 x 5)
%c0 = arith.constant 0: index
%c1 = arith.constant 1: index
%c10 = arith.constant 10: index
%c5 = arith.constant 5: index
%A = memref.alloc(%c10, %c5) : memref
scf.for %i = %c0 to %c10 step %c1 {
scf.for %j = %c0 to %c5 step %c1 {
%off = arith.muli %i, %c10 : index
%val_index = arith.addi %j, %off : index
%val_i32 = arith.index_cast %val_index : index to i32
%val = arith.sitofp %val_i32 : i32 to f32
memref.store %val, %A[%i, %j] : memref
}
}
vector.print str "\nSOURCE MEMREF _BEFORE_ CAST: \n"
%A_cast = memref.cast %A : memref to memref<*xf32>
call @printMemrefF32(%A_cast) : (memref<*xf32>) -> ()

// 2. Create a 3 x 3 view
%c3 = arith.constant 3 : index
%A_subview = memref.reinterpret_cast %A to
offset : [0], sizes : [%c3, %c3], strides : [%c5, 1]
: memref to memref>

vector.print str "\nSOURCE MEMREF _AFTER_ CAST: \n"
%A_subview_cast = memref.cast %A_subview : memref> to memref<*xf32>
call @printMemrefF32(%A_subview_cast) : (memref<*xf32>) -> ()

// 3. Set up idx vector ([0, 1, 2, 3, 4, 5, 6, 7])
%indices = vector.step : vector<8xindex>

// 4. Set up pass thru vector.
%u = arith.constant -7.0: f32
%pass = vector.broadcast %u : f32 to vector<8xf32>

// 5. Set up masks.
%all = vector.constant_mask [8] : vector<8xi1>

// 6. "Gather"
// CHECK: LOAD VIA vector.gather
%g5 = call @gather8_strided(%A_subview, %indices, %all, %pass)
: (memref>, vector<8xindex>, vector<8xi1>, vector<8xf32>)
-> (vector<8xf32>)
vector.print str "\nLOAD VIA vector.gather (depnds on lowering!): "
vector.print %g5 : vector<8xf32>

// 7. "Load"
// CHECK: DIRECT LOAD VIA vector.load:
%l1 = vector.load %A_subview[%c0, %c0] : memref>, vector<8xf32>
vector.print str "\nDIRECT LOAD VIA vector.load: "
vector.print %l1 : vector<8xf32>

memref.dealloc %A : memref
return
}
func.func private @printMemrefF32(%ptr : memref<*xf32>)
```

To run:
```
$ cp file.mlir /mlir/test/Integration/Dialect/Vector/CPU/gather.mlir
$ cd && bin/llvm-lit -va /mlir/test/Integration/Dialect/Vector/CPU/gather.mlir
```

**OUTPUT**

For the strided view:
```bash
[[0, 1, 2],
[10, 11, 12],
[20, 21, 22]]
```
the lowering via `--test-vector-gather-lowering` produces (all elements within the strided view):
```bash
(0, 1, 2, 10, 11, 12, 20, 21)
```
while direct lowering produces (some elements _outside_ the strided view):
```baah
(0, 1, 2, 3, 4, 10, 11, 12)
```
For reference, vector.load from the same base also produces:
```bash
(0, 1, 2, 3, 4, 10, 11, 12)
```
I currently find the “logical memref view” interpretation more natural (see also https://github.com/llvm/llvm-project/pull/181357) for `vector.gather`, but the immediate issue is that the semantics need to be made explicit so both lowering paths agree.

**REFERENCES**
* https://discourse.llvm.org/t/rfc-semantics-of-vector-gather-indices-with-strided-memrefs/
* https://github.com/llvm/llvm-project/pull/181357
* https://github.com/llvm/llvm-project/pull/184706

**EDITS**
If you notice any errors, please let me know and I will update this description accordingly.

CC @Groverkss @dcaballe

Contributor guide

Open the contributing guide

Research direction

Start with the reproducer in mlir/test/Integration/Dialect/Vector/CPU/gather.mlir and run it with llvm-lit to compare the two lowering pipelines on the strided memref. Read the linked RFC and referenced discussions to establish the intended vector.gather semantics, then verify that direct and test-vector-gather-lowering paths agree, including the noted vector.scatter impact.

Written by the indexing model from the issue text.

Assessment

Domain
compilers
Issue type
Bug
Difficulty
5/5
Estimated time
Over a week
Activity status
Quiet
Clarity
Needs clarification
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.