[Mlir-commits] [mlir] [Linalg][Vectorization] Add support for linalg vectorization of a tensor.extract case (PR #107922)

[Andrzej Warzyński]

================
@@ -810,6 +810,23 @@ static Value calculateGatherOffset(RewriterBase &rewriter,
 
 enum VectorMemoryAccessKind { ScalarBroadcast, Contiguous, Gather };
 
+/// Find the non constant dim in a linalgOp. This is used for finding contiguous
+/// loads and it is expected that only one dim will be non constant, if thats
+/// not the case this function will assert.
+static uint64_t getNonUnitLoopDim(LinalgOp linalgOp) {
+  uint64_t nonUnitDim = 0;
+  uint64_t countNonUnitDim = 0;
+  for (auto tripCount : llvm::enumerate(linalgOp.getStaticLoopRanges())) {
+    if (tripCount.value() != 1) {
+      nonUnitDim = tripCount.index();
+      countNonUnitDim++;
+    }
+  }
+  assert(countNonUnitDim == 1 &&
+         "Expected only one non unit loop dim in this linalg op");
----------------
banach-space wrote:

As you have noticed, this condition does not hold for "dynamically" shaped Linalg Ops. This is fine, because:
* For statically shaped Linalg Ops, we infer vector sizes (for the vectoriser) from static loop ranges. Hence we do require that there's only one non-unit dim in static loop ranges for the corresponding memory access to qualify as a contiguous load.
* For dynamically shaped Linalg Ops, we need to specify the output vector sizes (for the vectoriser) ourselves. These need to be 1-D vectors for the corresponding memory access to be considered contiguous. However, it also means that this requirement on the loop ranges can be relaxed (because it no longer defines the memory access pattern). 

Does this make sense? 

```suggestion
  assert(linalgOp.hasDynamicShape() || countNonUnitDim == 1 &&
         "For statically shaped Linalg Ops, only one non-unit loop dim is expected");
```

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