[LLVM] Add support for printing and parsing symbolic address spaces

Author: jurahul

llvm/include/llvm/TargetParser/NVPTXAddressSpaces.h

@@ -0,0 +1,33 @@
+//===----------------------------------------------------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// Address space definitions for NVPTX target.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_TARGETPARSER_NVPTX_ADDRESS_SPACES_H
+#define LLVM_TARGETPARSER_NVPTX_ADDRESS_SPACES_H
+
+namespace llvm::NVPTX {
+
+using AddressSpaceUnderlyingType = unsigned int;
+enum AddressSpace : AddressSpaceUnderlyingType {
+  Generic = 0,
+  Global = 1,
+  Shared = 3,
+  Const = 4,
+  Local = 5,
+  SharedCluster = 7,
+
+  // NVPTX Backend Private:
+  Param = 101
+};
+
+} // namespace llvm::NVPTX
+
+#endif // LLVM_TARGETPARSER_NVPTX_ADDRESS_SPACES_H

llvm/include/llvm/TargetParser/Triple.h

@@ -1362,6 +1362,9 @@ class Triple {
   /// Compute the LLVM IR data layout string based on the triple. Some targets
   /// customize the layout based on the ABIName string.
   LLVM_ABI std::string computeDataLayout(StringRef ABIName = "") const;
+
+  LLVM_ABI StringRef getAddressSpaceName(unsigned AS) const;
+  LLVM_ABI std::optional<unsigned> getAddressSpaceNumber(StringRef Name) const;
 };
 
 } // End llvm namespace

llvm/lib/AsmParser/LLParser.cpp

@@ -1965,6 +1965,9 @@ bool LLParser::parseOptionalAddrSpace(unsigned &AddrSpace, unsigned DefaultAS) {
         AddrSpace = M->getDataLayout().getDefaultGlobalsAddressSpace();
       } else if (AddrSpaceStr == "P") {
         AddrSpace = M->getDataLayout().getProgramAddressSpace();
+      } else if (std::optional<unsigned> AS =
+                     M->getTargetTriple().getAddressSpaceNumber(AddrSpaceStr)) {
+        AddrSpace = *AS;
       } else {
         return tokError("invalid symbolic addrspace '" + AddrSpaceStr + "'");
       }

llvm/lib/IR/AsmWriter.cpp

