[NVPTX][NFC] Simplify printing initialization of aggregates

[ikudrin]

Summary

This simplifies NVPTXAsmPrinter::AggBuffer and its usage. It is also a preparation for D127504.

[tra]

llvm/lib/Target/NVPTX/NVPTXAsmPrinter.h

139–144

Nit: printAsBytes/ printAsWords or just printBytes/printWords.

[ikudrin]

Thanks!

• Rename printInBytes() -> printBytes(), printInWords() -> printWords()`
• Remove an raw_ostream argument in the constructor of AggBuffer; pass the stream directly to printBytes()/printWords()
• Extract printSymbol() in this patch to further simplify D127504
• isNonGenericPointer -> isGenericPointer

llvm/lib/Target/NVPTX/NVPTXAsmPrinter.h

class AggBuffer {
     // those symbol addresses. For the i-th symbol address in the
     //aggregate, its corresponding 4-byte or 8-byte elements in 'buffer'
     // are filled with 0s. symbolPosInBuffer[i-1] records its position
     // in 'buffer', and Symbols[i-1] records the Value*.
     //
     // Once we have this AggBuffer setup, we can choose how to print
     // it out.
   public:
-    unsigned numSymbols;   // number of symbol addresses
+    // number of symbol addresses
+    unsigned numSymbols() const { return Symbols.size(); }
 
   private:
     const unsigned size;   // size of the buffer in bytes
     std::vector<unsigned char> buffer; // the buffer
     SmallVector<unsigned, 4> symbolPosInBuffer;
     SmallVector<const Value *, 4> Symbols;
     // SymbolsBeforeStripping[i] is the original form of Symbols[i] before
     // stripping pointer casts, i.e.,
     // Symbols[i] == SymbolsBeforeStripping[i]->stripPointerCasts().
     //
     // We need to keep these values because AggBuffer::print decides whether to
     // emit a "generic()" cast for Symbols[i] depending on the address space of
     // SymbolsBeforeStripping[i].
     SmallVector<const Value *, 4> SymbolsBeforeStripping;
     unsigned curpos;
-    raw_ostream &O;
     NVPTXAsmPrinter &AP;
     bool EmitGeneric;
 
   public:
-    AggBuffer(unsigned size, raw_ostream &O, NVPTXAsmPrinter &AP)
-        : size(size), buffer(size), O(O), AP(AP) {
+    AggBuffer(unsigned size, NVPTXAsmPrinter &AP)
+        : size(size), buffer(size), AP(AP) {
       curpos = 0;
-      numSymbols = 0;
       EmitGeneric = AP.EmitGeneric;
     }
 
     // Copy Num bytes from Ptr.
     // if Bytes > Num, zero fill up to Bytes.
     unsigned addBytes(unsigned char *Ptr, int Num, int Bytes) {
       assert((curpos + Num) <= size);
       assert((curpos + Bytes) <= size);
unsigned addZeros(int Num) {
       }
       return curpos;
     }
 
     void addSymbol(const Value *GVar, const Value *GVarBeforeStripping) {
       symbolPosInBuffer.push_back(curpos);
       Symbols.push_back(GVar);
       SymbolsBeforeStripping.push_back(GVarBeforeStripping);
-      numSymbols++;
     }
 
-    void print() {
-      if (numSymbols == 0) {
-        // print out in bytes
-        for (unsigned i = 0; i < size; i++) {
-          if (i)
-            O << ", ";
-          O << (unsigned int) buffer[i];
-        }
-      } else {
-        // print out in 4-bytes or 8-bytes
-        unsigned int pos = 0;
-        unsigned int nSym = 0;
-        unsigned int nextSymbolPos = symbolPosInBuffer[nSym];
-        unsigned int nBytes = 4;
-        if (static_cast<const NVPTXTargetMachine &>(AP.TM).is64Bit())
-          nBytes = 8;
-        for (pos = 0; pos < size; pos += nBytes) {
-          if (pos)
-            O << ", ";
-          if (pos == nextSymbolPos) {
-            const Value *v = Symbols[nSym];
-            const Value *v0 = SymbolsBeforeStripping[nSym];
-            if (const GlobalValue *GVar = dyn_cast<GlobalValue>(v)) {
-              MCSymbol *Name = AP.getSymbol(GVar);
-              PointerType *PTy = dyn_cast<PointerType>(v0->getType());
-              bool IsNonGenericPointer = false; // Is v0 a non-generic pointer?
-              if (PTy && PTy->getAddressSpace() != 0) {
-                IsNonGenericPointer = true;
-              }
-              if (EmitGeneric && !isa<Function>(v) && !IsNonGenericPointer) {
-                O << "generic(";
-                Name->print(O, AP.MAI);
-                O << ")";
-              } else {
-                Name->print(O, AP.MAI);
-              }
-            } else if (const ConstantExpr *CExpr = dyn_cast<ConstantExpr>(v0)) {
-              const MCExpr *Expr =
-                AP.lowerConstantForGV(cast<Constant>(CExpr), false);
-              AP.printMCExpr(*Expr, O);
-            } else
-              llvm_unreachable("symbol type unknown");
-            nSym++;
-            if (nSym >= numSymbols)
-              nextSymbolPos = size + 1;
-            else
-              nextSymbolPos = symbolPosInBuffer[nSym];
-          } else if (nBytes == 4)
-            O << *(unsigned int *)(&buffer[pos]);
-          else
-            O << *(unsigned long long *)(&buffer[pos]);
-        }
-      }
-    }
+    void printBytes(raw_ostream &os);
+    void printWords(raw_ostream &os);
+
+  private:
+    void printSymbol(unsigned nSym, raw_ostream &os);
   };
 
   friend class AggBuffer;
 
 private:
   StringRef getPassName() const override { return "NVPTX Assembly Printer"; }
 
   const Function *F;
   std::string CurrentFnName;

llvm/lib/Target/NVPTX/NVPTXAsmPrinter.cpp

case Type::FixedVectorTyID:
       ElementSize = DL.getTypeStoreSize(ETy);
       // Ptx allows variable initilization only for constant and
       // global state spaces.
       if (((PTy->getAddressSpace() == ADDRESS_SPACE_GLOBAL) ||
            (PTy->getAddressSpace() == ADDRESS_SPACE_CONST)) &&
           GVar->hasInitializer()) {
         const Constant *Initializer = GVar->getInitializer();
         if (!isa<UndefValue>(Initializer) && !Initializer->isNullValue()) {
-          AggBuffer aggBuffer(ElementSize, O, *this);
+          AggBuffer aggBuffer(ElementSize, *this);
           bufferAggregateConstant(Initializer, &aggBuffer);
-          if (aggBuffer.numSymbols) {
-            if (static_cast<const NVPTXTargetMachine &>(TM).is64Bit()) {
-              O << " .u64 ";
-              getSymbol(GVar)->print(O, MAI);
-              O << "[";
-              O << ElementSize / 8;
-            } else {
-              O << " .u32 ";
-              getSymbol(GVar)->print(O, MAI);
-              O << "[";
-              O << ElementSize / 4;
-            }
-            O << "]";
+          if (aggBuffer.numSymbols()) {
+            unsigned int ptrSize = MAI->getCodePointerSize();
+            O << " .u" << ptrSize * 8 << " ";
+            O << "[" << ElementSize / ptrSize << "] = {";
+            aggBuffer.printWords(O);
+            O << "}";
           } else {
             O << " .b8 ";
             getSymbol(GVar)->print(O, MAI);
-            O << "[";
-            O << ElementSize;
-            O << "]";
-          }
-          O << " = {";
-          aggBuffer.print();
+            O << "[" << ElementSize << "] = {";
+            aggBuffer.printBytes(O);
             O << "}";
+          }
         } else {
           O << " .b8 ";
           getSymbol(GVar)->print(O, MAI);
           if (ElementSize) {
             O << "[";
             O << ElementSize;
             O << "]";
           }
case Type::FixedVectorTyID:
       break;
     default:
       llvm_unreachable("type not supported yet");
     }
   }
   O << ";\n";
 }
 
+void NVPTXAsmPrinter::AggBuffer::printSymbol(unsigned nSym, raw_ostream &os) {
+  const Value *v = Symbols[nSym];
+  const Value *v0 = SymbolsBeforeStripping[nSym];
+  if (const GlobalValue *GVar = dyn_cast<GlobalValue>(v)) {
+    MCSymbol *Name = AP.getSymbol(GVar);
+    PointerType *PTy = dyn_cast<PointerType>(v0->getType());
+    // Is v0 a generic pointer?
+    bool isGenericPointer = PTy && PTy->getAddressSpace() == 0;
+    if (EmitGeneric && isGenericPointer && !isa<Function>(v)) {
+      os << "generic(";
+      Name->print(os, AP.MAI);
+      os << ")";
+    } else {
+      Name->print(os, AP.MAI);
+    }
+  } else if (const ConstantExpr *CExpr = dyn_cast<ConstantExpr>(v0)) {
+    const MCExpr *Expr = AP.lowerConstantForGV(cast<Constant>(CExpr), false);
+    AP.printMCExpr(*Expr, os);
+  } else
+    llvm_unreachable("symbol type unknown");
+}
+
+void NVPTXAsmPrinter::AggBuffer::printBytes(raw_ostream &os) {
+  for (unsigned int pos = 0; pos < size; ++pos) {
+    if (pos)
+      os << ", ";
+    os << (unsigned int)buffer[pos];
+  }
+}
+
+void NVPTXAsmPrinter::AggBuffer::printWords(raw_ostream &os) {
+  unsigned int ptrSize = AP.MAI->getCodePointerSize();
+  symbolPosInBuffer.push_back(size);
+  unsigned int nSym = 0;
+  unsigned int nextSymbolPos = symbolPosInBuffer[nSym];
+  assert(nextSymbolPos % ptrSize == 0);
+  for (unsigned int pos = 0; pos < size; pos += ptrSize) {
+    if (pos)
+      os << ", ";
+    if (pos == nextSymbolPos) {
+      printSymbol(nSym, os);
+      nextSymbolPos = symbolPosInBuffer[++nSym];
+      assert(nextSymbolPos % ptrSize == 0);
+      assert(nextSymbolPos >= pos + ptrSize);
+    } else if (ptrSize == 4)
+      os << *(uint32_t *)(&buffer[pos]);
+    else
+      os << *(uint64_t *)(&buffer[pos]);
+  }
+}
+
 void NVPTXAsmPrinter::emitDemotedVars(const Function *f, raw_ostream &O) {
   if (localDecls.find(f) == localDecls.end())
     return;
 
   std::vector<const GlobalVariable *> &gvars = localDecls[f];
 
   const NVPTXTargetMachine &NTM = static_cast<const NVPTXTargetMachine &>(TM);
   const NVPTXSubtarget &STI =