fixup and address comments

Author: AlexMaclean

llvm/lib/Target/NVPTX/NVPTXISelDAGToDAG.cpp

@@ -2183,25 +2183,29 @@ bool NVPTXDAGToDAGISel::tryStoreRetval(SDNode *N) {
 }
 
 // Helpers for constructing opcode (ex: NVPTX::StoreParamV4F32_iiri)
-#define getOpcV2H(ty, op0, op1) NVPTX::StoreParamV2##ty##_##op0##op1
+#define getOpcV2H(ty, opKind0, opKind1)                                        \
+  NVPTX::StoreParamV2##ty##_##opKind0##opKind1
 
-#define getOpcV2H1(ty, op0, op1)                                               \
-  (op1) ? getOpcV2H(ty, op0, i) : getOpcV2H(ty, op0, r)
+#define getOpcV2H1(ty, opKind0, isImm1)                                        \
+  (isImm1) ? getOpcV2H(ty, opKind0, i) : getOpcV2H(ty, opKind0, r)
 
 #define getOpcodeForVectorStParamV2(ty, isimm)                                 \
   (isimm[0]) ? getOpcV2H1(ty, i, isimm[1]) : getOpcV2H1(ty, r, isimm[1])
 
-#define getOpcV4H(ty, op0, op1, op2, op3)                                      \
-  NVPTX::StoreParamV4##ty##_##op0##op1##op2##op3
+#define getOpcV4H(ty, opKind0, opKind1, opKind2, opKind3)                      \
+  NVPTX::StoreParamV4##ty##_##opKind0##opKind1##opKind2##opKind3
 
-#define getOpcV4H3(ty, op0, op1, op2, op3)                                     \
-  (op3) ? getOpcV4H(ty, op0, op1, op2, i) : getOpcV4H(ty, op0, op1, op2, r)
+#define getOpcV4H3(ty, opKind0, opKind1, opKind2, isImm3)                      \
+  (isImm3) ? getOpcV4H(ty, opKind0, opKind1, opKind2, i)                       \
+           : getOpcV4H(ty, opKind0, opKind1, opKind2, r)
 
-#define getOpcV4H2(ty, op0, op1, op2, op3)                                     \
-  (op2) ? getOpcV4H3(ty, op0, op1, i, op3) : getOpcV4H3(ty, op0, op1, r, op3)
+#define getOpcV4H2(ty, opKind0, opKind1, isImm2, isImm3)                       \
+  (isImm2) ? getOpcV4H3(ty, opKind0, opKind1, i, isImm3)                       \
+           : getOpcV4H3(ty, opKind0, opKind1, r, isImm3)
 
-#define getOpcV4H1(ty, op0, op1, op2, op3)                                     \
-  (op1) ? getOpcV4H2(ty, op0, i, op2, op3) : getOpcV4H2(ty, op0, r, op2, op3)
+#define getOpcV4H1(ty, opKind0, isImm1, isImm2, isImm3)                        \
+  (isImm1) ? getOpcV4H2(ty, opKind0, i, isImm2, isImm3)                        \
+           : getOpcV4H2(ty, opKind0, r, isImm2, isImm3)
 
 #define getOpcodeForVectorStParamV4(ty, isimm)                                 \
   (isimm[0]) ? getOpcV4H1(ty, i, isimm[1], isimm[2], isimm[3])                 \
@@ -2211,10 +2215,10 @@ bool NVPTXDAGToDAGISel::tryStoreRetval(SDNode *N) {
   (n == 2) ? getOpcodeForVectorStParamV2(ty, isimm)                            \
            : getOpcodeForVectorStParamV4(ty, isimm)
 
-static std::optional<unsigned>
-pickOpcodeForVectorStParam(SmallVector<SDValue, 8> &Ops, unsigned NumElts,
-                           MVT::SimpleValueType MemTy, SelectionDAG *CurDAG,
-                           SDLoc DL) {
+static unsigned pickOpcodeForVectorStParam(SmallVector<SDValue, 8> &Ops,
+                                           unsigned NumElts,
+                                           MVT::SimpleValueType MemTy,
+                                           SelectionDAG *CurDAG, SDLoc DL) {
   // Determine which inputs are registers and immediates make new operators
   // with constant values
   SmallVector<bool, 4> IsImm(NumElts, false);
@@ -2244,19 +2248,31 @@ pickOpcodeForVectorStParam(SmallVector<SDValue, 8> &Ops, unsigned NumElts,
   case MVT::i32:
     return getOpcodeForVectorStParam(NumElts, I32, IsImm);
   case MVT::i64:
-    if (NumElts == 4)
-      return std::nullopt;
+    assert(NumElts == 2 && "MVT too large for NumElts > 2");
     return getOpcodeForVectorStParamV2(I64, IsImm);
   case MVT::f32:
     return getOpcodeForVectorStParam(NumElts, F32, IsImm);
   case MVT::f64:
-    if (NumElts == 4)
-      return std::nullopt;
+    assert(NumElts == 2 && "MVT too large for NumElts > 2");
     return getOpcodeForVectorStParamV2(F64, IsImm);
+
+  // These cases don't support immediates, just use the all register version
+  // and generate moves.
+  case MVT::i1:
+    return (NumElts == 2) ? NVPTX::StoreParamV2I8_rr
+                          : NVPTX::StoreParamV4I8_rrrr;
   case MVT::f16:
+  case MVT::bf16:
+    return (NumElts == 2) ? NVPTX::StoreParamV2I16_rr
+                          : NVPTX::StoreParamV4I16_rrrr;
   case MVT::v2f16:
+  case MVT::v2bf16:
+  case MVT::v2i16:
+  case MVT::v4i8:
+    return (NumElts == 2) ? NVPTX::StoreParamV2I32_rr
+                          : NVPTX::StoreParamV4I32_rrrr;
   default:
-    return std::nullopt;
+    llvm_unreachable("Cannot select st.param for unknown MemTy");
   }
 }
 
@@ -2271,10 +2287,10 @@ bool NVPTXDAGToDAGISel::tryStoreParam(SDNode *N) {
   SDValue Glue = N->getOperand(N->getNumOperands() - 1);
 
   // How many elements do we have?
-  unsigned NumElts = 1;
+  unsigned NumElts;
   switch (N->getOpcode()) {
   default:
-    return false;
+    llvm_unreachable("Unexpected opcode");
   case NVPTXISD::StoreParamU32:
   case NVPTXISD::StoreParamS32:
   case NVPTXISD::StoreParam:
@@ -2300,12 +2316,12 @@ bool NVPTXDAGToDAGISel::tryStoreParam(SDNode *N) {
   // Determine target opcode
   // If we have an i1, use an 8-bit store. The lowering code in
   // NVPTXISelLowering will have already emitted an upcast.
-  std::optional<unsigned> Opcode = 0;
+  std::optional<unsigned> Opcode;
   switch (N->getOpcode()) {
   default:
     switch (NumElts) {
     default:
-      return false;
+      llvm_unreachable("Unexpected NumElts");
     case 1: {
       MVT::SimpleValueType MemTy = Mem->getMemoryVT().getSimpleVT().SimpleTy;
       SDValue Imm = Ops[0];
@@ -2357,8 +2373,6 @@ bool NVPTXDAGToDAGISel::tryStoreParam(SDNode *N) {
       break;
     }
     }
-    if (!Opcode)
-      return false;
     break;
   // Special case: if we have a sign-extend/zero-extend node, insert the
   // conversion instruction first, and use that as the value operand to