CUDA: unions in kernel arguments not copied completely if member contains padding

[mkuron]

Summary

I have a CUDA kernel that takes a union as one of its parameters. One of the union's member types is a struct (called Ptrs in my minimal working example) that contains some padding due to the memory alignment requirements of its own members. I initialize the union as its other member type (Data in the MWE) and pass it to the kernel launch. Inside the kernel however, some of the contents of Data appear to be zeroed out -- notably those corresponding to the padding region of Ptrs. It seems like Clang looks at the padding inside one of the union's member types and assumes that all the members have holes in the same places.

When compiling the CUDA kernel, Clang should make no assumption about which one of the member types is currently stored in a union that it gets as an argument, but it clearly does and it's the wrong assumption in this case. We initially found this bug in conjunction with thrust and mpark::variant, but it can be observed with plain CUDA code too as in the MWE.

Versions

I reproduced this bug with multiple Clang versions between 7.0.0 and 14.0.0-rc1, as well as the current master branch. Nvidia's nvcc compiler does not exhibit such an issue.

Starting with d0615a9, my MWE required a minor modification to make sure that the argument copy wasn't optimized out. This optimization is not generally possible (e.g. when using thrust), so we cannot rely on it accidentally fixing the real issue.

Minimal working example

Compile with clang++ -O3 -std=c++14 --cuda-gpu-arch=sm_70 mwe.cu -o mwe.cu -lcudart. The issue can also be observed without -O3. Add -DFIX_IT to get a struct without padding, in which case the issue does not occur.

#include <algorithm>
#include <iostream>
#include <vector>

#include <cuda.h>

struct Ptrs
{
  char s1, s2, s3, s4, s5, s6, s7;
#ifdef FIX_IT
  char s8;
#endif
  void const * data;
};

template <class T, unsigned N>
struct Data
{
  T data[N];
};

template <class T, unsigned N>
union DataType
{
  Data<T,N> d;
  Ptrs p;
};

__global__ void transform(DataType<unsigned, 4> umap, unsigned n, unsigned * result)
{
  const unsigned i = threadIdx.x;

  if(n > 0xffff) {
    /* This condition is never true, but the next line ensures that `umap` is copied into writable memory.
     * Clang 13 (https://github.com/llvm/llvm-project/commit/d0615a93bb6d7aedc43323dc8957fe57e86ed8ae)
     * introduced an optimization that will otherwise mask the issue we are trying to demonstrate.
     */
    umap.d.data[i] = 0;
  }

  if(i < n)
  {
    result[i] = umap.d.data[i];
  }
}

bool test()
{
  std::vector<unsigned> map{23U, 0xffffffff, 42U, 13U};
  DataType<unsigned, 4> umap;
  Data<unsigned, 4> d;
  std::copy(map.begin(), map.end(), &d.data[0]);
  umap.d = d;

  const unsigned mod = map.size();

  std::vector<unsigned> h_values(16);

  unsigned * values = nullptr;
  cudaMalloc(&values, h_values.size() * sizeof(unsigned));
  transform<<<1, h_values.size()>>>(umap, h_values.size(), values);
  cudaMemcpy(h_values.data(), values, h_values.size() * sizeof(unsigned), cudaMemcpyDeviceToHost);
  cudaFree(values);

  bool good = true;
  for(int i = 0; i < mod; ++i)
  {
    std::cout << h_values[i];
    if (h_values[i] != map[i])
    {
      std::cout << " wrong";
      good = false;
    }
    std::cout << std::endl;
  }

  return good;
}

int main()
{
    bool good = test();
    return (good ? 0 : 1);
}

Assembly

To see what is happening, compile the MWE with and without -DFIX_IT and compare the NVPTX assembly: https://godbolt.org/z/W7PY9q3sj . You can see that in the former case, it generates ld.param and st.local instructions that cover the entire size of the union, while in the latter case it skips those bytes that are padding.

[Artem-B]

Interesting. Generated IR does not have Data in the DataType: : https://godbolt.org/z/YhbPbecjx

%union.DataType = type { %struct.Ptrs }
%struct.Ptrs = type { i8, i8, i8, i8, i8, i8, i8, i8* }

When it's time to lower the byval arguments, LLVM only sees Ptrs and that's what it ends up copying. It has no idea that the parameter is actually a union. FIX_IT just fills the gap in the union, but does not really fix the root cause.

This is indeed problematic. There's a fundamental disconnect here -- union-ness is handled by the front-end, but when LLVM needs to generate IR to copy the argument deep down in LLVM, we have no idea that there's any other data out there other than the single union field picked by the front-end.

I guess one thing we could try is to convince the front-end to generate union types in LLVM in a way that the type covers the whole union storage.

[Artem-B]

Actually, we don't need the front-end changes. memcpy of the the byval parameter instead of load/store should work, as the front-end-generated type should always have the right size and that's all we care about.

[mkuron]

Thanks for your quick reply, @Artem-B.

when LLVM needs to generate IR to copy the argument deep down in LLVM, we have no idea that there's any other data out there other than the single union field picked by the front-end.

That seems dangerous in general, not only in a CUDA context. If information about one (the largest?) of the union's members is available to a backend, it might make optimizations based on that. Which may be wrong if a union contains one of its other member types, like in the present example. Can we be sure that neither other parts of the NVPTX backend nor other backends ever look inside a %union type and use any of its properties beyond the overall size -- like member size, member alignment, etc.?

