[Mlir-commits] [mlir] [mlir][VectorToXeGPU] Un-flatten gather/scatter in convert-vector-to-xegpu (PR #218292)

[Charitha Saumya]

================
@@ -974,6 +975,151 @@ struct ContractionLowering : public OpRewritePattern<vector::ContractionOp> {
   }
 };
 
+// Returns the `vector.shape_cast` that flattened a value of type `ndType` into
+// `flat`, if that is how `flat` was produced.
+static vector::ShapeCastOp getFlattenCast(Value flat, VectorType ndType) {
+  auto shapeCast = flat.getDefiningOp<vector::ShapeCastOp>();
+  if (shapeCast && shapeCast.getSourceVectorType() == ndType)
+    return shapeCast;
+  return nullptr;
+}
+
+static DenseElementsAttr getDenseConstant(Value flat) {
+  DenseElementsAttr elements;
+  if (matchPattern(flat, m_Constant(&elements)))
+    return elements;
+  return nullptr;
+}
+
+static vector::BroadcastOp getSplatBroadcast(Value flat) {
+  auto broadcast = flat.getDefiningOp<vector::BroadcastOp>();
+  if (broadcast && !isa<VectorType>(broadcast.getSourceType()))
+    return broadcast;
+  return nullptr;
+}
+
+// Returns true if `unflatten` can reshape `flat` to `ndType`.
+static bool canUnflatten(Value flat, VectorType ndType) {
+  return getFlattenCast(flat, ndType) || getDenseConstant(flat) ||
+         getSplatBroadcast(flat);
+}
+
+// Reshape `flat` to `ndType` without introducing a 1-D to N-D
+// `vector.shape_cast`. Only handles the forms a frontend actually emits when
+// flattening: the flattening cast itself, a constant, and a splat.
+// `canUnflatten` must hold.
+static Value unflatten(PatternRewriter &rewriter, Value flat,
----------------
charithaintc wrote:

it seems shape_cast folder already does this.

worth trying with following changes:
```
static Value unflatten(PatternRewriter &rewriter, Value flat, VectorType ndType) {
  return vector::ShapeCastOp::create(rewriter, flat.getLoc(), ndType, flat);
}
```
and in the driver
```
patterns.add<UnflattenGatherScatter<vector::GatherOp>,
             UnflattenGatherScatter<vector::ScatterOp>>(ctx);
vector::ShapeCastOp::getCanonicalizationPatterns(patterns, ctx);
```

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