[NVPTX] Promote i24, i40, i48 and i56 to next power-of-two register when passing

[kjetilkjeka]

Summary

Today llc will crash when attempting to use non-power-of-two integer types as function arguments or returns. This patch enables passing non standard integer values in functions by promoting them before store and truncating after load.

The main motivation of implementing this change is that rust casts small structs (less than pointer size) into an integer of the same size. As an example, if a struct contains three u8 then it will be passed as an i24. This patch is a step towards enabling rust compilation to ptx while retaining the target independent optimizations.

The tests reflects that it is mostly the multiple of 8 integers less than 64 that is of interest. I have locally tested some non-multiple-of-eight integers but decided against writing tests for them as nothing should really rely on them. Let me know if you want a few of those in addition.

More context can be found in my original github issue

---

This is my first LLVM contribution and I hope I have done everything by your contribution guide. Let me know if I should fix anything and I will do my best to get it done. I'm also not very familiar with the LLVM codebase yet and I have not received any external input on the content of this patch beyond the test passing. If anything looks fishy it is probably something I have misunderstood.

[kjetilkjeka]

When I run the failing test on my machine with <build>/tools/clang/tools/extra/clangd/unittests/./ClangdTests --gtest_filter=TUSchedulerTests.PreambleThrottle it passes even for the exact same pre-merge revision as the CI uses. Is this some sort of sporadic failure?