I guess one thing we could try is to convince the front-end to generate union types in LLVM in a way that the type covers the whole union storage.

Something like %union.DataType = type { [16 x i8] }?

[Artem-B]

Can we be sure that neither other parts of the NVPTX backend nor other backends ever look inside a %union type and use any of its properties beyond the overall size -- like member size, member alignment, etc.?

I believe the answer is yes. Front-end is expected to generate semantically valid IR which for unions would include appropriate bitcasts, data transfers, etc. If that would be broken, I would expect to see fairly widespread failures in regular C++ code.

This particular issue is a special case where NVPTX back-end has to add additional code w/o the information available to the front-end. If that copy were to be generated by clang, it would copy all the data. E.g. https://godbolt.org/z/aa3r7nqvj -- regardless of whether we copy the union as one field or the whole union, clang always generates a memcpy of the whole union. If it needs to pass a union initialized with a field that has a padding gap, it still copies over the whole union: https://godbolt.org/z/jxbxsqMP6

Doing the same during lowering of byval arguments appears to be the right fix for this issue.

[mkuron]

Thanks for your detailed explanation, @Artem-B. Did you get a chance yet to look at how to fix this? I was actually able to put together a simple patch (see below) that resolves the issue for my minimal example code, but leads to CUDA illegal memory access errors when I try it on my actual code which involves Thrust and is generally more complex. So I will need to rely on you to get a complete patch plus an LLVM unit test.

diff --git a/llvm/lib/Target/NVPTX/NVPTXLowerArgs.cpp b/llvm/lib/Target/NVPTX/NVPTXLowerArgs.cpp
index 6183019de43d..46d6109b6bb0 100644
--- a/llvm/lib/Target/NVPTX/NVPTXLowerArgs.cpp
+++ b/llvm/lib/Target/NVPTX/NVPTXLowerArgs.cpp
@@ -92,12 +92,14 @@
 #include "NVPTXTargetMachine.h"
 #include "NVPTXUtilities.h"
 #include "MCTargetDesc/NVPTXBaseInfo.h"
+#include "llvm/Analysis/TargetTransformInfo.h"
 #include "llvm/Analysis/ValueTracking.h"
 #include "llvm/IR/Function.h"
 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/Module.h"
 #include "llvm/IR/Type.h"
 #include "llvm/Pass.h"
+#include "llvm/Transforms/Utils/LowerMemIntrinsics.h"

 #define DEBUG_TYPE "nvptx-lower-args"

@@ -130,6 +132,10 @@ public:
     return "Lower pointer arguments of CUDA kernels";
   }

+  void getAnalysisUsage(AnalysisUsage &AU) const override {
+    AU.addRequired<TargetTransformInfoWrapperPass>();
+  }
+
 private:
   const NVPTXTargetMachine *TM;
 };
@@ -290,10 +296,14 @@ void NVPTXLowerArgs::handleByValParam(Argument *Arg) {
   // Be sure to propagate alignment to this load; LLVM doesn't know that NVPTX
   // addrspacecast preserves alignment.  Since params are constant, this load is
   // definitely not volatile.
-  LoadInst *LI =
-      new LoadInst(StructType, ArgInParam, Arg->getName(),
-                   /*isVolatile=*/false, AllocA->getAlign(), FirstInst);
-  new StoreInst(LI, AllocA, FirstInst);
+  const uint64_t Size = DL.getTypeSizeInBits(Arg->getType());
+  LLVMContext &Ctx = Func->getParent()->getContext();
+  const TargetTransformInfo &TTI =
+      getAnalysis<TargetTransformInfoWrapperPass>().getTTI(*Func);
+  createMemCpyLoopKnownSize(FirstInst, ArgInParam, AllocA,
+                            ConstantInt::get(Type::getInt64Ty(Ctx), Size),
+                            AllocA->getAlign(), AllocA->getAlign(),
+                            false, false, TTI);
 }

 void NVPTXLowerArgs::markPointerAsGlobal(Value *Ptr) {

[Artem-B]

I was planning to incorporate this into a work-in-progress patch along with other improvements in NVPTXLowerArgs.cpp but it's got pushed back by other work.

I think the crash is due to attempt to copy 8x the amount we want to copy as Size we've got is in bits.

const uint64_t Size = DL.getTypeSizeInBits(Arg->getType());

You should verify that the memcpy produces ld.param instructions and not the regular ld.

Also, memcpy will be sub-optimal performance-wise as it would not be able to use 128-bit loads/stores. In the end we may want to use load/store on an opaque array type of the size that we need to copy.

[Artem-B]

Looks like argument lowering is only part of the problem. We also need to address the same issue with function call lowering so the arguments we pass and return are also copied completely

E.g: https://godbolt.org/z/d111d6asP. You can see that we copy only struct fields, but not the gaps between them.

[mkuron]

I was planning to incorporate this into a work-in-progress patch along with other improvements in NVPTXLowerArgs.cpp

@Artem-B, what's the status of this work-in-progress patch? We're still running into this bug, so I was wondering whether there might have been some progress in the past year.

[Artem-B]

Unfortunately, none. It got preempted by other work and then slipped through the cracks. :-(

I probably would not have much time to deal with it till August.

[mkuron]

@Artem-B, no problem. August or even September aligns well with our own plans, so I'm looking forward to any solution you might come up with.

[mkuron]

@Artem-B, any update? I suspect this bug might become more relevant to other users besides myself, now that libcu++ added cuda::variant in NVIDIA/cccl#1076.