Value traits support (GPUOpen-Drivers#121)

Co-authored-by: Valentin Churavy <[email protected]>

Author: sstipano

example/ExampleDialect.td

@@ -152,6 +152,10 @@ def SizeOfOp : ExampleOp<"sizeof", [Memory<[]>, NoUnwind, WillReturn]> {
   let results = (outs I64:$result);
   let arguments = (ins type:$sizeof_type);
 
+  let value_traits = [
+    (NoCapture $sizeof_type),
+  ];
+
   let summary = "size of a given type";
   let description = [{
     Returns the store size of the given type in bytes.
@@ -235,6 +239,10 @@ def IExtOp : ExampleOp<"iext", [Memory<[]>, NoUnwind, WillReturn]> {
 def StreamReduceOp : OpClass<ExampleDialect> {
   let arguments = (ins Ptr:$ptr, I64:$count, value:$initial);
 
+  let value_traits = [
+    (NoCapture $ptr)
+  ];
+
   let summary = "family of operations that reduce some array in memory";
   let description = [{
     Illustrate the use of the OpClass feature.
@@ -336,3 +344,20 @@ def NoDescriptionOp : Op<ExampleDialect, "no.description.op", [WillReturn]> {
 
   let summary = "Some summary";
 }
+
+def BufferCompareOp : Op<ExampleDialect, "buffer.compare.op", [WillReturn, NoUnwind]> {
+  let results = (outs I32:$ret);
+  let arguments = (ins Ptr:$lhs, Ptr:$rhs);
+
+  let value_traits = [
+    (NoCapture $lhs),
+    (NoCapture $rhs),
+    (NoUndef $lhs),
+    (NoUndef $ret),
+  ];
+
+  let summary = "demonstrate how multiple parameter attributes are added";
+  let description = [{
+    Both arguments get a parameter attribute, as well as return value
+  }];
+}

include/llvm-dialects/Dialect/Dialect.td

@@ -322,29 +322,42 @@ multiclass AttrEnum<string cppType_> {
 /// Traits generally map to llvm::Attributes.
 // ============================================================================
 
-class Trait;
+class TraitProperty;
+def FnTrait : TraitProperty;
+def ParamTrait : TraitProperty;
+def RetTrait : TraitProperty;
 
-class LlvmEnumAttributeTrait<string llvmEnum_> : Trait {
+class Trait<list<TraitProperty> P> {
+  list<TraitProperty> Properties = P;
+}
+
+class LlvmEnumAttributeTrait<string llvmEnum_, list<TraitProperty> P> : Trait<P> {
   string llvmEnum = llvmEnum_;
 }
 
-def NoUnwind : LlvmEnumAttributeTrait<"NoUnwind">;
-def WillReturn : LlvmEnumAttributeTrait<"WillReturn">;
-def NoReturn : LlvmEnumAttributeTrait<"NoReturn">;
-def NoRecurse : LlvmEnumAttributeTrait<"NoRecurse">;
-def NoSync : LlvmEnumAttributeTrait<"NoSync">;
-def NoFree : LlvmEnumAttributeTrait<"NoFree">;
-def MustProgress : LlvmEnumAttributeTrait<"MustProgress">;
-def NoCallback : LlvmEnumAttributeTrait<"NoCallback">;
-def NoDuplicate : LlvmEnumAttributeTrait<"NoDuplicate">;
-def NoBuiltin : LlvmEnumAttributeTrait<"NoBuiltin">;
-def Builtin : LlvmEnumAttributeTrait<"Builtin">;
-def InlineHint : LlvmEnumAttributeTrait<"InlineHint">;
-def AlwaysInline : LlvmEnumAttributeTrait<"AlwaysInline">;
-def Cold : LlvmEnumAttributeTrait<"Cold">;
-def Hot : LlvmEnumAttributeTrait<"Hot">;
-def Convergent : LlvmEnumAttributeTrait<"Convergent">;
-def Speculatable : LlvmEnumAttributeTrait<"Speculatable">;
+def NoUnwind : LlvmEnumAttributeTrait<"NoUnwind", [FnTrait]>;
+def WillReturn : LlvmEnumAttributeTrait<"WillReturn", [FnTrait]>;
+def NoReturn : LlvmEnumAttributeTrait<"NoReturn", [FnTrait]>;
+def NoRecurse : LlvmEnumAttributeTrait<"NoRecurse", [FnTrait]>;
+def NoSync : LlvmEnumAttributeTrait<"NoSync", [FnTrait]>;
+def NoFree : LlvmEnumAttributeTrait<"NoFree", [FnTrait]>;
+def MustProgress : LlvmEnumAttributeTrait<"MustProgress", [FnTrait]>;
+def NoCallback : LlvmEnumAttributeTrait<"NoCallback", [FnTrait]>;
+def NoDuplicate : LlvmEnumAttributeTrait<"NoDuplicate", [FnTrait]>;
+def NoBuiltin : LlvmEnumAttributeTrait<"NoBuiltin", [FnTrait]>;
+def Builtin : LlvmEnumAttributeTrait<"Builtin", [FnTrait]>;
+def InlineHint : LlvmEnumAttributeTrait<"InlineHint", [FnTrait]>;
+def AlwaysInline : LlvmEnumAttributeTrait<"AlwaysInline", [FnTrait]>;
+def Cold : LlvmEnumAttributeTrait<"Cold", [FnTrait]>;
+def Hot : LlvmEnumAttributeTrait<"Hot", [FnTrait]>;
+def Convergent : LlvmEnumAttributeTrait<"Convergent", [FnTrait]>;
+def Speculatable : LlvmEnumAttributeTrait<"Speculatable", [FnTrait]>;
+
+def NoCapture : LlvmEnumAttributeTrait<"NoCapture", [ParamTrait]>;
+def ReadOnly : LlvmEnumAttributeTrait<"ReadOnly", [ParamTrait]>;
+
+def NoUndef : LlvmEnumAttributeTrait<"NoUndef", [ParamTrait, RetTrait]>;
+def NonNull : LlvmEnumAttributeTrait<"NonNull", [ParamTrait, RetTrait]>;
 
 /// Represent the LLVM `memory(...)` attribute as the OR (or union) of memory
 /// effects. An empty effects list means the operation does not access memory
@@ -358,7 +371,7 @@ def Speculatable : LlvmEnumAttributeTrait<"Speculatable">;
 /// Example: `Memory<[(ref), (mod ArgMem, InaccessibleMem)]>` means the
 /// operation may read from any kind of memory and write to argument and
 /// inaccessible memory.
-class Memory<list<dag> effects_> : Trait {
+class Memory<list<dag> effects_> : Trait<[FnTrait]> {
   list<dag> effects = effects_;
 }
 
@@ -394,6 +407,8 @@ class OpClass<Dialect dialect_> : OpClassBase {
 
   dag arguments = ?;
 
+  list<dag> value_traits = [];
+
   string summary = ?;
   string description = ?;
 }
@@ -412,6 +427,8 @@ class Op<Dialect dialect_, string mnemonic_, list<Trait> traits_> {
   dag arguments = ?;
   dag results = ?;
 
+  list<dag> value_traits = [];
+
   list<dag> verifier = [];
 
   string summary = ?;

include/llvm-dialects/TableGen/Dialects.h

@@ -19,6 +19,7 @@
 #pragma once
 
 #include <memory>
+#include <unordered_map>
 
 #include "llvm-dialects/TableGen/Common.h"
 #include "llvm/ADT/ArrayRef.h"
@@ -75,7 +76,7 @@ class GenDialectsContext {
   void init(RecordKeeperTy &records,
             const llvm::DenseSet<llvm::StringRef> &dialects);
 
-  Trait *getTrait(RecordTy *traitRec);
+  Trait *getTrait(RecordTy *traitRec, int idx = -1);
   Predicate *getPredicate(const llvm::Init *init, llvm::raw_ostream &errs);
   Attr *getAttr(RecordTy *record, llvm::raw_ostream &errs);
   OpClass *getOpClass(RecordTy *opClassRec);
@@ -97,6 +98,7 @@ class GenDialectsContext {
   const llvm::Init *m_any = nullptr;
   bool m_attrsComplete = false;
   llvm::DenseMap<RecordTy *, std::unique_ptr<Trait>> m_traits;
+  llvm::DenseMap<RecordTy *, std::unordered_map<int, std::unique_ptr<Trait>>> m_valueTraits;
   llvm::DenseMap<const llvm::Init *, std::unique_ptr<Predicate>> m_predicates;
   llvm::DenseMap<RecordTy *, std::unique_ptr<Attr>> m_attrs;
   llvm::DenseMap<RecordTy *, std::unique_ptr<OpClass>> m_opClasses;

include/llvm-dialects/TableGen/Operations.h

@@ -67,6 +67,11 @@ class OperationBase {
   void emitArgumentAccessorDefinitions(llvm::raw_ostream &out,
                                        FmtContext &fmt) const;
 
+  void parseValueTraits(llvm::raw_ostream &errs, RecordTy *record,
+                        GenDialectsContext &context);
+
+  std::vector<Trait *> traits;
+
 protected:
   bool init(llvm::raw_ostream &errs, GenDialectsContext &context,
             RecordTy *record);
@@ -105,7 +110,6 @@ class Operation : public OperationBase {
   std::string mnemonic;
   std::string summary;
   std::string description;
-  std::vector<Trait *> traits;
 
   std::vector<NamedValue> results;
 

include/llvm-dialects/TableGen/Traits.h

@@ -40,16 +40,19 @@ class Trait {
   enum class Kind : uint8_t {
     LlvmAttributeTrait_First,
     LlvmEnumAttributeTrait = LlvmAttributeTrait_First,
+    LlvmEnumFnAttributeTrait,
+    LlvmEnumRetAttributeTrait,
+    LlvmEnumParamAttributeTrait,
     LlvmMemoryAttributeTrait,
     LlvmAttributeTrait_Last = LlvmMemoryAttributeTrait,
   };
 
   static std::unique_ptr<Trait> fromRecord(GenDialectsContext *context,
-                                           RecordTy *record);
+                                           RecordTy *record, int idx = 0);
 
   virtual ~Trait() = default;
 
-  virtual void init(GenDialectsContext *context, RecordTy *record);
+  virtual void init(GenDialectsContext *context, RecordTy *record, int idx);
 
   Kind getKind() const { return m_kind; }
   RecordTy *getRecord() const { return m_record; }

lib/TableGen/Dialects.cpp

@@ -79,14 +79,14 @@ void GenDialect::finalize(raw_ostream &errs) {
 GenDialectsContext::GenDialectsContext() = default;
 GenDialectsContext::~GenDialectsContext() = default;
 
-Trait *GenDialectsContext::getTrait(RecordTy *traitRec) {
+Trait *GenDialectsContext::getTrait(RecordTy *traitRec, int idx) {
   if (!traitRec->isSubClassOf("Trait"))
     report_fatal_error(Twine("Trying to use '") + traitRec->getName() +
                        "' as a trait, but it is not a subclass of 'Trait'");
 
-  auto &result = m_traits[traitRec];
+  auto &result = idx < 0 ? m_traits[traitRec] : m_valueTraits[traitRec][idx];
   if (!result)
-    result = Trait::fromRecord(this, traitRec);
+    result = Trait::fromRecord(this, traitRec, idx);
   return result.get();
 }
 

lib/TableGen/GenDialect.cpp

@@ -353,9 +353,11 @@ void llvm_dialects::genDialectDefs(raw_ostream &out, RecordKeeperTy &records) {
   if (!dialect->attribute_lists_empty()) {
     FmtContextScope scope{fmt};
     fmt.addSubst("attrBuilder", "attrBuilder");
+    fmt.addSubst("argAttrList", "argAttrList");
 
     for (const auto &enumeratedTraits : enumerate(dialect->attribute_lists())) {
       out << tgfmt("{\n  ::llvm::AttrBuilder $attrBuilder{context};\n", &fmt);
+      out << tgfmt(" ::llvm::AttributeList $argAttrList;\n", &fmt);
 
       for (const Trait *trait : enumeratedTraits.value()) {
         if (auto *llvmAttribute = dyn_cast<LlvmAttributeTrait>(trait)) {
@@ -365,8 +367,8 @@ void llvm_dialects::genDialectDefs(raw_ostream &out, RecordKeeperTy &records) {
         }
       }
 
-      out << tgfmt("m_attributeLists[$0] = ::llvm::AttributeList::get(context, "
-                   "::llvm::AttributeList::FunctionIndex, $attrBuilder);\n}\n",
+      out << tgfmt("m_attributeLists[$0] = "
+                   "$argAttrList.addFnAttributes(context, $attrBuilder);\n}\n",
                    &fmt, enumeratedTraits.index());
     }
   }

lib/TableGen/Operations.cpp

@@ -22,6 +22,7 @@
 #include "llvm-dialects/TableGen/Dialects.h"
 #include "llvm-dialects/TableGen/Format.h"
 
+#include "llvm/Support/ErrorHandling.h"
 #include "llvm/TableGen/Record.h"
 
 using namespace llvm;
@@ -90,6 +91,11 @@ bool OperationBase::init(raw_ostream &errs, GenDialectsContext &context,
 
   m_arguments = std::move(*arguments);
 
+  if (m_superclass && m_superclass->traits.size() > 0)
+    traits = m_superclass->traits;
+
+  parseValueTraits(errs, record, context);
+
   // Don't allow any other arguments if the superclass already uses
   // variadic arguments, as the arguments will be appended to the arguments of
   // the superclass.
@@ -297,6 +303,47 @@ void OperationBase::emitArgumentAccessorDefinitions(llvm::raw_ostream &out,
   }
 }
 
+void OperationBase::parseValueTraits(raw_ostream &errs, RecordTy *record,
+                                     GenDialectsContext &context) {
+  const DagInit *insDag = record->getValueAsDag("arguments");
+  std::unordered_map<std::string, unsigned> nameToIndexMap;
+  for (unsigned i = 0; i < insDag->getNumArgs(); ++i) {
+    StringRef name = insDag->getArgNameStr(i);
+    nameToIndexMap[name.str()] = i + 1;
+  }
+
+  const RecordVal *outsVal = record->getValue("results");
+  if (outsVal) {
+    const DagInit *DI = cast<DagInit>(outsVal->getValue());
+    if (DI->getNumArgs() > 0) {
+      StringRef name = DI->getArgNameStr(0);
+      nameToIndexMap[name.str()] = 0;
+    }
+  }
+
+  const ListInit *List = record->getValueAsListInit("value_traits");
+  for (const Init *I : List->getValues()) {
+    if (const DagInit *DI = dyn_cast<DagInit>(I)) {
+      if (DI->getNumArgs() != 1) {
+        errs << "value_traits " << *DI << " is missing argument name";
+        return;
+      }
+
+      StringRef name = DI->getArgNameStr(0);
+
+      if (const DefInit *Op = dyn_cast<DefInit>(DI->getOperator())) {
+        traits.push_back(
+            context.getTrait(Op->getDef(), nameToIndexMap[name.str()]));
+      } else {
+        errs << "value_traits " << *DI << " is not of form (Trait $arg)";
+        return;
+      }
+    } else {
+      report_fatal_error("value_traits was not a list of DAG's");
+    }
+  }
+}
+
 std::unique_ptr<OpClass> OpClass::parse(raw_ostream &errs,
                                         GenDialectsContext &context,
                                         RecordTy *record) {