[Mlir-commits] [mlir] [mlir][vector] Migrate drop-lead-unit-dim to shape_cast (PR #196206)

[Erick Ochoa Lopez]

================
@@ -444,13 +499,16 @@ mlir::vector::castAwayContractionLeadingOneDim(vector::ContractionOp contractOp,
     }
     newIndexingMaps.push_back(AffineMap::get(map.getNumDims() - 1, 0, results,
                                              contractOp.getContext()));
-    // Extract if its a valid extraction, otherwise use the operand
-    // without extraction.
-    newOperands.push_back(validExtract
-                              ? vector::ExtractOp::create(rewriter, loc,
-                                                          operands[it.index()],
-                                                          splatZero(dropDim))
-                              : operands[it.index()]);
+    if (needsShapeCast) {
+      auto operandType = cast<VectorType>(operands[it.index()].getType());
+      if (operandType.getRank() < dropDim ||
+          !hasNonScalableUnitLeadingDims(operandType, dropDim))
+        return failure();
----------------
amd-eochoalo wrote:

This is returning failure after a (guarded) create or fold. Are these mutually exclusive?

```cxx
      if (transposeNeeded) {
        map = AffineMap::get(map.getNumDims(), 0, transposeResults,
                             contractOp.getContext());
        if (transposeNonOuterUnitDims) {
          operands[it.index()] = rewriter.createOrFold<vector::TransposeOp>(
              loc, operands[it.index()], perm);
        }
      }
    }

   // ...

    // here:
    return failure();
```

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