[llvm] [NVPTX] Use v2.u64 to load/store 128-bit values (PR #136638)

[Artem Belevich via llvm-commits]

================
@@ -3201,28 +3197,31 @@ NVPTXTargetLowering::LowerSTOREVector(SDValue Op, SelectionDAG &DAG) const {
   Ops.push_back(N->getOperand(0));
 
   // Then the split values
-  assert(NumElts <= ValVT.getVectorNumElements() &&
-         "NumElts should not increase, only decrease or stay the same.");
-  if (NumElts < ValVT.getVectorNumElements()) {
-    // If the number of elements has decreased, getVectorLoweringShape has
-    // upsized the element types
-    assert(EltVT.isVector() && EltVT.getSizeInBits() == 32 &&
-           EltVT.getVectorNumElements() <= 4 && "Unexpected upsized type.");
+  if (EltVT.isVector()) {
+    assert(EVT(EltVT.getVectorElementType()) == ValVT.getVectorElementType());
+    assert(NumElts * EltVT.getVectorNumElements() ==
+           ValVT.getVectorNumElements());
     // Combine individual elements into v2[i,f,bf]16/v4i8 subvectors to be
     // stored as b32s
-    unsigned NumEltsPerSubVector = EltVT.getVectorNumElements();
-    for (unsigned i = 0; i < NumElts; ++i) {
+    const unsigned NumEltsPerSubVector = EltVT.getVectorNumElements();
+    for (const auto I : llvm::seq(NumElts)) {
       SmallVector<SDValue, 4> SubVectorElts;
-      DAG.ExtractVectorElements(Val, SubVectorElts, i * NumEltsPerSubVector,
+      DAG.ExtractVectorElements(Val, SubVectorElts, I * NumEltsPerSubVector,
                                 NumEltsPerSubVector);
       SDValue SubVector = DAG.getBuildVector(EltVT, DL, SubVectorElts);
       Ops.push_back(SubVector);
----------------
Artem-B wrote:

Nit: We could just push_back getBuildVector result directly.

https://github.com/llvm/llvm-project/pull/136638