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

[Adam Siemieniuk]

================
@@ -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");
+  auto sgMapB = (*this)->getAttrOfType<xegpu::SGMapAttr>("sg_map_b");
+  auto sgMapC = (*this)->getAttrOfType<xegpu::SGMapAttr>("sg_map_c");
+
+  // If sg_maps are not present, then the operation is in VC mode.
+  if (!sgMapA && !sgMapB && !sgMapC) {
+    if (lhsRank != 2 || (rhsRank != 2 && rhsRank != 3) || resultRank != 2)
+      return emitOpError(
+          "expecting lhs and result to be a 2D vector, and rhs to be either "
+          "2D or 3D (packed) vector.");
+    auto bK = rhsRank == 3 ? rhsShape[0] * rhsShape[2] : rhsShape[0];
+    if (bK != lhsShape[1])
+      return emitOpError("K-dimension mismatch.");
+    if (lhsShape[0] != resultShape[0])
+      return emitOpError("M-dimension mismatch.");
+    if (rhsShape[1] != resultShape[1])
+      return emitOpError("N-dimension mismatch.");
+    return success();
+  }
+  // Otherwise, in SIMT mode we expect sg_map attributes for all operands and
+  // result of DPAS operation.
+  if (!sgMapA || !sgMapB || !sgMapC)
+    return emitOpError("sg_map attributes for all operands and outputs are "
+                       "expected in SIMT xegpu::Dpas operation");
+
+  // In SIMT mode, All data fragments must be 2D
----------------
adam-smnk wrote:

Shouldn't 3D technically be also fine? Looking at IMEX XeGPU docs

AFAIK, 2D is also fine for SIMD with VNNI packing being injected automatically later.
Maybe we could drop 3D operands altogether to simplify usage? @chencha3 @Jianhui-Li

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