[NVPTX] Use the mask() operator to initialize packed structs with pointers

[ikudrin]

Summary

The current implementation assumes that all pointers used in the initialization of an aggregate are aligned according to the pointer size of the target; that might not be so if the object is packed. In that case, an array of .u8 should be used and pointers should be decorated with the mask() operator.

The operator was introduced in PTX ISA 7.1, so an error is issued if the case is detected for an earlier version.

[tra]

llvm/lib/Target/NVPTX/NVPTXAsmPrinter.cpp

1173–1174

What will happen if we have an aggregate with a mix of packed and unpacked fields and we only want to init with a pointer located in the unpacked part of it? Will something like this trigger the error?

%ti = type <{ i8, i32 }>
%to = type { i8, %ti, i32, void ()* }
@n = addrspace(1) global %to {
    i8 12,
    %ti <{ i8 56, i32 78 }>,
    i32 34,
    void()* @func
    }

It's one of the test cases below with another pointer added to non-packed top-level struct.

llvm/test/CodeGen/NVPTX/packed-aggr.ll

51

Interesting. Right now we're printing '0' for the function pointer, which is indeed not good at all. https://godbolt.org/z/KEexEar5o

[ikudrin]

llvm/lib/Target/NVPTX/NVPTXAsmPrinter.cpp

1173–1174

Unfortunately, yes, that would trigger the error. On the other hand, supporting that case would require a deeper analysis of the initialization expression itself, not only the type, because there can be ptrtoint operators. I'm not sure if we need that complication, considering that CUDA 11.1 is almost 2 years old already.

llvm/test/CodeGen/NVPTX/packed-aggr.ll

51

That's true. Currently, the implementation does not expect pointers to be unaligned.

[ikudrin]

Hi @tra Artem, do you think the said complicated deeper analysis is profitable?

llvm/lib/Target/NVPTX/NVPTXAsmPrinter.cpp

1173–1174

Unfortunately, yes, that would trigger the error. On the other hand, supporting that case would require a deeper analysis of the initialization expression itself, not only the type, because there can be ptrtoint operators. I'm not sure if we need that complication, considering that CUDA 11.1 is almost 2 years old already.

[tra]

llvm/lib/Target/NVPTX/NVPTXAsmPrinter.cpp

1173–1174

I think we need to fix this. I'm not convinced that recursively checking fields' 'packed'-ness makes sense here.

I think we only care whether the structure with the field we need to init is packed or not. If it's not packed, then per-field alignment guarantees still apply and that's all we need.
E.g. in the example above, it does not matter whether any of the fields itself are packed or not. The pointer field alignment is only affected by its containing aggregate itself.

Perhaps I'm missing something. If we do need recursive check, could you give me an example where that would make a difference?

[ikudrin]

Rework the implementation so that the mask() operator is used only if it is indeed required.

[tra]

Nice. I like this implementation better.

llvm/lib/Target/NVPTX/NVPTXAsmPrinter.cpp

1180

We should probably have a test case triggering this.

Also, we may want to move the check into printInBytes, as otherwise it may be possible to invoke it when mask() is not supported.

1183

requires at least

1236–1237

It may be a bit easier to read this way:

bool isGenericPointer = PTy && PTy->getAddressSpace() == 0;
  if (EmitGeneric  && isGenericPointer && !isa<Function>(v)) {
     // wrap in generic();
 } else {
     ...
 }

Comment above would need to be updated, too.

1265

I'd add a comment that we're generating a per-byte mask operator here with an example illustrating it.
PTX's mask() is weird due to the fact that it actually has no mask in it, so an explanation would be useful here.

llvm/lib/Target/NVPTX/NVPTXAsmPrinter.h

84–87

I'd suggest rephrasing it as "all are aligned", which would closer match our intent here vs "none are non-aligned" we have now which is equivalent, but is a bit harder to read.

std::all_of( ... return pos % ptrSize == 0; )

[ikudrin]

Thanks!

• Add a comment about the mask() operator
• Update the error message
• std::none_of() -> std::all_of()

[ikudrin]

• 0xff -> 0xFF in the comment

llvm/lib/Target/NVPTX/NVPTXAsmPrinter.cpp

1180

We should probably have a test case triggering this.

The first test in packed-aggr.ll checks the error.

Also, we may want to move the check into printInBytes, as otherwise it may be possible to invoke it when mask() is not supported.

That would require passing GVar and STI there, only to check the requirements and to print the error message. Not sure it's worth it.

1236–1237

Done in D129773

[tra]

LGTM with few nits.

llvm/lib/Target/NVPTX/NVPTXAsmPrinter.cpp

1275

