[Mlir-commits] [mlir] [mlir][xegpu] Improve XeGPU op verification logic for SIMT flavor and update tests. (PR #127920)

[Charitha Saumya]

================
@@ -610,20 +620,61 @@ void UpdateOffsetOp::build(OpBuilder &builder, OperationState &state,
 LogicalResult DpasOp::verify() {
   int64_t lhsRank = getLhsType().getRank();
   int64_t rhsRank = getRhsType().getRank();
-
-  if (lhsRank != 2 || (rhsRank != 2 && rhsRank != 3))
-    return emitOpError("expecting lhs to be a 2D vector, and rhs to be either "
-                       "2D or 3D (packed) vector.");
-
+  int64_t resultRank = getResultType().getRank();
   auto lhsShape = getLhsType().getShape();
   auto rhsShape = getRhsType().getShape();
-  auto bK = rhsRank == 3 ? rhsShape[0] * rhsShape[2] : rhsShape[0];
-  if (bK != lhsShape[1])
+  auto resultShape = getResultType().getShape();
+
+  auto sgMapA = (*this)->getAttrOfType<xegpu::SGMapAttr>("sg_map_a");
----------------
charithaintc wrote:

good point. I have added an optional attribute. 

Initially I decided against it because non-XeGPU ops will also need these maps (say vector.transpose) and we will still need to obtain those maps during distribution. But after discussion with team we decided to add optional attributes for DPAS (all collaborative ops in general must have it)

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