[IR] Add support for memory attribute

nikic · SvyatoslavScherbina · commit 7fbe3cfa7529 · 2024-06-25T13:53:55.000+02:00
This implements IR and bitcode support for the memory attribute, as specified in https://reviews.llvm.org/D135597. The new attribute is not used for anything yet (and as such, the old memory attributes are unaffected). Differential Revision: https://reviews.llvm.org/D135592 (cherry picked from commit e9754f0)
diff --git a/llvm/include/llvm/AsmParser/LLParser.h b/llvm/include/llvm/AsmParser/LLParser.h
@@ -42,6 +42,7 @@ namespace llvm {
   class Comdat;
   class MDString;
   class MDNode;
+  class MemoryEffects;
   struct SlotMapping;
 
   /// ValID - Represents a reference of a definition of some sort with no type.
@@ -284,6 +285,7 @@ namespace llvm {
     bool parseOptionalDerefAttrBytes(lltok::Kind AttrKind, uint64_t &Bytes);
     bool parseOptionalUWTableKind(UWTableKind &Kind);
     bool parseAllocKind(AllocFnKind &Kind);
+    Optional<MemoryEffects> parseMemoryAttr();
     bool parseScopeAndOrdering(bool IsAtomic, SyncScope::ID &SSID,
                                AtomicOrdering &Ordering);
     bool parseScope(SyncScope::ID &SSID);
diff --git a/llvm/include/llvm/AsmParser/LLToken.h b/llvm/include/llvm/AsmParser/LLToken.h
@@ -183,6 +183,13 @@ enum Kind {
   kw_##DISPLAY_NAME,
 #include "llvm/IR/Attributes.inc"
 
+  // Memory attribute:
+  kw_read,
+  kw_write,
+  kw_readwrite,
+  kw_argmem,
+  kw_inaccessiblemem,
+
   kw_type,
   kw_opaque,
 
diff --git a/llvm/include/llvm/Bitcode/LLVMBitCodes.h b/llvm/include/llvm/Bitcode/LLVMBitCodes.h
@@ -690,6 +690,7 @@ enum AttributeKindCodes {
   ATTR_KIND_PRESPLIT_COROUTINE = 83,
   ATTR_KIND_FNRETTHUNK_EXTERN = 84,
   ATTR_KIND_SKIP_PROFILE = 85,
+  ATTR_KIND_MEMORY = 86,
 };
 
 enum ComdatSelectionKindCodes {
diff --git a/llvm/include/llvm/IR/Attributes.h b/llvm/include/llvm/IR/Attributes.h
@@ -42,6 +42,7 @@ class AttributeSetNode;
 class FoldingSetNodeID;
 class Function;
 class LLVMContext;
+class MemoryEffects;
 class Type;
 
 enum class AllocFnKind : uint64_t {
@@ -243,6 +244,9 @@ class Attribute {
   // Returns the allocator function kind.
   AllocFnKind getAllocKind() const;
 
+  /// Returns memory effects.
+  MemoryEffects getMemoryEffects() const;
+
   /// The Attribute is converted to a string of equivalent mnemonic. This
   /// is, presumably, for writing out the mnemonics for the assembly writer.
   std::string getAsString(bool InAttrGrp = false) const;
@@ -1220,6 +1224,9 @@ class AttrBuilder {
   // This turns the allocator kind into the form used internally in Attribute.
   AttrBuilder &addAllocKindAttr(AllocFnKind Kind);
 
+  /// Add memory effect attribute.
+  AttrBuilder &addMemoryAttr(MemoryEffects ME);
+
   ArrayRef<Attribute> attrs() const { return Attrs; }
 
   bool operator==(const AttrBuilder &B) const;
diff --git a/llvm/include/llvm/IR/Attributes.td b/llvm/include/llvm/IR/Attributes.td
@@ -126,6 +126,9 @@ def InReg : EnumAttr<"inreg", [ParamAttr, RetAttr]>;
 /// Build jump-instruction tables and replace refs.
 def JumpTable : EnumAttr<"jumptable", [FnAttr]>;
 
+/// Memory effects of the function.
+def Memory : IntAttr<"memory", [FnAttr]>;
+
 /// Function must be optimized for size first.
 def MinSize : EnumAttr<"minsize", [FnAttr]>;
 
diff --git a/llvm/include/llvm/IR/ModRef.h b/llvm/include/llvm/IR/ModRef.h
@@ -82,11 +82,6 @@ class MemoryEffects {
     return (uint32_t)Loc * BitsPerLoc;
   }
 
-  static auto locations() {
-    return enum_seq_inclusive(Location::ArgMem, Location::Other,
-                              force_iteration_on_noniterable_enum);
-  }
-
   MemoryEffects(uint32_t Data) : Data(Data) {}
 
   void setModRef(Location Loc, ModRefInfo MR) {
@@ -97,6 +92,12 @@ class MemoryEffects {
   friend raw_ostream &operator<<(raw_ostream &OS, MemoryEffects RMRB);
 
 public:
+  /// Returns iterator over all supported location kinds.
+  static auto locations() {
+    return enum_seq_inclusive(Location::ArgMem, Location::Other,
+                              force_iteration_on_noniterable_enum);
+  }
+
   /// Create MemoryEffects that can access only the given location with the
   /// given ModRefInfo.
   MemoryEffects(Location Loc, ModRefInfo MR) { setModRef(Loc, MR); }
@@ -147,6 +148,18 @@ class MemoryEffects {
     return FRMB;
   }
 
+  /// Create MemoryEffects from an encoded integer value (used by memory
+  /// attribute).
+  static MemoryEffects createFromIntValue(uint32_t Data) {
+    return MemoryEffects(Data);
+  }
+
+  /// Convert MemoryEffects into an encoded integer value (used by memory
+  /// attribute).
+  uint32_t toIntValue() const {
+    return Data;
+  }
+
   /// Get ModRefInfo for the given Location.
   ModRefInfo getModRef(Location Loc) const {
     return ModRefInfo((Data >> getLocationPos(Loc)) & LocMask);
diff --git a/llvm/lib/AsmParser/LLLexer.cpp b/llvm/lib/AsmParser/LLLexer.cpp
@@ -644,6 +644,12 @@ lltok::Kind LLLexer::LexIdentifier() {
   KEYWORD(DISPLAY_NAME);
 #include "llvm/IR/Attributes.inc"
 
+  KEYWORD(read);
+  KEYWORD(write);
+  KEYWORD(readwrite);
+  KEYWORD(argmem);
+  KEYWORD(inaccessiblemem);
+
   KEYWORD(type);
   KEYWORD(opaque);
 
diff --git a/llvm/lib/AsmParser/LLParser.cpp b/llvm/lib/AsmParser/LLParser.cpp
@@ -14,6 +14,7 @@
 #include "llvm/ADT/APSInt.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/None.h"
+#include "llvm/ADT/ScopeExit.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/AsmParser/LLToken.h"
@@ -36,6 +37,7 @@
 #include "llvm/IR/Intrinsics.h"
 #include "llvm/IR/LLVMContext.h"
 #include "llvm/IR/Metadata.h"
+#include "llvm/IR/ModRef.h"
 #include "llvm/IR/Module.h"
 #include "llvm/IR/Operator.h"
 #include "llvm/IR/Value.h"
@@ -1466,6 +1468,13 @@ bool LLParser::parseEnumAttribute(Attribute::AttrKind Attr, AttrBuilder &B,
     B.addAllocKindAttr(Kind);
     return false;
   }
+  case Attribute::Memory: {
+    Optional<MemoryEffects> ME = parseMemoryAttr();
+    if (!ME)
+      return true;
+    B.addMemoryAttr(*ME);
+    return false;
+  }
   default:
     B.addAttribute(Attr);
     Lex.Lex();
@@ -2187,6 +2196,87 @@ bool LLParser::parseAllocKind(AllocFnKind &Kind) {
   return false;
 }
 
+static Optional<MemoryEffects::Location> keywordToLoc(lltok::Kind Tok) {
+  switch (Tok) {
+  case lltok::kw_argmem:
+    return MemoryEffects::ArgMem;
+  case lltok::kw_inaccessiblemem:
+    return MemoryEffects::InaccessibleMem;
+  default:
+    return None;
+  }
+}
+
+static Optional<ModRefInfo> keywordToModRef(lltok::Kind Tok) {
+  switch (Tok) {
+  case lltok::kw_none:
+    return ModRefInfo::NoModRef;
+  case lltok::kw_read:
+    return ModRefInfo::Ref;
+  case lltok::kw_write:
+    return ModRefInfo::Mod;
+  case lltok::kw_readwrite:
+    return ModRefInfo::ModRef;
+  default:
+    return None;
+  }
+}
+
+Optional<MemoryEffects> LLParser::parseMemoryAttr() {
+  MemoryEffects ME = MemoryEffects::none();
+
+  // We use syntax like memory(argmem: read), so the colon should not be
+  // interpreted as a label terminator.
+  Lex.setIgnoreColonInIdentifiers(true);
+  auto _ = make_scope_exit([&] { Lex.setIgnoreColonInIdentifiers(false); });
+
+  Lex.Lex();
+  if (!EatIfPresent(lltok::lparen)) {
+    tokError("expected '('");
+    return None;
+  }
+
+  bool SeenLoc = false;
+  do {
+    Optional<MemoryEffects::Location> Loc = keywordToLoc(Lex.getKind());
+    if (Loc) {
+      Lex.Lex();
+      if (!EatIfPresent(lltok::colon)) {
+        tokError("expected ':' after location");
+        return None;
+      }
+    }
+
+    Optional<ModRefInfo> MR = keywordToModRef(Lex.getKind());
+    if (!MR) {
+      if (!Loc)
+        tokError("expected memory location (argmem, inaccessiblemem) "
+                 "or access kind (none, read, write, readwrite)");
+      else
+        tokError("expected access kind (none, read, write, readwrite)");
+      return None;
+    }
+
+    Lex.Lex();
+    if (Loc) {
+      SeenLoc = true;
+      ME = ME.getWithModRef(*Loc, *MR);
+    } else {
+      if (SeenLoc) {
+        tokError("default access kind must be specified first");
+        return None;
+      }
+      ME = MemoryEffects(*MR);
+    }
+
+    if (EatIfPresent(lltok::rparen))
+      return ME;
+  } while (EatIfPresent(lltok::comma));
+
+  tokError("unterminated memory attribute");
+  return None;
+}
+
 /// parseOptionalCommaAlign
 ///   ::=
 ///   ::= ',' align 4
diff --git a/llvm/lib/Bitcode/Reader/BitcodeReader.cpp b/llvm/lib/Bitcode/Reader/BitcodeReader.cpp
@@ -55,6 +55,7 @@
 #include "llvm/IR/IntrinsicsARM.h"
 #include "llvm/IR/LLVMContext.h"
 #include "llvm/IR/Metadata.h"
+#include "llvm/IR/ModRef.h"
 #include "llvm/IR/Module.h"
 #include "llvm/IR/ModuleSummaryIndex.h"
 #include "llvm/IR/Operator.h"
@@ -1880,6 +1881,8 @@ static Attribute::AttrKind getAttrFromCode(uint64_t Code) {
     return Attribute::InReg;
   case bitc::ATTR_KIND_JUMP_TABLE:
     return Attribute::JumpTable;
+  case bitc::ATTR_KIND_MEMORY:
+    return Attribute::Memory;
   case bitc::ATTR_KIND_MIN_SIZE:
     return Attribute::MinSize;
   case bitc::ATTR_KIND_NAKED:
@@ -2124,6 +2127,8 @@ Error BitcodeReader::parseAttributeGroupBlock() {
             B.addUWTableAttr(UWTableKind(Record[++i]));
           else if (Kind == Attribute::AllocKind)
             B.addAllocKindAttr(static_cast<AllocFnKind>(Record[++i]));
+          else if (Kind == Attribute::Memory)
+            B.addMemoryAttr(MemoryEffects::createFromIntValue(Record[++i]));
         } else if (Record[i] == 3 || Record[i] == 4) { // String attribute
           bool HasValue = (Record[i++] == 4);
           SmallString<64> KindStr;
diff --git a/llvm/lib/Bitcode/Writer/BitcodeWriter.cpp b/llvm/lib/Bitcode/Writer/BitcodeWriter.cpp
@@ -656,6 +656,8 @@ static uint64_t getAttrKindEncoding(Attribute::AttrKind Kind) {
     return bitc::ATTR_KIND_ALLOCATED_POINTER;
   case Attribute::AllocKind:
     return bitc::ATTR_KIND_ALLOC_KIND;
+  case Attribute::Memory:
+    return bitc::ATTR_KIND_MEMORY;
   case Attribute::Naked:
     return bitc::ATTR_KIND_NAKED;
   case Attribute::Nest:
diff --git a/llvm/lib/IR/Attributes.cpp b/llvm/lib/IR/Attributes.cpp
@@ -26,6 +26,7 @@
 #include "llvm/Config/llvm-config.h"
 #include "llvm/IR/Function.h"
 #include "llvm/IR/LLVMContext.h"
+#include "llvm/IR/ModRef.h"
 #include "llvm/IR/Type.h"
 #include "llvm/Support/Compiler.h"
 #include "llvm/Support/ErrorHandling.h"
@@ -383,6 +384,26 @@ AllocFnKind Attribute::getAllocKind() const {
   return AllocFnKind(pImpl->getValueAsInt());
 }
 
+MemoryEffects Attribute::getMemoryEffects() const {
+  assert(hasAttribute(Attribute::Memory) &&
+         "Can only call getMemoryEffects() on memory attribute");
+  return MemoryEffects::createFromIntValue(pImpl->getValueAsInt());
+}
+
+static const char *getModRefStr(ModRefInfo MR) {
+  switch (MR) {
+  case ModRefInfo::NoModRef:
+    return "none";
+  case ModRefInfo::Ref:
+    return "read";
+  case ModRefInfo::Mod:
+    return "write";
+  case ModRefInfo::ModRef:
+    return "readwrite";
+  }
+  llvm_unreachable("Invalid ModRefInfo");
+}
+
 std::string Attribute::getAsString(bool InAttrGrp) const {
   if (!pImpl) return {};
 
@@ -474,6 +495,48 @@ std::string Attribute::getAsString(bool InAttrGrp) const {
         .str();
   }
 
+  if (hasAttribute(Attribute::Memory)) {
+    std::string Result;
+    raw_string_ostream OS(Result);
+    bool First = true;
+    OS << "memory(";
+
+    MemoryEffects ME = getMemoryEffects();
+
+    // Print access kind for "other" as the default access kind. This way it
+    // will apply to any new location kinds that get split out of "other".
+    ModRefInfo OtherMR = ME.getModRef(MemoryEffects::Other);
+    if (OtherMR != ModRefInfo::NoModRef || ME.getModRef() == OtherMR) {
+      First = false;
+      OS << getModRefStr(OtherMR);
+    }
+
+    for (auto Loc : MemoryEffects::locations()) {
+      ModRefInfo MR = ME.getModRef(Loc);
+      if (MR == OtherMR)
+        continue;
+
+      if (!First)
+        OS << ", ";
+      First = false;
+
+      switch (Loc) {
+      case MemoryEffects::ArgMem:
+        OS << "argmem: ";
+        break;
+      case MemoryEffects::InaccessibleMem:
+        OS << "inaccessiblemem: ";
+        break;
+      case MemoryEffects::Other:
+        llvm_unreachable("This is represented as the default access kind");
+      }
+      OS << getModRefStr(MR);
+    }
+    OS << ")";
+    OS.flush();
+    return Result;
+  }
+
   // Convert target-dependent attributes to strings of the form:
   //
   //   "kind"
@@ -1723,6 +1786,10 @@ AttrBuilder &AttrBuilder::addUWTableAttr(UWTableKind Kind) {
   return addRawIntAttr(Attribute::UWTable, uint64_t(Kind));
 }
 
+AttrBuilder &AttrBuilder::addMemoryAttr(MemoryEffects ME) {
+  return addRawIntAttr(Attribute::Memory, ME.toIntValue());
+}
+
 AttrBuilder &AttrBuilder::addAllocKindAttr(AllocFnKind Kind) {
   return addRawIntAttr(Attribute::AllocKind, static_cast<uint64_t>(Kind));
 }
diff --git a/llvm/lib/Transforms/Utils/CodeExtractor.cpp b/llvm/lib/Transforms/Utils/CodeExtractor.cpp
@@ -922,6 +922,7 @@ Function *CodeExtractor::constructFunction(const ValueSet &inputs,
       case Attribute::WriteOnly:
       case Attribute::AllocKind:
       case Attribute::PresplitCoroutine:
+      case Attribute::Memory:
         continue;
       // Those attributes should be safe to propagate to the extracted function.
       case Attribute::AlwaysInline:
diff --git a/llvm/test/Assembler/memory-attribute-errors.ll b/llvm/test/Assembler/memory-attribute-errors.ll
diff --git a/llvm/test/Assembler/memory-attribute.ll b/llvm/test/Assembler/memory-attribute.ll
