[Mlir-commits] [mlir] [mlir][linalg] Collapse ND tensor.extract sources before gathering (PR #237308)

[Federico Bruzzone]

================
@@ -1495,6 +1510,18 @@ vectorizeAsLinalgGeneric(RewriterBase &rewriter, VectorizationState &state,
   if (linalgOp.getNumDpsInits() == 0)
     return failure();
 
+  // A static gather source needs a flattened extent that fits in int64.
+  for (auto extractOp : block->getOps<tensor::ExtractOp>()) {
+    auto sourceType = extractOp.getTensor().getType();
+    if (sourceType.hasStaticShape() &&
+        !ShapedType::tryGetNumElements(sourceType.getShape()) &&
+        getTensorExtractMemoryAccessPattern(
----------------
FedericoBruzzone wrote:

`getTensorExtractMemoryAccessPattern` is already run at line 1201. Can we avoid this somehow?

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