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

[Jianhui Li]

================
@@ -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,
+                       VectorType ndType) {
+  assert(cast<VectorType>(flat.getType()).getRank() == 1 &&
+         "expected a 1-D vector");
+
+  if (auto shapeCast = getFlattenCast(flat, ndType))
+    return shapeCast.getSource();
+
+  if (DenseElementsAttr elements = getDenseConstant(flat))
+    return arith::ConstantOp::create(rewriter, flat.getLoc(), ndType,
+                                     elements.reshape(ndType));
+
+  auto broadcast = getSplatBroadcast(flat);
+  assert(broadcast && "expected canUnflatten to hold");
+  return vector::BroadcastOp::create(rewriter, flat.getLoc(), ndType,
+                                     broadcast.getSource());
+}
+
+// Restore the N-D form of a flattened `vector.gather` / `vector.scatter`.
+//
+// XeGPU layouts are expressed in terms of the N-D shape of the accessed data,
+// so a flattened gather/scatter forces layout propagation to reason through the
+// surrounding `vector.shape_cast` ops. That adds complexity and tends to yield
+// layouts that lower to unoptimized code.
+//
+// Before:
+//   %cst = arith.constant dense<0.0> : vector<8192xbf16>
+//   %fi = vector.shape_cast %idx : vector<128x64xindex> to vector<8192xindex>
+//   %fm = vector.shape_cast %mask : vector<128x64xi1> to vector<8192xi1>
----------------
Jianhui-Li wrote:

These are not supported. 

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