@@ -543,7 +543,8 @@ namespace {
 
 class TypePrinting {
 public:
-  TypePrinting(const Module *M = nullptr) : DeferredM(M) {}
+  TypePrinting(const Module *M = nullptr)
+      : M(M), TypesIncorporated(M == nullptr) {}
 
   TypePrinting(const TypePrinting &) = delete;
   TypePrinting &operator=(const TypePrinting &) = delete;
@@ -564,7 +565,8 @@ class TypePrinting {
   void incorporateTypes();
 
   /// A module to process lazily when needed. Set to nullptr as soon as used.
-  const Module *DeferredM;
+  const Module *M;
+  bool TypesIncorporated;
 
   TypeFinder NamedTypes;
 
@@ -605,11 +607,11 @@ bool TypePrinting::empty() {
 }
 
 void TypePrinting::incorporateTypes() {
-  if (!DeferredM)
+  if (TypesIncorporated)
     return;
 
-  NamedTypes.run(*DeferredM, false);
-  DeferredM = nullptr;
+  NamedTypes.run(*M, false);
+  TypesIncorporated = true;
 
   // The list of struct types we got back includes all the struct types, split
   // the unnamed ones out to a numbering and remove the anonymous structs.
@@ -630,6 +632,20 @@ void TypePrinting::incorporateTypes() {
   NamedTypes.erase(NextToUse, NamedTypes.end());
 }
 
+static void printAddressSpace(const Module *M, unsigned AS, raw_ostream &OS,
+                              StringRef Prefix = " ", StringRef Suffix = "",
+                              bool ForcePrint = false) {
+  if (AS == 0 && !ForcePrint)
+    return;
+  OS << Prefix << "addrspace(";
+  StringRef ASName = M ? M->getTargetTriple().getAddressSpaceName(AS) : "";
+  if (!ASName.empty())
+    OS << "\"" << ASName << "\"";
+  else
+    OS << AS;
+  OS << ")" << Suffix;
+}
+
 /// Write the specified type to the specified raw_ostream, making use of type
 /// names or up references to shorten the type name where possible.
 void TypePrinting::print(Type *Ty, raw_ostream &OS) {
@@ -686,8 +702,7 @@ void TypePrinting::print(Type *Ty, raw_ostream &OS) {
   case Type::PointerTyID: {
     PointerType *PTy = cast<PointerType>(Ty);
     OS << "ptr";
-    if (unsigned AddressSpace = PTy->getAddressSpace())
-      OS << " addrspace(" << AddressSpace << ')';
+    printAddressSpace(M, PTy->getAddressSpace(), OS);
     return;
   }
   case Type::ArrayTyID: {
@@ -3896,10 +3911,10 @@ void AssemblyWriter::printGlobal(const GlobalVariable *GV) {
   printThreadLocalModel(GV->getThreadLocalMode(), Out);
   StringRef UA = getUnnamedAddrEncoding(GV->getUnnamedAddr());
   if (!UA.empty())
-      Out << UA << ' ';
+    Out << UA << ' ';
 
-  if (unsigned AddressSpace = GV->getType()->getAddressSpace())
-    Out << "addrspace(" << AddressSpace << ") ";
+  printAddressSpace(GV->getParent(), GV->getType()->getAddressSpace(), Out,
+                    /*Prefix=*/"", /*Suffix=*/" ");
   if (GV->isExternallyInitialized()) Out << "externally_initialized ";
   Out << (GV->isConstant() ? "constant " : "global ");
   TypePrinter.print(GV->getValueType(), Out);
@@ -4174,9 +4189,10 @@ void AssemblyWriter::printFunction(const Function *F) {
   // a module with a non-zero program address space or if there is no valid
   // Module* so that the file can be parsed without the datalayout string.
   const Module *Mod = F->getParent();
-  if (F->getAddressSpace() != 0 || !Mod ||
-      Mod->getDataLayout().getProgramAddressSpace() != 0)
-    Out << " addrspace(" << F->getAddressSpace() << ")";
+  bool ForcePrintAddressSpace =
+      !Mod || Mod->getDataLayout().getProgramAddressSpace() != 0;
+  printAddressSpace(Mod, F->getAddressSpace(), Out, /*Prefix=*/" ",
+                    /*Suffix=*/"", ForcePrintAddressSpace);
   if (Attrs.hasFnAttrs())
     Out << " #" << Machine.getAttributeGroupSlot(Attrs.getFnAttrs());
   if (F->hasSection()) {
@@ -4358,17 +4374,14 @@ static void maybePrintCallAddrSpace(const Value *Operand, const Instruction *I,
     return;
   }
   unsigned CallAddrSpace = Operand->getType()->getPointerAddressSpace();
-  bool PrintAddrSpace = CallAddrSpace != 0;
-  if (!PrintAddrSpace) {
-    const Module *Mod = getModuleFromVal(I);
-    // We also print it if it is zero but not equal to the program address space
-    // or if we can't find a valid Module* to make it possible to parse
-    // the resulting file even without a datalayout string.
-    if (!Mod || Mod->getDataLayout().getProgramAddressSpace() != 0)
-      PrintAddrSpace = true;
-  }
-  if (PrintAddrSpace)
-    Out << " addrspace(" << CallAddrSpace << ")";
+  const Module *Mod = getModuleFromVal(I);
+  // We also print it if it is zero but not equal to the program address space
+  // or if we can't find a valid Module* to make it possible to parse
+  // the resulting file even without a datalayout string.
+  bool ForcePrintAddrSpace =
+      !Mod || Mod->getDataLayout().getProgramAddressSpace() != 0;
+  printAddressSpace(Mod, CallAddrSpace, Out, /*Prefix=*/" ", /*Suffix=*/"",
+                    ForcePrintAddrSpace);
 }
 
 // This member is called for each Instruction in a function..
@@ -4735,9 +4748,8 @@ void AssemblyWriter::printInstruction(const Instruction &I) {
       Out << ", align " << A->value();
     }
 
-    unsigned AddrSpace = AI->getAddressSpace();
-    if (AddrSpace != 0)
-      Out << ", addrspace(" << AddrSpace << ')';
+    printAddressSpace(AI->getModule(), AI->getAddressSpace(), Out,
+                      /*Prefix=*/", ");
   } else if (isa<CastInst>(I)) {
     if (Operand) {
       Out << ' ';

llvm/lib/Target/NVPTX/NVPTX.h

@@ -20,6 +20,8 @@
 #include "llvm/Support/AtomicOrdering.h"
 #include "llvm/Support/CodeGen.h"
 #include "llvm/Target/TargetMachine.h"
+#include "llvm/TargetParser/NVPTXAddressSpaces.h"
+
 namespace llvm {
 class FunctionPass;
 class MachineFunctionPass;
@@ -176,19 +178,6 @@ enum Scope : ScopeUnderlyingType {
   LASTSCOPE = DefaultDevice
 };
 
-using AddressSpaceUnderlyingType = unsigned int;
-enum AddressSpace : AddressSpaceUnderlyingType {
-  Generic = 0,
-  Global = 1,
-  Shared = 3,
-  Const = 4,
-  Local = 5,
-  SharedCluster = 7,
-
-  // NVPTX Backend Private:
-  Param = 101
-};
-
 namespace PTXLdStInstCode {
 enum FromType { Unsigned = 0, Signed, Float, Untyped };
 } // namespace PTXLdStInstCode

llvm/lib/TargetParser/Triple.cpp

@@ -17,6 +17,7 @@
 #include "llvm/TargetParser/ARMTargetParser.h"
 #include "llvm/TargetParser/ARMTargetParserCommon.h"
 #include "llvm/TargetParser/Host.h"
+#include "llvm/TargetParser/NVPTXAddressSpaces.h"
 #include <cassert>
 #include <cstring>
 using namespace llvm;
@@ -2361,6 +2362,56 @@ ExceptionHandling Triple::getDefaultExceptionHandling() const {
   return ExceptionHandling::None;
 }
 
+static StringRef getNVPTXAddressSpaceName(unsigned AS) {
+  switch (AS) {
+  case NVPTX::Generic:
+    return "generic";
+  case NVPTX::Global:
+    return "global";
+  case NVPTX::Shared:
+    return "shared";
+  case NVPTX::Const:
+    return "const";
+  case NVPTX::Local:
+    return "local";
+  case NVPTX::SharedCluster:
+    return "shared_cluster";
+  default:
+    return "";
+  }
+}
+
+static std::optional<unsigned> getNVPTXAddressSpaceNumber(StringRef Name) {
+  return StringSwitch<std::optional<unsigned>>(Name)
+      .Case("generic", NVPTX::Generic)
+      .Case("global", NVPTX::Global)
+      .Case("shared", NVPTX::Shared)
+      .Case("const", NVPTX::Const)
+      .Case("local", NVPTX::Local)
+      .Case("shared_cluster", NVPTX::SharedCluster)
+      .Default(std::nullopt);
+}
+
+StringRef Triple::getAddressSpaceName(unsigned AS) const {
+  switch (getArch()) {
+  case Triple::nvptx:
+  case Triple::nvptx64:
+    return getNVPTXAddressSpaceName(AS);
+  default:
+    return "";
+  }
+}
+
+std::optional<unsigned> Triple::getAddressSpaceNumber(StringRef Name) const {
+  switch (getArch()) {
+  case Triple::nvptx:
+  case Triple::nvptx64:
+    return getNVPTXAddressSpaceNumber(Name);
+  default:
+    return std::nullopt;
+  }
+}
+
 // HLSL triple environment orders are relied on in the front end
 static_assert(Triple::Vertex - Triple::Pixel == 1,
               "incorrect HLSL stage order");

llvm/test/Assembler/symbolic-addrspace-nvptx.ll

@@ -0,0 +1,54 @@
+;; Check that we can parse/print symbolic addres space for NVPTX triple. 
+; RUN: split-file %s %t --leading-lines
+; RUN: llvm-as < %t/num-to-sym.ll | llvm-dis | FileCheck %t/num-to-sym.ll
+; RUN: llvm-as < %t/sym-to-sym.ll | llvm-dis | FileCheck %t/sym-to-sym.ll
+; RUN: not llvm-as < %t/invalid-sym.ll 2>&1 | FileCheck %t/invalid-sym.ll
+
+;--- num-to-sym.ll
+target triple = "nvptx64-nvidia-cuda"
+
+; CHECK: @str0 = private addrspace("global") constant
+@str0 = private addrspace(1) constant [4 x i8] c"str\00"
+
+; CHECK: @gvs = external addrspace("shared") global i32
+@gvs = external addrspace(3) global i32
+
+; CHECK: @gv_unknown_as = external addrspace(100) global i32
+@gv_unknown_as = external addrspace(100) global i32
+
+define void @foo() {
+  ; CHECK: %alloca = alloca i32, align 4, addrspace("local")
+  %alloca = alloca i32, addrspace(5)
+  ; CHECK: store i32 0, ptr addrspace("local") %alloca
+  store i32 0, ptr addrspace(5) %alloca
+  ; CHECK: store i32 3, ptr addrspace("shared") @gvs
+  store i32 3, ptr addrspace(3) @gvs
+  ret void
+}
+
+;--- sym-to-sym.ll
+target triple = "nvptx64-nvidia-cuda"
+
+; CHECK: @str0 = private addrspace("global") constant
+@str0 = private addrspace("global") constant [4 x i8] c"str\00"
+
+; CHECK: @gvs = external addrspace("shared") global i32
+@gvs = external addrspace("shared") global i32
+
+; CHECK: @gv_unknown_as = external addrspace(100) global i32
+@gv_unknown_as = external addrspace(100) global i32
+
+define void @foo() {
+  ; CHECK: %alloca = alloca i32, align 4, addrspace("local")
+  %alloca = alloca i32, addrspace("local")
+  ; CHECK: store i32 0, ptr addrspace("local") %alloca
+  store i32 0, ptr addrspace("local") %alloca
+  ; CHECK: store i32 3, ptr addrspace("shared") @gvs
+  store i32 3, ptr addrspace("shared") @gvs
+  ret void
+}
+
+;--- invalid-sym.ll
+target triple = "nvptx64-nvidia-cuda"
+; CHECK: error: invalid symbolic addrspace 'ram'
+@str0 = private addrspace("ram") constant [4 x i8] c"str\00"