Nit: splitting pos incrementing between the for loop and the body looks a bit odd.

It may be a bit better to increment pos where we actually produce output:
+=1 where we do byte-by-byte output and += ptrSize where we output whole pointer.

llvm/lib/Target/NVPTX/NVPTXAsmPrinter.h

85

Nit: we would use llvm::all_of(symbolPosInBuffer, ...) and make it a bit more concise.

llvm/test/CodeGen/NVPTX/packed-aggr.ll

4

I've missed this test case when I reviewed previous revision of this patch. ERR and RUN kind of blend together.
Maybe change it to ERROR to make it easier to see.
I'd also move it down to where the s1 is defined.

[ikudrin]

Thanks! I'm going to land that on Monday.

• Fix nits

[tra]

Looks like the patch got reverted and reapplied. Could you tell me what went wrong and what was changed to fix it?

[ikudrin]

In D127504#3666230, @tra wrote:

Looks like the patch got reverted and reapplied. Could you tell me what went wrong and what was changed to fix it?

The test failed on BE bots. Direct reading u32 and u64 values through a C-style pointer cast made unexpected results, thus I had to replace them with support::endian::read32le() and read64le().

llvm/lib/Target/NVPTX/NVPTXAsmPrinter.h

 
 namespace llvm {
 
 class MCOperand;
 
 class LLVM_LIBRARY_VISIBILITY NVPTXAsmPrinter : public AsmPrinter {
 
   class AggBuffer {
-    // Used to buffer the emitted string for initializing global
-    // aggregates.
+    // Used to buffer the emitted string for initializing global aggregates.
     //
-    // Normally an aggregate (array, vector or structure) is emitted
-    // as a u8[]. However, if one element/field of the aggregate
-    // is a non-NULL address, then the aggregate is emitted as u32[]
-    // or u64[].
-    //
-    // We first layout the aggregate in 'buffer' in bytes, except for
-    // 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*.
+    // Normally an aggregate (array, vector, or structure) is emitted as a u8[].
+    // However, if either element/field of the aggregate is a non-NULL address,
+    // and all such addresses are properly aligned, then the aggregate is
+    // emitted as u32[] or u64[]. In the case of unaligned addresses, the
+    // aggregate is emitted as u8[], and the mask() operator is used for all
+    // pointers.
+    //
+    // We first layout the aggregate in 'buffer' in bytes, except for 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.
+    // Once we have this AggBuffer setup, we can choose how to print it out.
   public:
     // number of symbol addresses
     unsigned numSymbols() const { return Symbols.size(); }
 
+    bool allSymbolsAligned(unsigned ptrSize) const {
+      return std::all_of(symbolPosInBuffer.begin(), symbolPosInBuffer.end(),
+                         [=](unsigned pos) { return pos % ptrSize == 0; });
+    }
+
   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().

llvm/lib/Target/NVPTX/NVPTXAsmPrinter.cpp

 #include "llvm/MC/MCInstrDesc.h"
 #include "llvm/MC/MCStreamer.h"
 #include "llvm/MC/MCSymbol.h"
 #include "llvm/MC/TargetRegistry.h"
 #include "llvm/Support/Casting.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/ErrorHandling.h"
 #include "llvm/Support/MachineValueType.h"
+#include "llvm/Support/NativeFormatting.h"
 #include "llvm/Support/Path.h"
 #include "llvm/Support/raw_ostream.h"
 #include "llvm/Target/TargetLoweringObjectFile.h"
 #include "llvm/Target/TargetMachine.h"
 #include "llvm/Transforms/Utils/UnrollLoop.h"
 #include <cassert>
 #include <cstdint>
 #include <cstring>
case Type::FixedVectorTyID:
       // 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, *this);
           bufferAggregateConstant(Initializer, &aggBuffer);
           if (aggBuffer.numSymbols()) {
             unsigned int ptrSize = MAI->getCodePointerSize();
+            if (ElementSize % ptrSize ||
+                !aggBuffer.allSymbolsAligned(ptrSize)) {
+              // Print in bytes and use the mask() operator for pointers.
+              if (!STI.hasMaskOperator())
+                report_fatal_error(
+                    "initialized packed aggregate with pointers '" +
+                    GVar->getName() +
+                    "' requires at least PTX ISA version 7.1");
+              O << " .u8 ";
+              getSymbol(GVar)->print(O, MAI);
+              O << "[" << ElementSize << "] = {";
+              aggBuffer.printBytes(O);
+              O << "}";
+            } else {
               O << " .u" << ptrSize * 8 << " ";
               getSymbol(GVar)->print(O, MAI);
               O << "[" << ElementSize / ptrSize << "] = {";
               aggBuffer.printWords(O);
               O << "}";
+            }
           } else {
             O << " .b8 ";
             getSymbol(GVar)->print(O, MAI);
             O << "[" << ElementSize << "] = {";
             aggBuffer.printBytes(O);
             O << "}";
           }
         } else {

 
 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) {
+  unsigned int ptrSize = AP.MAI->getCodePointerSize();
+  symbolPosInBuffer.push_back(size);
+  unsigned int nSym = 0;
+  unsigned int nextSymbolPos = symbolPosInBuffer[nSym];
   for (unsigned int pos = 0; pos < size; ++pos) {
     if (pos)
       os << ", ";
+    if (pos != nextSymbolPos) {
       os << (unsigned int)buffer[pos];
+      continue;
+    }
+    // Generate a per-byte mask() operator for the symbol, which looks like:
+    //   .global .u8 addr[] = {0xff(foo), 0xff00(foo), 0xff0000(foo), ...};
+    // See https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#initializers
+    std::string symText;
+    llvm::raw_string_ostream oss(symText);
+    printSymbol(nSym, oss);
+    for (unsigned i = 0; i < ptrSize; ++i) {
+      if (i)
+        os << ", ";
+      llvm::write_hex(os, 0xFFULL << i * 8, HexPrintStyle::PrefixUpper);
+      os << "(" << symText << ")";
+    }
+    pos += ptrSize - 1;
+    nextSymbolPos = symbolPosInBuffer[++nSym];
+    assert(pos < nextSymbolPos);
   }
 }
 
 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];

llvm/lib/Target/NVPTX/NVPTXSubtarget.h

public:
   bool hasAtomScope() const { return SmVersion >= 60; }
   bool hasAtomBitwise64() const { return SmVersion >= 32; }
   bool hasAtomMinMax64() const { return SmVersion >= 32; }
   bool hasLDG() const { return SmVersion >= 32; }
   inline bool hasHWROT32() const { return SmVersion >= 32; }
   bool hasImageHandles() const;
   bool hasFP16Math() const { return SmVersion >= 53; }
   bool allowFP16Math() const;
+  bool hasMaskOperator() const { return PTXVersion >= 71; }
   unsigned int getSmVersion() const { return SmVersion; }
   std::string getTargetName() const { return TargetName; }
 
   unsigned getPTXVersion() const { return PTXVersion; }
 
   NVPTXSubtarget &initializeSubtargetDependencies(StringRef CPU, StringRef FS);
   void ParseSubtargetFeatures(StringRef CPU, StringRef TuneCPU, StringRef FS);
 };
 
 } // End llvm namespace
 
 #endif

