[Mlir-commits] [mlir] [MLIR][XeGPU] Extend SGMapAttr and Add ConvertLayoutOp (PR #132425)

[Adam Siemieniuk]

================
@@ -68,73 +68,65 @@ LogicalResult ScatterTensorDescAttr::verify(
 }
 
 //===----------------------------------------------------------------------===//
-// XeGPU_SGMapAttr
+// XeGPU_LayoutAttr
 //===----------------------------------------------------------------------===//
-namespace {
-template <typename T, unsigned N>
-LogicalResult parseIntArrayField(::mlir::AsmParser &parser,
-                                 llvm::SmallVector<T, N> &result,
-                                 llvm::StringRef fieldName) {
-  if (failed(parser.parseKeyword(fieldName))) {
-    parser.emitError(parser.getCurrentLocation(),
-                     "unexpected field name. Expected " + fieldName + ".");
-    return failure();
+LogicalResult
+LayoutAttr::verify(llvm::function_ref<mlir::InFlightDiagnostic()> emitError,
+                   DenseI32ArrayAttr sg_layout, DenseI32ArrayAttr sg_data,
+                   DenseI32ArrayAttr inst_data, DenseI32ArrayAttr lane_layout,
+                   DenseI32ArrayAttr lane_data, DenseI32ArrayAttr order) {
+
+  // A valid layout must include at least one of sg_layout and lane_layout.
+  // sg_layout is essential for Workgroup layout, while lane_layout is
+  // required for Subgroup layout.
+  if (!sg_layout && !lane_layout) {
+    return emitError() << "expected at least one of sg_layout or lane_layout";
   }
 
-  if (failed(parser.parseEqual())) {
-    parser.emitError(parser.getCurrentLocation(), "expected '=' sign.");
-    return failure();
+  if (sg_layout && lane_layout && sg_layout.size() != lane_layout.size()) {
+    return emitError()
+           << "expected sg_layout and lane_layout having the same rank";
   }
 
-  auto elemParser = [&]() -> llvm::ParseResult {
-    uint32_t elem = 0;
-    auto res = parser.parseInteger(elem);
-    result.push_back(elem);
-    return res;
-  };
-
-  return parser.parseCommaSeparatedList(AsmParser::Delimiter::Square,
-                                        elemParser, fieldName);
-}
-} // namespace
-
-mlir::Attribute SGMapAttr::parse(::mlir::AsmParser &parser,
-                                 ::mlir::Type attrType) {
-  if (failed(parser.parseLess()))
-    return {};
-
-  llvm::SmallVector<uint32_t, 2> wi_layout, wi_data;
-  if (failed(parseIntArrayField(parser, wi_layout, "wi_layout")))
-    return {};
-
-  if (failed(parser.parseComma()))
-    return {};
+  // sg_data is optional for Workgroup layout, but its presence requires
+  // sg_layout.
+  if (sg_data) {
+    if (!sg_layout)
+      return emitError() << "expected sg_layout being used with sg_data";
+    if (sg_data.size() != sg_layout.size())
+      return emitError()
+             << "expected sg_data having the same rank as sg_layout";
+  }
 
-  if (failed(parseIntArrayField(parser, wi_data, "wi_data")))
-    return {};
+  // inst_data is optional for Subgroup layout, but its presence requires
+  // lane_layout.
+  if (inst_data) {
+    if (!lane_layout)
+      return emitError() << "expected lane_layout being used with inst_data";
+    if (inst_data.size() != lane_layout.size())
+      return emitError()
+             << "expected inst_data having the same rank as lane_layout";
+  }
 
-  return SGMapAttr::getChecked(
-      [&]() { return parser.emitError(parser.getNameLoc()); },
-      parser.getContext(), wi_layout, wi_data);
-}
+  // lane_data is optional for Subgroup layout, but its presence requires
+  // lane_layout.
+  if (lane_data) {
+    if (!lane_layout)
+      return emitError() << "expected lane_layout being used with lane_data";
+    if (lane_data.size() != lane_layout.size())
+      return emitError()
+             << "expected lane_data having the same rank as lane_layout";
+  }
 
-void SGMapAttr::print(::mlir::AsmPrinter &printer) const {
-  printer << "<";
-  printer.printKeywordOrString("wi_layout");
-  printer << " = [" << getWiLayout() << "], ";
-  printer.printKeywordOrString("wi_data");
-  printer << " = [" << getWiData() << "]";
-  printer << ">";
-}
+  if (order) {
+    if (!sg_layout && !lane_layout)
+      return emitError()
+             << "expected sg_layout/lane_layout being used with order";
+    if (order.size() != sg_layout.size() && order.size() != lane_layout.size())
----------------
adam-smnk wrote:

Won't this check crash if one of the layouts is missing?

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