(Also, if I understand llvm correctly, the result of this test should not be changed by the nvptx changes I made)
`

[tra]

General nit: patch should be submitted on phabricator using large context. Please see: https://llvm.org/docs/Phabricator.html#requesting-a-review-via-the-web-interface

All i*-param.ll tests should probably be combined into one test file. They do pretty much the same thing.

Also, there should probably be test cases for integer sizes that are not multiples of 8. E.g i49.

llvm/lib/Target/NVPTX/NVPTXAsmPrinter.cpp

1500–1503

Size roundup calculation appears to be repeated in many places and could be extracted into a helper function.

llvm/lib/Target/NVPTX/NVPTXISelLowering.cpp

216

This looks a bit odd -- checking for the same boundary in two places, implicitly skipping power-of-2 sized integers...

I'd restructure it a bit to look like this:

if (VT.isScalarInteger()) {
   auto n = VT.getFixedSizeInBits();
   if (isPowerOf2_32(n))
     return false;
   switch(PowerOf2Ceil(n)){
      default: return false; // Covers i1 and integers larger than i64
      case 2:
      case 4:
      case 8:       *PromotedVT = MVT::i8; break;
      case 16:     *PromotedVT = MVT::i16; break;
      case 32:     *PromotedVT = MVT::i32; break;
      case 64:     *PromotedVT = MVT::i64; break;
    }
    return true
}

[kjetilkjeka]

Thanks for the review!

I have tried to address all the requested changes and created the new diff with the large context.

I added a couple of non standard tests, but LLVM seem to handle the promote/truncate in many different ways for the different non standard integers so I'm not sure at what level of detail it makes sense to test for here. Let me know if you want more tests.

[tra]

I have tried to address all the requested changes and created the new diff with the large context.

Thank you. It makes patch reviewing much easier.

I added a couple of non standard tests, but LLVM seem to handle the promote/truncate in many different ways for the different non standard integers so I'm not sure at what level of detail it makes sense to test for here. Let me know if you want more tests.

We want to make sure that param loads/stores are done using correct sizes. We do not really care how LLVM does extension/truncation of non-power-of-2 sized integers.
So, in practice we only need to check for relevant ld.param/st.param instructions.
We don't need to enumerate all possible sizes. Something like i3/i11/i23/i47 should be sufficient to exercise the code in this patch.
We do have test/CodeGen/NVPTX/param-load-store.ll so the new test cases should probably go there as well.

BTW, for IR test cases here it's convenient to get the function to call itself, so you do not need both caller and callee and can test all relevant functionality in one place.

llvm/lib/Target/NVPTX/NVPTXISelLowering.cpp

217

I think we may simplify it a bit further. Get rid of this check, always set PromotedVT to the appropriate size and return true if the type changed.

switch(PowerOf2Ceil(n)){
      default: llvm_unreachable();
      case 1:      *PromotedVT = MVT::i1; break;
      case 2:
      case 4:
      case 8:       *PromotedVT = MVT::i8; break;
      case 16:     *PromotedVT = MVT::i16; break;
      case 32:     *PromotedVT = MVT::i32; break;
      case 64:     *PromotedVT = MVT::i64; break;
    }
    return EVT(*PromotedVT) != VT;

1327–1330

Use promoteScalarArgumentSize() here, too?

1381–1384

ditto.

1555–1562

This could also use promoteScalarArgumentSize().

if ((VT.isInteger() || VT.isFloatingPoint())
  TypeSize = promoteScalarArgumentSize(TypeSize*8)/8.

1690–1693

ditto.

[kjetilkjeka]

Thanks for being so responsive and helpful! It helps a lot being my first time attempting to contribute to LLVM.

I added two small questions to your comments as inline comments.

We want to make sure that param loads/stores are done using correct sizes. We do not really care how LLVM does extension/truncation of non-power-of-2 sized integers.
So, in practice we only need to check for relevant ld.param/st.param instructions.
We don't need to enumerate all possible sizes. Something like i3/i11/i23/i47 should be sufficient to exercise the code in this patch.

Should I then also remove the attempts of testing the truncating/extending to make all tests the same?

llvm/lib/Target/NVPTX/NVPTXISelLowering.cpp

217

I will fix this together with the change from the other comment when I have gotten your opinion on where to put the helper from the other comment but first just a question about the case default.

Isn't it better to keep the case default: return false; to allow for attempting to promote the EVT and doing somewhere else if there is not a proper way to do it? I'm thinking more of the future scenario where things like i99 and i1000 is supported and you might first try to promote the integer type and if it fail then instead try to split it up into parts. I guess it doesn't make a big difference for this patch because things like i99 or i1000 will fail to compile both before and after the patch anyway. The only practical difference is the type of error you get if you try to feed these kind of integers to llc.

1327–1330

We would then need to put promoteScalarArgumentSize in a place that would be visible from both files (NVPTXSelLowering.cpp and NVPTXAsmPrinter.cpp). Would it make sense to use NVPTXUtilities.cpp for that?

[tra]

In D129291#3647790, @kjetilkjeka wrote:

Should I then also remove the attempts of testing the truncating/extending to make all tests the same?

If you want to. Or it could be cleaned up in a separate patch. Leaving existing tests as is is fine, too.

llvm/lib/Target/NVPTX/NVPTXISelLowering.cpp

217

Isn't it better to keep the case default: return false; to allow for attempting to promote the EVT and doing somewhere else if there is not a proper way to do it?

If there's a valid use case for it, then sure. On the other hand, if we are only expected to handle specific set of types (i1-i64), then I'd prefer to know right away if we ever end up with something that we can't handle. It's better to fail right away rather than ignore it and then have to debug it the hard way when things go wrong somewhere downstream. If we do grow a legitimate need to pass through other types unpromoted it will be easy enough to change.

1327–1330

I think you could just put it as an inline function into NVPTXUtilities.h

[kjetilkjeka]

I have implement your review comments.

Looking at the tests from test/CodeGen/NVPTX/param-load-store.ll helped a lot and I think the quality of the tests are better now.

I also discovered a bug when implementing the i2 and i3 test related to promoting the type to i8 but using a i16 as the container but it should be fixed now.

[tra]

llvm/test/CodeGen/NVPTX/param-load-store.ll

559–562 ↗

(On Diff #444662)

This, and other tests dealing with multiple loads/stores of parts of the argument may be fragile. First, the order of loads would not necessarily be guaranteed. Second -- the way LLVM reconstructs the value from parts may also change.

I would stick with a set of CHECK-DAG: ld.param.XX {{.*}} [test_i23_param_0+Y]. All we care about here is that we load the right set of bits. We can assume that reconstructing the integer is handled by appropriate tests already.

[kjetilkjeka]

I would stick with a set of CHECK-DAG: ld.param.XX {{.*}} [test_i23_param_0+Y]. All we care about here is that we load the right set of bits. We can assume that reconstructing the integer is handled by appropriate tests already.

I realize that you say that all we care about is loading the right set of bits. But I assume we would like to check it for returns/args on both the caller and callee side? These four cases is basically what I'm doing now in addition to checking that the function is actually called. I have removed the checking of truncating/promoting as you mentioned.

More specifically I'm checking the reads from the parameter in the call is done for the correct amount of bytes with CHECK-DAG. While the pass of the return is done with the promoted type. I'm also checking the pass and return from the call is being done with the correctly promoted type. I'm not checking anything related to truncating or promoting the integers before or after the load/store to param space.

Let me know if you want even less checks in the tests, like in fact only checking the read is being done with the correct amount of bytes, and I will fix it.

[tra]

LGTM with one test nit.

In D129291#3654804, @kjetilkjeka wrote:

I would stick with a set of CHECK-DAG: ld.param.XX {{.*}} [test_i23_param_0+Y]. All we care about here is that we load the right set of bits. We can assume that reconstructing the integer is handled by appropriate tests already.

I realize that you say that all we care about is loading the right set of bits. But I assume we would like to check it for returns/args on both the caller and callee side? These four cases is basically what I'm doing now in addition to checking that the function is actually called. I have removed the checking of truncating/promoting as you mentioned.

Let me rephrase -- we only care about loads/stores from the param space. That does apply to both load of parameters from the function arguments and storing them to pass parameters and return values.
I think we are in agreement on what we need to check. The comment was largely about checks for cvt, shl and or instructions we match (e.g a few still remain in test_i24).

llvm/test/CodeGen/NVPTX/param-load-store.ll

567 ↗

(On Diff #444947)

we want CHECK-DAG on loads here.

Nit: no need to check logical ops.

[kjetilkjeka]

Seems like I forgot changing the i24 tests. It should be fixed now.

I still do not know why these CI tests are failing? I don't think it should have anything to do with this change though? I assume it's OK since you have not commented on it.

I hope that you will apply the patch to the code-base since I obviously do not have write access.

Thanks for helping and reviewing!

[tra]

Sanitizer test failures are AFAICT unrelated.

I will land your patch a bit later this week. Thank you for contributing these changes.

llvm/lib/Target/NVPTX/NVPTXAsmPrinter.cpp

       // PTX, so its size must be adjusted here, too.
       if (size < 32)
         size = 32;
+      if (size > 32 && size < 64)
+        size = 64;
 
       O << ".param .b" << size << " func_retval0";
     } else if (isa<PointerType>(Ty)) {

 
       for (unsigned j = 0, je = elems; j != je; ++j) {
         unsigned sz = elemtype.getSizeInBits();
-        if (elemtype.isInteger() && (sz < 32))
-          sz = 32;
+        if (elemtype.isInteger()) {
+          if (sz <= 32)
+            sz = 32;
+          else if (sz <= 64)
+            sz = 64;
+        }
         O << ".reg .b" << sz << " func_retval" << idx;
         if (j < je - 1)
           O << ", ";

       unsigned sz = 0;
       if (isa<IntegerType>(Ty)) {
         sz = cast<IntegerType>(Ty)->getBitWidth();
-        if (sz < 32)
+        if (sz <= 32)
           sz = 32;
+        else if (sz <= 64)
+          sz = 64;
       } else if (isa<PointerType>(Ty))
         sz = thePointerTy.getSizeInBits();
       else if (Ty->isHalfTy())

 
         for (unsigned j = 0, je = elems; j != je; ++j) {
           unsigned sz = elemtype.getSizeInBits();
-          if (elemtype.isInteger() && (sz < 32))
-            sz = 32;
+          if (elemtype.isInteger()) {
+            if (sz <= 32)
+              sz = 32;
+            else if (sz <= 64)
+              sz = 64;
+          }
           O << "\t.reg .b" << sz << " ";
           printParamName(I, paramIndex, O);
           if (j < je - 1)

llvm/lib/Target/NVPTX/NVPTXISelLowering.cpp

   }
 }
 
+/// PromoteScalarIntegerPTX
+/// Used to make sure the arguments/returns are suitable for passing
+/// and promote them to a larger size if they're not.
+///
+/// The promoted type is placed in \p PromoteVT if the function returns true.
+static bool PromoteScalarIntegerPTX(const EVT &VT, MVT *PromotedVT) {
+  if (VT.isScalarInteger()) {
+    if (1 < VT.getFixedSizeInBits() && VT.getFixedSizeInBits() < 8) {
+      *PromotedVT = MVT::i8;
+      return true;
+    } else if (8 < VT.getFixedSizeInBits() && VT.getFixedSizeInBits() < 16) {
+      *PromotedVT = MVT::i16;
+      return true;
+    } else if (16 < VT.getFixedSizeInBits() && VT.getFixedSizeInBits() < 32) {
+      *PromotedVT = MVT::i32;
+      return true;
+    } else if (32 < VT.getFixedSizeInBits() && VT.getFixedSizeInBits() < 64) {
+      *PromotedVT = MVT::i64;
+      return true;
+    }
+  }
+  return false;
+}
+
 // Check whether we can merge loads/stores of some of the pieces of a
 // flattened function parameter or return value into a single vector
 // load/store.

       // PTX, so its size must be adjusted here, too.
       if (size < 32)
         size = 32;
+      else if (size > 32 && size < 64)
+        size = 64;
 
       O << ".param .b" << size << " _";
     } else if (isa<PointerType>(retTy)) {

         sz = cast<IntegerType>(Ty)->getBitWidth();
         if (sz < 32)
           sz = 32;
+        else if (sz > 32 && sz < 64)
+          sz = 64;
       } else if (isa<PointerType>(Ty)) {
         sz = PtrVT.getSizeInBits();
       } else if (Ty->isHalfTy())

         // size. FP16 is loaded/stored using i16, so it's handled
         // here as well.
         TypeSize = 4;
+      } else if (VT.isInteger() && TypeSize > 4 && TypeSize < 8) {
+        TypeSize = 8;
       }
       SDValue DeclareScalarParamOps[] = {
           Chain, DAG.getConstant(ParamCount, dl, MVT::i32),

       }
 
       SDValue StVal = OutVals[OIdx];
+
+      MVT PromotedVT;
+      if (PromoteScalarIntegerPTX(EltVT, &PromotedVT)) {
+        llvm::ISD::NodeType Ext =
+            Outs[OIdx].Flags.isSExt() ? ISD::SIGN_EXTEND : ISD::ZERO_EXTEND;
+        StVal = DAG.getNode(Ext, dl, PromotedVT, StVal);
+        EltVT = EVT(PromotedVT);
+      }
+
       if (IsByVal) {
         auto PtrVT = getPointerTy(DL);
         SDValue srcAddr = DAG.getNode(ISD::ADD, dl, PtrVT, StVal,

       // Scalar needs to be at least 32bit wide
       if (resultsz < 32)
         resultsz = 32;
+      else if (resultsz > 32 && resultsz < 64)
+        resultsz = 64;
       SDVTList DeclareRetVTs = DAG.getVTList(MVT::Other, MVT::Glue);
       SDValue DeclareRetOps[] = { Chain, DAG.getConstant(1, dl, MVT::i32),
                                   DAG.getConstant(resultsz, dl, MVT::i32),

       EVT TheLoadType = VTs[i];
       EVT EltType = Ins[i].VT;
       Align EltAlign = commonAlignment(RetAlign, Offsets[i]);
+      MVT PromotedVT;
+
+      if (PromoteScalarIntegerPTX(TheLoadType, &PromotedVT)) {
+        TheLoadType = EVT(PromotedVT);
+        EltType = EVT(PromotedVT);
+        needTruncate = true;
+      }
+
       if (ExtendIntegerRetVal) {
         TheLoadType = MVT::i32;
         EltType = MVT::i32;

             // v2f16 was loaded as an i32. Now we must bitcast it back.
             else if (EltVT == MVT::v2f16)
               Elt = DAG.getNode(ISD::BITCAST, dl, MVT::v2f16, Elt);
+
+            // If a promoted integer type is used, truncate down to the original
+            MVT PromotedVT;
+            if (PromoteScalarIntegerPTX(EltVT, &PromotedVT)) {
+              Elt = DAG.getNode(ISD::TRUNCATE, dl, EltVT, Elt);
+            }
+
             // Extend the element if necessary (e.g. an i8 is loaded
             // into an i16 register)
             if (Ins[InsIdx].VT.isInteger() &&

     return Chain;
 
   const DataLayout &DL = DAG.getDataLayout();
+  SmallVector<SDValue, 16> PromotedOutVals;
   SmallVector<EVT, 16> VTs;
   SmallVector<uint64_t, 16> Offsets;
   ComputePTXValueVTs(*this, DL, RetTy, VTs, &Offsets);
   assert(VTs.size() == OutVals.size() && "Bad return value decomposition");
 
+  for (unsigned i = 0, e = VTs.size(); i != e; ++i) {
+    SDValue PromotedOutVal = OutVals[i];
+    MVT PromotedVT;
+    if (PromoteScalarIntegerPTX(VTs[i], &PromotedVT)) {
+      llvm::ISD::NodeType Ext =
+          Outs[i].Flags.isSExt() ? ISD::SIGN_EXTEND : ISD::ZERO_EXTEND;
+      PromotedOutVal = DAG.getNode(Ext, dl, PromotedVT, PromotedOutVal);
+      VTs[i] = EVT(PromotedVT);
+    }
+    PromotedOutVals.push_back(PromotedOutVal);
+  }
+
   auto VectorInfo = VectorizePTXValueVTs(
       VTs, Offsets,
       RetTy->isSized() ? getFunctionParamOptimizedAlign(&F, RetTy, DL)

       StoreOperands.push_back(DAG.getConstant(Offsets[i], dl, MVT::i32));
     }
 
-    SDValue RetVal = OutVals[i];
+    SDValue OutVal = OutVals[i];
+    SDValue RetVal = PromotedOutVals[i];
+
     if (ExtendIntegerRetVal) {
       RetVal = DAG.getNode(Outs[i].Flags.isSExt() ? ISD::SIGN_EXTEND
                                                   : ISD::ZERO_EXTEND,
                            dl, MVT::i32, RetVal);
-    } else if (RetVal.getValueSizeInBits() < 16) {
+    } else if (OutVal.getValueSizeInBits() < 16) {
       // Use 16-bit registers for small load-stores as it's the
       // smallest general purpose register size supported by NVPTX.
       RetVal = DAG.getNode(ISD::ANY_EXTEND, dl, MVT::i16, RetVal);

llvm/test/CodeGen/NVPTX/i24-param.ll

+; RUN: llc < %s -march=nvptx -mcpu=sm_20 -verify-machineinstrs | FileCheck %s
+; RUN: %if ptxas %{ llc < %s -march=nvptx -mcpu=sm_20 -verify-machineinstrs | %ptxas-verify %}
+
+target datalayout = "e-p:32:32:32-i1:8:8-i8:8:8-i16:16:16-i32:32:32-i64:64:64-f32:32:32-f64:64:64-v16:16:16-v32:32:32-v64:64:64-v128:128:128-n16:32:64"
+
+; CHECK: .visible .func  (.param .b32 func_retval0) callee
+; CHECK: .param .b32 callee_param_0
+define i24 @callee(i24 %a) {
+  %val = alloca i24, align 4
+  store i24 %a, i24* %val, align 4
+  %ret = load i24, i24* %val, align 1
+; CHECK: ld.param.u8
+; CHECK: ld.param.u16
+; CHECK: st.param.b32
+  ret i24 %ret
+}
+
+; CHECK: .visible .func caller
+define void @caller(i24* %a) {
+  %val = load i24, i24* %a
+  %ret = call i24 @callee(i24 %val)
+; CHECK: ld.param.b32
+; CHECK: st.u16
+; CHECK: st.u8
+  store i24 %ret, i24* %a
+  ret void
+}

llvm/test/CodeGen/NVPTX/i40-param.ll

+; RUN: llc < %s -march=nvptx -mcpu=sm_20 -verify-machineinstrs | FileCheck %s
+; RUN: %if ptxas %{ llc < %s -march=nvptx -mcpu=sm_20 -verify-machineinstrs | %ptxas-verify %}
+
+target datalayout = "e-p:32:32:32-i1:8:8-i8:8:8-i16:16:16-i32:32:32-i64:64:64-f32:32:32-f64:64:64-v16:16:16-v32:32:32-v64:64:64-v128:128:128-n16:32:64"
+
+; CHECK: .visible .func  (.param .b64 func_retval0) callee
+; CHECK: .param .b64 callee_param_0
+define i40 @callee(i40 %a) {
+  %val = alloca i40, align 8
+  store i40 %a, i40* %val, align 8
+  %ret = load i40, i40* %val, align 1
+; CHECK: ld.param.u8
+; CHECK: ld.param.u32
+; CHECK: st.param.b64
+  ret i40 %ret
+}
+
+; CHECK: .visible .func caller
+define void @caller(i40* %a) {
+  %val = load i40, i40* %a
+  %ret = call i40 @callee(i40 %val)
+; CHECK: ld.param.b64
+; CHECK: st.u32
+; CHECK: st.u8
+  store i40 %ret, i40* %a
+  ret void
+}

llvm/test/CodeGen/NVPTX/i48-param.ll

+; RUN: llc < %s -march=nvptx -mcpu=sm_20 -verify-machineinstrs | FileCheck %s
+; RUN: %if ptxas %{ llc < %s -march=nvptx -mcpu=sm_20 -verify-machineinstrs | %ptxas-verify %}
+
+target datalayout = "e-p:32:32:32-i1:8:8-i8:8:8-i16:16:16-i32:32:32-i64:64:64-f32:32:32-f64:64:64-v16:16:16-v32:32:32-v64:64:64-v128:128:128-n16:32:64"
+
+; CHECK: .visible .func  (.param .b64 func_retval0) callee
+; CHECK: .param .b64 callee_param_0
+define i48 @callee(i48 %a) {
+  %val = alloca i48, align 8
+  store i48 %a, i48* %val, align 8
+  %ret = load i48, i48* %val, align 1
+; CHECK: ld.param.u16
+; CHECK: ld.param.u32
+; CHECK: st.param.b64
+  ret i48 %ret
+}
+
+; CHECK: .visible .func caller
+define void @caller(i48* %a) {
+  %val = load i48, i48* %a
+  %ret = call i48 @callee(i48 %val)
+; CHECK: ld.param.b64
+; CHECK: st.u32
+; CHECK: st.u16
+  store i48 %ret, i48* %a
+  ret void
+}

llvm/test/CodeGen/NVPTX/i56-param.ll

+; RUN: llc < %s -march=nvptx -mcpu=sm_20 -verify-machineinstrs | FileCheck %s
+; RUN: %if ptxas %{ llc < %s -march=nvptx -mcpu=sm_20 -verify-machineinstrs | %ptxas-verify %}
+
+target datalayout = "e-p:32:32:32-i1:8:8-i8:8:8-i16:16:16-i32:32:32-i64:64:64-f32:32:32-f64:64:64-v16:16:16-v32:32:32-v64:64:64-v128:128:128-n16:32:64"
+
+; CHECK: .visible .func  (.param .b64 func_retval0) callee
+; CHECK: .param .b64 callee_param_0
+define i56 @callee(i56 %a) {
+  %val = alloca i56, align 8
+  store i56 %a, i56* %val, align 8
+  %ret = load i56, i56* %val, align 1
+; CHECK: ld.param.u8
+; CHECK: ld.param.u16
+; CHECK: ld.param.u32
+; CHECK: st.param.b64
+  ret i56 %ret
+}
+
+; CHECK: .visible .func caller
+define void @caller(i56* %a) {
+  %val = load i56, i56* %a
+  %ret = call i56 @callee(i56 %val)
+; CHECK: ld.param.b64
+; CHECK: st.u32
+; CHECK: st.u8
+; CHECK: st.u16
+  store i56 %ret, i56* %a
+  ret void
+}