[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(
+            extractOp, linalgOp,
+            state.getCanonicalVecType(extractOp.getType())) ==
+            VectorMemoryAccessKind::Gather)
+      return failure();
----------------
FedericoBruzzone wrote:

Could you please say in the comment why the check runs up front?

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