[mlir] Support OptionalAttr inside a StructAttr

[tberghammer]

Summary

Previously many of the generated methods were crashing when the user
tried to pass a nullptr (translating to a null attribute) into optional
fields.

[tberghammer]

Fix attribute accessor to not crash for null values

[antiagainst]

Nice, thanks!

[rriddle]

mlir/tools/mlir-tblgen/StructsGen.cpp

147

nit: remove trivial braces here and below.

[tberghammer]

Address review comments

[tberghammer]

Re-try addressing review comments

mlir/tools/mlir-tblgen/StructsGen.cpp

)";
   os << llvm::formatv(fieldStart, fields.size());
 
   const char *getFieldInfo = R"(
   assert({0});
   auto {0}_id = mlir::Identifier::get("{0}", context);
   fields.emplace_back({0}_id, {0});
 )";
 
+  const char *getFieldInfoOptional = R"(
+  if ({0}) {
+    auto {0}_id = mlir::Identifier::get("{0}", context);
+    fields.emplace_back({0}_id, {0});
+  }
+)";
+
   for (auto field : fields) {
+    if (field.getType().isOptional()) {
+      os << llvm::formatv(getFieldInfoOptional, field.getName());
+    } else {
       os << llvm::formatv(getFieldInfo, field.getName());
     }
+  }
 
   const char *getEndInfo = R"(
   Attribute dict = mlir::DictionaryAttr::get(fields, context);
   return dict.dyn_cast<{0}>();
 }
 )";
   os << llvm::formatv(getEndInfo, structName);
 }
 
 static void emitClassofDef(llvm::StringRef structName,
                            llvm::ArrayRef<mlir::tblgen::StructFieldAttr> fields,
                            raw_ostream &os) {
   const char *classofInfo = R"(
 bool {0}::classof(mlir::Attribute attr))";
 
   const char *classofInfoHeader = R"(
-   auto derived = attr.dyn_cast<mlir::DictionaryAttr>();
-   if (!derived)
-     return false;
-   if (derived.size() != {0})
-     return false;
+  if (!attr)
+    return false;
+  int empty_optionals = 0;
 )";
 
   os << llvm::formatv(classofInfo, structName) << " {";
-  os << llvm::formatv(classofInfoHeader, fields.size());
+  os << llvm::formatv(classofInfoHeader);
 
   FmtContext fctx;
   const char *classofArgInfo = R"(
   auto {0} = derived.get("{0}");
   if (!{0} || !({1}))
     return false;
 )";
+  const char *classofArgInfoOptional = R"(
+  auto {0} = derived.get("{0}");
+  if (!{0})
+    ++empty_optionals;
+  else if (!({1}))
+    return false;
+)";
   for (auto field : fields) {
     auto name = field.getName();
     auto type = field.getType();
     std::string condition =
         std::string(tgfmt(type.getConditionTemplate(), &fctx.withSelf(name)));
+    if (type.isOptional()) {
+      os << llvm::formatv(classofArgInfoOptional, name, condition);
+    } else {
       os << llvm::formatv(classofArgInfo, name, condition);
     }
+  }
 
   const char *classofEndInfo = R"(
-  return true;
+  return derived.size() + empty_optionals == {0};
 }
 )";
-  os << classofEndInfo;
+  os << llvm::formatv(classofEndInfo, fields.size());
 }
 
 static void
 emitAccessorDef(llvm::StringRef structName,
                 llvm::ArrayRef<mlir::tblgen::StructFieldAttr> fields,
                 raw_ostream &os) {
   const char *fieldInfo = R"(
 {0} {2}::{1}() const {
   auto derived = this->cast<mlir::DictionaryAttr>();
   auto {1} = derived.get("{1}");
   assert({1} && "attribute not found.");
   assert({1}.isa<{0}>() && "incorrect Attribute type found.");
   return {1}.cast<{0}>();
 }
 )";
+  const char *fieldInfoOptional = R"(
+{0} {2}::{1}() const {
+  auto derived = this->cast<mlir::DictionaryAttr>();
+  auto {1} = derived.get("{1}");
+  if (!{1})
+    return nullptr;
+  assert({1}.isa<{0}>() && "incorrect Attribute type found.");
+  return {1}.cast<{0}>();
+}
+)";
   for (auto field : fields) {
     auto name = field.getName();
     auto type = field.getType();
     auto storage = type.getStorageType();
+    if (type.isOptional())
+      os << llvm::formatv(fieldInfoOptional, storage, name, structName);
+    else
       os << llvm::formatv(fieldInfo, storage, name, structName);
   }
 }
 
 static void emitStructDef(const Record &structDef, raw_ostream &os) {
   StructAttr structAttr(&structDef);
   StringRef cppNamespace = structAttr.getCppNamespace();
   StringRef structName = structAttr.getStructClassName();
   mlir::tblgen::FmtContext ctx;

mlir/unittests/TableGen/StructsGenTest.cpp

static test::TestStruct getTestStruct(mlir::MLIRContext *context) {
   auto integerAttr = mlir::IntegerAttr::get(integerType, 127);
 
   auto floatType = mlir::FloatType::getF32(context);
   auto floatAttr = mlir::FloatAttr::get(floatType, 0.25);
 
   auto elementsType = mlir::RankedTensorType::get({2, 3}, integerType);
   auto elementsAttr =
       mlir::DenseIntElementsAttr::get(elementsType, {1, 2, 3, 4, 5, 6});
+  auto optionalAttr = nullptr;
 
-  return test::TestStruct::get(integerAttr, floatAttr, elementsAttr, context);
+  return test::TestStruct::get(integerAttr, floatAttr, elementsAttr,
+                               optionalAttr, context);
 }
 
 // Validates that test::TestStruct::classof correctly identifies a valid
 // test::TestStruct.
 TEST(StructsGenTest, ClassofTrue) {
   mlir::MLIRContext context;
   auto structAttr = getTestStruct(&context);
   ASSERT_TRUE(test::TestStruct::classof(structAttr));
TEST(StructsGenTest, GetElements) {
   auto denseAttr = returnedAttr.dyn_cast<mlir::DenseElementsAttr>();
   ASSERT_TRUE(denseAttr);
 
   for (const auto &valIndexIt : llvm::enumerate(denseAttr.getIntValues())) {
     EXPECT_EQ(valIndexIt.value(), valIndexIt.index() + 1);
   }
 }
 
+TEST(StructsGenTest, EmptyOptional) {
+  mlir::MLIRContext context;
+  auto structAttr = getTestStruct(&context);
+  EXPECT_EQ(structAttr.sample_optional_integer(), nullptr);
+}
+
 } // namespace mlir

mlir/unittests/TableGen/structs.td

 
 def Test_Dialect : Dialect {
   let name = "test";
 }
 
 def Test_Struct : StructAttr<"TestStruct", Test_Dialect, [
                 StructFieldAttr<"sample_integer", I32Attr>,
                 StructFieldAttr<"sample_float", F32Attr>,
-                StructFieldAttr<"sample_elements", I32ElementsAttr>] > {
+                StructFieldAttr<"sample_elements", I32ElementsAttr>,
+                StructFieldAttr<"sample_optional_integer",
+                                OptionalAttr<I32Attr>>] > {
   let description = "Structure for test data";
 }