[mlir:ODS] Generate unwrapped operation attribute setters

[rriddle]

Summary

This allows for setting an attribute using the underlying C++ type,
which is generally much nicer to interact with than the attribute type.

[lattner]

This will remove so much boilerplate, thank you!

[jpienaar]

Nice, thanks! Yeah UnitAttr is unfortunately special.

[Mogball]

mlir/tools/mlir-tblgen/OpDefinitionsGen.cpp

1120

Is this the best way we have for detecting UnitAttr?

[ayzhuang]

mlir/tools/mlir-tblgen/OpDefinitionsGen.cpp

1155

*attrValue => (*attrValue)

mlir/include/mlir/Dialect/Affine/IR/AffineOps.td

let extraClassDeclaration = [{
 
     /// Returns elements of the loop lower bound.
     AffineMap getLowerBoundMap(unsigned pos);
     operand_range getLowerBoundsOperands();
     AffineValueMap getLowerBoundsValueMap();
 
     /// Sets elements of the loop lower bound.
     void setLowerBounds(ValueRange operands, AffineMap map);
-    void setLowerBoundsMap(AffineMap map);
 
     /// Returns elements of the loop upper bound.
     AffineMap getUpperBoundMap(unsigned pos);
     operand_range getUpperBoundsOperands();
     AffineValueMap getUpperBoundsValueMap();
 
     /// Sets elements fo the loop upper bound.
     void setUpperBounds(ValueRange operands, AffineMap map);
-    void setUpperBoundsMap(AffineMap map);
 
     void setSteps(ArrayRef<int64_t> newSteps);
 
     /// Returns attribute names to use in op construction. Not expected to be
     /// used directly.
     static StringRef getReductionsAttrStrName() { return "reductions"; }
     static StringRef getLowerBoundsMapAttrStrName() { return "lowerBoundsMap"; }
     static StringRef getLowerBoundsGroupsAttrStrName() {

mlir/lib/Dialect/Affine/IR/AffineOps.cpp

void AffineParallelOp::setUpperBounds(ValueRange ubOperands, AffineMap map) {
 
   SmallVector<Value, 4> newOperands(getLowerBoundsOperands());
   newOperands.append(ubOperands.begin(), ubOperands.end());
   (*this)->setOperands(newOperands);
 
   setUpperBoundsMapAttr(AffineMapAttr::get(map));
 }
 
-void AffineParallelOp::setLowerBoundsMap(AffineMap map) {
-  AffineMap lbMap = getLowerBoundsMap();
-  assert(lbMap.getNumDims() == map.getNumDims() &&
-         lbMap.getNumSymbols() == map.getNumSymbols());
-  (void)lbMap;
-  setLowerBoundsMapAttr(AffineMapAttr::get(map));
-}
-
-void AffineParallelOp::setUpperBoundsMap(AffineMap map) {
-  AffineMap ubMap = getUpperBoundsMap();
-  assert(ubMap.getNumDims() == map.getNumDims() &&
-         ubMap.getNumSymbols() == map.getNumSymbols());
-  (void)ubMap;
-  setUpperBoundsMapAttr(AffineMapAttr::get(map));
-}
-
 void AffineParallelOp::setSteps(ArrayRef<int64_t> newSteps) {
   setStepsAttr(getBodyBuilder().getI64ArrayAttr(newSteps));
 }
 
 LogicalResult AffineParallelOp::verify() {
   auto numDims = getNumDims();
   if (getLowerBoundsGroups().getNumElements() != numDims ||
       getUpperBoundsGroups().getNumElements() != numDims ||

mlir/lib/Dialect/Arith/IR/ArithOps.cpp

const std::pair<Pred, Pred> invPreds[] = {
         {Pred::slt, Pred::sgt}, {Pred::sgt, Pred::slt}, {Pred::sle, Pred::sge},
         {Pred::sge, Pred::sle}, {Pred::ult, Pred::ugt}, {Pred::ugt, Pred::ult},
         {Pred::ule, Pred::uge}, {Pred::uge, Pred::ule}, {Pred::eq, Pred::eq},
         {Pred::ne, Pred::ne},
     };
     Pred origPred = getPredicate();
     for (auto pred : invPreds) {
       if (origPred == pred.first) {
-        setPredicateAttr(CmpIPredicateAttr::get(getContext(), pred.second));
+        setPredicate(pred.second);
         Value lhs = getLhs();
         Value rhs = getRhs();
         getLhsMutable().assign(rhs);
         getRhsMutable().assign(lhs);
         return getResult();
       }
     }
     llvm_unreachable("unknown cmpi predicate kind");

mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp

OpFoldResult LLVM::GEPOp::fold(ArrayRef<Attribute> operands) {
   }
   if (changed) {
     SmallVector<int32_t> rawConstantIndices;
     SmallVector<Value> dynamicIndices;
     destructureIndices(getSourceElementType(), gepArgs, rawConstantIndices,
                        dynamicIndices);
 
     getDynamicIndicesMutable().assign(dynamicIndices);
-    setRawConstantIndicesAttr(
-        DenseI32ArrayAttr::get(getContext(), rawConstantIndices));
+    setRawConstantIndices(rawConstantIndices);
     return Value{*this};
   }
 
   return {};
 }
 
 //===----------------------------------------------------------------------===//
 // LLVMDialect initialization, type parsing, and registration.

mlir/lib/Target/LLVMIR/ConvertFromLLVMIR.cpp

if (gv->hasInitializer() && !valueAttr) {
     currentEntryBlock = b.createBlock(&r);
     b.setInsertionPoint(currentEntryBlock, currentEntryBlock->begin());
     Value v = processConstant(gv->getInitializer());
     if (!v)
       return nullptr;
     b.create<ReturnOp>(op.getLoc(), ArrayRef<Value>({v}));
   }
   if (gv->hasAtLeastLocalUnnamedAddr())
-    op.setUnnamedAddrAttr(UnnamedAddrAttr::get(
-        context, convertUnnamedAddrFromLLVM(gv->getUnnamedAddr())));
+    op.setUnnamedAddr(convertUnnamedAddrFromLLVM(gv->getUnnamedAddr()));
   if (gv->hasSection())
-    op.setSectionAttr(b.getStringAttr(gv->getSection()));
+    op.setSection(gv->getSection());
 
   return globals[gv] = op;
 }
 
 Value Importer::processConstant(llvm::Constant *c) {
   OpBuilder bEntry(currentEntryBlock, currentEntryBlock->begin());
   if (Attribute attr = getConstantAsAttr(c)) {
     // These constants can be represented as attributes.
if (auto *type = f->getParamInAllocaType(arg.index())) {
           NamedAttribute(b.getStringAttr(LLVMDialect::getInAllocaAttrName()),
                          TypeAttr::get(mlirType)));
     }
 
     fop.setArgAttrs(arg.index(), argAttrs);
   }
 
   if (FlatSymbolRefAttr personality = getPersonalityAsAttr(f))
-    fop->setAttr(b.getStringAttr("personality"), personality);
+    fop.setPersonalityAttr(personality);
   else if (f->hasPersonalityFn())
     emitWarning(UnknownLoc::get(context),
                 "could not deduce personality, skipping it");
 
   if (f->hasGC())
-    fop.setGarbageCollectorAttr(b.getStringAttr(f->getGC()));
+    fop.setGarbageCollector(StringRef(f->getGC()));
 
   // Handle Function attributes.
   processFunctionAttributes(f, fop);
 
   if (f->isDeclaration())
     return success();
 
   // Eagerly create all blocks.

mlir/test/mlir-tblgen/op-attribute.td

 // DEF-NEXT:       return some-const-builder-call(::mlir::Builder((*this)->getContext()), 4.2).some-convert-from-storage();
 // DEF-NEXT:   return attr.some-convert-from-storage();
 
 // Test setter methods
 // ---
 
 // DEF:      void AOp::setAAttrAttr(some-attr-kind attr) {
 // DEF-NEXT:   (*this)->setAttr(getAAttrAttrName(), attr);
+// DEF:      void AOp::setAAttr(some-return-type attrValue) {
+// DEF-NEXT:   (*this)->setAttr(getAAttrAttrName(), some-const-builder-call(::mlir::Builder(getContext()), attrValue));
 // DEF:      void AOp::setBAttrAttr(some-attr-kind attr) {
 // DEF-NEXT:   (*this)->setAttr(getBAttrAttrName(), attr);
+// DEF:      void AOp::setBAttr(some-return-type attrValue) {
+// DEF-NEXT:   (*this)->setAttr(getBAttrAttrName(), some-const-builder-call(::mlir::Builder(getContext()), attrValue));
 // DEF:      void AOp::setCAttrAttr(some-attr-kind attr) {
 // DEF-NEXT:   (*this)->setAttr(getCAttrAttrName(), attr);
+// DEF:      void AOp::setCAttr(::llvm::Optional<some-return-type> attrValue) {
+// DEF-NEXT:   if (attrValue)
+// DEF-NEXT:     return (*this)->setAttr(getCAttrAttrName(), some-const-builder-call(::mlir::Builder(getContext()), *attrValue));
+// DEF-NEXT:   (*this)->removeAttr(getCAttrAttrName());
 
 // Test remove methods
 // ---
 
 // DEF: ::mlir::Attribute AOp::removeCAttrAttr() {
 // DEF-NEXT: return (*this)->removeAttr(getCAttrAttrName());
 
 // Test build methods

mlir/tools/mlir-tblgen/OpDefinitionsGen.cpp

 // Returns true if we can use unwrapped value for the given `attr` in builders.
 static bool canUseUnwrappedRawValue(const tblgen::Attribute &attr) {
   return attr.getReturnType() != attr.getStorageType() &&
          // We need to wrap the raw value into an attribute in the builder impl
          // so we need to make sure that the attribute specifies how to do that.
          !attr.getConstBuilderTemplate().empty();
 }
 
+/// Build an attribute from a parameter value using the constant builder.
+static std::string constBuildAttrFromParam(const tblgen::Attribute &attr,
+                                           FmtContext &fctx,
+                                           StringRef paramName) {
+  std::string builderTemplate = attr.getConstBuilderTemplate().str();
+
+  // For StringAttr, its constant builder call will wrap the input in
+  // quotes, which is correct for normal string literals, but incorrect
+  // here given we use function arguments. So we need to strip the
+  // wrapping quotes.
+  if (StringRef(builderTemplate).contains("\"$0\""))
+    builderTemplate = replaceAllSubstrs(builderTemplate, "\"$0\"", "$0");
+
+  return tgfmt(builderTemplate, &fctx, paramName).str();
+}
+
 namespace {
 /// Metadata on a registered attribute. Given that attributes are stored in
 /// sorted order on operations, we can use information from ODS to deduce the
 /// number of required attributes less and and greater than each attribute,
 /// allowing us to search only a subrange of the attributes in ODS-generated
 /// getters.
 struct AttributeMetadata {
   /// The attribute name.
auto emitAttrWithStorageType = [&](StringRef setterName, StringRef getterName,
     auto *method =
         opClass.addMethod("void", setterName + "Attr",
                           MethodParameter(attr.getStorageType(), "attr"));
     if (method)
       method->body() << formatv("  (*this)->setAttr({0}AttrName(), attr);",
                                 getterName);
   };
 
+  // Generate a setter that accepts the underlying C++ type as opposed to the
+  // attribute type.
+  auto emitAttrWithReturnType = [&](StringRef setterName, StringRef getterName,
+                                    Attribute attr) {
+    Attribute baseAttr = attr.getBaseAttr();
+    if (!canUseUnwrappedRawValue(baseAttr))
+      return;
+    FmtContext fctx;
+    fctx.withBuilder("::mlir::Builder(getContext())");
+    bool isUnitAttr = attr.getAttrDefName() == "UnitAttr";
+    bool isOptional = attr.isOptional();
+
+    auto createMethod = [&](const Twine &paramType) {
+      return opClass.addMethod("void", setterName,
+                               MethodParameter(paramType.str(), "attrValue"));
+    };
+
+    // Build the method using the correct parameter type depending on
+    // optionality.
+    Method *method = nullptr;
+    if (isUnitAttr)
+      method = createMethod("bool");
+    else if (isOptional)
+      method =
+          createMethod("::llvm::Optional<" + baseAttr.getReturnType() + ">");
+    else
+      method = createMethod(attr.getReturnType());
+    if (!method)
+      return;
+
+    // If the value isn't optional, just set it directly.
+    if (!isOptional) {
+      method->body() << formatv(
+          "  (*this)->setAttr({0}AttrName(), {1});", getterName,
+          constBuildAttrFromParam(attr, fctx, "attrValue"));
+      return;
+    }
+
+    // Otherwise, we only set if the provided value is valid. If it isn't, we
+    // remove the attribute.
+
+    // TODO: Handle unit attr parameters specially, given that it is treated as
+    // optional but not in the same way as the others (i.e. it uses bool over
+    // Optional<>).
+    StringRef paramStr = isUnitAttr ? "attrValue" : "*attrValue";
+    const char *optionalCodeBody = R"(
+    if (attrValue)
+      return (*this)->setAttr({0}AttrName(), {1});
+    (*this)->removeAttr({0}AttrName());)";
+    method->body() << formatv(
+        optionalCodeBody, getterName,
+        constBuildAttrFromParam(baseAttr, fctx, paramStr));
+  };
+
   for (const NamedAttribute &namedAttr : op.getAttributes()) {
     if (namedAttr.attr.isDerivedAttr())
       continue;
-    for (auto names : llvm::zip(op.getSetterNames(namedAttr.name),
-                                op.getGetterNames(namedAttr.name)))
-      emitAttrWithStorageType(std::get<0>(names), std::get<1>(names),
-                              namedAttr.attr);
+    for (auto [setterName, getterName] :
+         llvm::zip(op.getSetterNames(namedAttr.name),
+                   op.getGetterNames(namedAttr.name))) {
+      emitAttrWithStorageType(setterName, getterName, namedAttr.attr);
+      emitAttrWithReturnType(setterName, getterName, namedAttr.attr);
+    }
   }
 }
 
 void OpEmitter::genOptionalAttrRemovers() {
   // Generate methods for removing optional attributes, instead of having to
   // use the string interface. Enables better compile time verification.
   auto emitRemoveAttr = [&](StringRef name) {
     auto upperInitial = name.take_front().upper();
for (const auto &namedAttr : op.getAttributes()) {
     if (emitNotNullCheck)
       body.indent() << formatv("if ({0}) ", namedAttr.name) << "{\n";
 
     if (isRawValueAttr && canUseUnwrappedRawValue(attr)) {
       // If this is a raw value, then we need to wrap it in an Attribute
       // instance.
       FmtContext fctx;
       fctx.withBuilder("odsBuilder");
-
-      std::string builderTemplate = std::string(attr.getConstBuilderTemplate());
-
-      // For StringAttr, its constant builder call will wrap the input in
-      // quotes, which is correct for normal string literals, but incorrect
-      // here given we use function arguments. So we need to strip the
-      // wrapping quotes.
-      if (StringRef(builderTemplate).contains("\"$0\""))
-        builderTemplate = replaceAllSubstrs(builderTemplate, "\"$0\"", "$0");
-
-      std::string value =
-          std::string(tgfmt(builderTemplate, &fctx, namedAttr.name));
       body << formatv("  {0}.addAttribute({1}AttrName({0}.name), {2});\n",
-                      builderOpState, op.getGetterName(namedAttr.name), value);
+                      builderOpState, op.getGetterName(namedAttr.name),
+                      constBuildAttrFromParam(attr, fctx, namedAttr.name));
     } else {
       body << formatv("  {0}.addAttribute({1}AttrName({0}.name), {2});\n",
                       builderOpState, op.getGetterName(namedAttr.name),
                       namedAttr.name);
     }
     if (emitNotNullCheck)
       body.unindent() << "  }\n";
   }