llvm/test/CodeGen/NVPTX/packed-aggr.ll

+; RUN: not --crash llc < %s -march=nvptx -mcpu=sm_20 -mattr=+ptx70 2>&1 | \
+; RUN:   FileCheck %s --check-prefix=ERR
+; ERR: initialized packed aggregate with pointers 's1' requires at least PTX ISA version 7.1
+
+; RUN: llc < %s -march=nvptx -mcpu=sm_20 -mattr=+ptx71 | \
+; RUN:   FileCheck %s --check-prefixes=CHECK,CHECK32
+; RUN: llc < %s -march=nvptx64 -mcpu=sm_20 -mattr=+ptx71 | \
+; RUN:   FileCheck %s --check-prefixes=CHECK,CHECK64
+; RUN: %if ptxas-11.1 %{ llc < %s -march=nvptx -mcpu=sm_20 -mattr=+ptx71 | %ptxas-verify %}
+; RUN: %if ptxas-11.1 %{ llc < %s -march=nvptx64 -mcpu=sm_20 -mattr=+ptx71 | %ptxas-verify %}
+
+;; Test that packed structs with symbol references are represented using the
+;; mask() operator.
+
+declare void @func()
+@p = addrspace(1) global i8 0
+; CHECK: .extern .func func
+; CHECK: .u8 p;
+
+%t1 = type <{ i16, i8*, i8, void ()*, i8*, i32 }>
+@s1 = addrspace(1) global %t1 <{
+; CHECK32: .global .align 1 .u8 s1[19] = {
+; CHECK64: .global .align 1 .u8 s1[31] = {
+    i16 12,
+; CHECK-SAME:   12, 0,
+    i8* addrspacecast (i8 addrspace(1)* @p to i8*),
+; CHECK-SAME:   0xFF(generic(p)), 0xFF00(generic(p)), 0xFF0000(generic(p)), 0xFF000000(generic(p)),
+; CHECK64-SAME: 0xFF00000000(generic(p)), 0xFF0000000000(generic(p)), 0xFF000000000000(generic(p)), 0xFF00000000000000(generic(p)),
+    i8 34,
+; CHECK-SAME:   34
+    void ()* @func,
+; CHECK-SAME:   0xFF(func), 0xFF00(func), 0xFF0000(func), 0xFF000000(func),
+; CHECK64-SAME: 0xFF00000000(func), 0xFF0000000000(func), 0xFF000000000000(func), 0xFF00000000000000(func),
+    i8* addrspacecast (i8 addrspace(1)* getelementptr (i8, i8 addrspace(1)* @p, i32 3) to i8*),
+; CHECK-SAME:   0xFF(generic(p)+3), 0xFF00(generic(p)+3), 0xFF0000(generic(p)+3), 0xFF000000(generic(p)+3),
+; CHECK64-SAME: 0xFF00000000(generic(p)+3), 0xFF0000000000(generic(p)+3), 0xFF000000000000(generic(p)+3), 0xFF00000000000000(generic(p)+3),
+    i32 56 }>, align 1
+; CHECK-SAME:   56, 0, 0, 0};
+
+;; Test a case than an unaligned pointer is in a nested struct.
+
+%t2i = type <{ void ()* }>
+%t2o = type { i8, %t2i, i32 }
+@s2 = addrspace(1) global %t2o {
+; CHECK32: .global .align 8 .u8 s2[12] = {
+; CHECK64: .global .align 8 .u8 s2[16] = {
+    i8 12,
+; CHECK-SAME:   12,
+    %t2i <{ void()* @func }>,
+; CHECK-SAME:   0xFF(func), 0xFF00(func), 0xFF0000(func), 0xFF000000(func),
+; CHECK64-SAME: 0xFF00000000(func), 0xFF0000000000(func), 0xFF000000000000(func), 0xFF00000000000000(func),
+    i32 34}
+; CHECK-SAME:   0, 0, 0,
+; CHECK-SAME:   34, 0, 0, 0};
+
+;; Test that a packed struct which size is not multiple of the pointer size
+;; is printed in bytes and uses the mask() operator for pointers even though
+;; the pointers are aligned.
+
+%t3 = type <{ void ()*, i8 }>
+@s3 = addrspace(1) global %t3 <{
+; CHECK32: .global .align 1 .u8 s3[5] = {
+; CHECK64: .global .align 1 .u8 s3[9] = {
+    void ()* @func,
+; CHECK-SAME:   0xFF(func), 0xFF00(func), 0xFF0000(func), 0xFF000000(func),
+; CHECK64-SAME: 0xFF00000000(func), 0xFF0000000000(func), 0xFF000000000000(func), 0xFF00000000000000(func),
+    i8 56 }>, align 1
+; CHECK-SAME:   56};
+
+;; Test that a packed struct with aligned pointers is printed in words.
+
+%t4 = type <{ void ()*, i64 }>
+@s4 = addrspace(1) global %t4 <{
+; CHECK32: .global .align 1 .u32 s4[3] = {
+; CHECK64: .global .align 1 .u64 s4[2] = {
+    void()* @func,
+; CHECK-SAME:   func,
+    i64 15}>, align 1
+; CHECK32-SAME: 15, 0};
+; CHECK64-SAME: 15};
+
+;; Test that a packed struct with unaligned pointers inside an array is handled.
+
+%t5 = type <{ void ()*, i16 }>
+@a5 = addrspace(1) global [2 x %t5] [%t5 <{ void()* @func, i16 5 }>, %t5 <{ void()* @func, i16 9 }> ]
+; CHECK32: .global .align 8 .u8 a5[12] = {
+; CHECK32-SAME: 0xFF(func), 0xFF00(func), 0xFF0000(func), 0xFF000000(func), 5, 0,
+; CHECK32-SAME: 0xFF(func), 0xFF00(func), 0xFF0000(func), 0xFF000000(func), 9, 0};
+; CHECK64: .global .align 8 .u8 a5[20] = {
+; CHECK64-SAME: 0xFF(func), 0xFF00(func), 0xFF0000(func), 0xFF000000(func),
+; CHECK64-SAME: 0xFF00000000(func), 0xFF0000000000(func), 0xFF000000000000(func), 0xFF00000000000000(func),
+; CHECK64-SAME: 5, 0,
+; CHECK64-SAME: 0xFF(func), 0xFF00(func), 0xFF0000(func), 0xFF000000(func),
+; CHECK64-SAME: 0xFF00000000(func), 0xFF0000000000(func), 0xFF000000000000(func), 0xFF00000000000000(func),
+; CHECK64-SAME: 9, 0};