[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");
+  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
----------------
charithaintc wrote:

> Thanks for the link. I'll try to catch up on the updates 👍
> 
> Ideally, I'd prefer SIMD to always have a valid path down to SIMT as it's a natural stepping stone in lowering from higher level abstraction. And right now, using 3D operands seems like a one-way street. It would be great to unify representation for both paths assuming VNNI packing is handled before or after - even if the details differ.
> 
For SIMD we deal with 2D vectors for block operations. So with VNNI it becomes 3D. In SIMT, we deal with 1D data fragments (which can directly go to LLVM vectors) so with VNNI it becomes 2D. I think this flow is natural. Maybe @Jianhui-Li  can explain this with more clarity. 
> For example, we could either keep the layout plain for both or fold packing factor into innermost dimension like in AMX dialect (https://mlir.llvm.org/docs/Dialects/AMX/#amxtile_mulf-amxtilemulfop) (not saying it's best design but an option).
Yes. It maybe possible to clear up the SIMD path in future and get rid of 3D representation there. But in general, SIMT path seems clean with this 2D representation for VNNI. 



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