[Mlir-commits] [mlir] [mlir]Add resultSegmentSizes/operandSegmentSizes to prop-dict. (PR #211222)

[Natanael Cintean]

================
@@ -1313,6 +1313,38 @@ ::mlir::MLIRContext *ctx = dict.getContext();
 (void)ctx;
 )decl";
 
+  // `operandSegmentSizes`/`resultSegmentSizes` are trait-injected properties
+  // not enumerated by `op.getProperties()`, so they need to be special-cased
+  // here, mirroring the handling in `setPropertiesFromAttr` in
+  // OpDefinitionsGen.cpp.
+  //
+  // {0}: segment sizes property name
+  // {1}: isRequired
+  const char *segmentSizesFromAttrFmt = R"decl(
+auto {0}AttrName = ::mlir::StringAttr::get(ctx, "{0}");
+usedKeys.insert({0}AttrName);
+auto attr = dict.get({0}AttrName);
+if (!attr && {1}) {{
+  emitError() << "expected key entry for {0} in DictionaryAttr to set "
+             "Properties.";
+  return ::mlir::failure();
+}
+if (attr && ::mlir::failed(::mlir::convertFromAttribute(prop.{0}, attr, [&]() {{
+      return emitError() << "for `{0}`: ";
+    })))
+  return ::mlir::failure();
+)decl";
+  if (op.getTrait("::mlir::OpTrait::AttrSizedOperandSegments")) {
+    auto scope = body.scope("{\n", "}\n", /*indent=*/true);
+    body << formatv(segmentSizesFromAttrFmt, "operandSegmentSizes",
+                    fmt.allOperands);
+  }
+  if (op.getTrait("::mlir::OpTrait::AttrSizedResultSegments")) {
+    auto scope = body.scope("{\n", "}\n", /*indent=*/true);
+    body << formatv(segmentSizesFromAttrFmt, "resultSegmentSizes",
+                    fmt.allResultTypes);
+  }
----------------
natanael-cintean wrote:

@joker-eph , I went a bit back and forward on this, as limiting to non inferable case means that an op with inferable results AND resultSegmentSizes property will still trigger the Unknown key error. In the initial implementation it would be checked for type, inserted in dict and ignored.
But in the end it is probably best to only touch the problematic case.

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