[llvm] [LLVM][NVPTX] Add support for tcgen05.ld.red Instruction (PR #175919)

[Pradeep Kumar via llvm-commits]

================
@@ -3005,6 +3007,122 @@ static SDValue lowerPrmtIntrinsic(SDValue Op, SelectionDAG &DAG) {
   return getPRMT(A, B, Selector, DL, DAG, Mode);
 }
 
+static unsigned getTcgen05LdRedID(Intrinsic::ID IID) {
+  switch (IID) {
+  case Intrinsic::nvvm_tcgen05_ld_red_16x32bx2_x128_f32:
+    return NVPTXISD::TCGEN05_LD_RED_16x32bx2_X128_F32;
+  case Intrinsic::nvvm_tcgen05_ld_red_16x32bx2_x16_f32:
+    return NVPTXISD::TCGEN05_LD_RED_16x32bx2_X16_F32;
+  case Intrinsic::nvvm_tcgen05_ld_red_16x32bx2_x2_f32:
+    return NVPTXISD::TCGEN05_LD_RED_16x32bx2_X2_F32;
+  case Intrinsic::nvvm_tcgen05_ld_red_16x32bx2_x32_f32:
+    return NVPTXISD::TCGEN05_LD_RED_16x32bx2_X32_F32;
+  case Intrinsic::nvvm_tcgen05_ld_red_16x32bx2_x4_f32:
+    return NVPTXISD::TCGEN05_LD_RED_16x32bx2_X4_F32;
+  case Intrinsic::nvvm_tcgen05_ld_red_16x32bx2_x64_f32:
+    return NVPTXISD::TCGEN05_LD_RED_16x32bx2_X64_F32;
+  case Intrinsic::nvvm_tcgen05_ld_red_16x32bx2_x8_f32:
+    return NVPTXISD::TCGEN05_LD_RED_16x32bx2_X8_F32;
+  case Intrinsic::nvvm_tcgen05_ld_red_32x32b_x128_f32:
+    return NVPTXISD::TCGEN05_LD_RED_32x32b_X128_F32;
+  case Intrinsic::nvvm_tcgen05_ld_red_32x32b_x16_f32:
+    return NVPTXISD::TCGEN05_LD_RED_32x32b_X16_F32;
+  case Intrinsic::nvvm_tcgen05_ld_red_32x32b_x2_f32:
+    return NVPTXISD::TCGEN05_LD_RED_32x32b_X2_F32;
+  case Intrinsic::nvvm_tcgen05_ld_red_32x32b_x32_f32:
+    return NVPTXISD::TCGEN05_LD_RED_32x32b_X32_F32;
+  case Intrinsic::nvvm_tcgen05_ld_red_32x32b_x4_f32:
+    return NVPTXISD::TCGEN05_LD_RED_32x32b_X4_F32;
+  case Intrinsic::nvvm_tcgen05_ld_red_32x32b_x64_f32:
+    return NVPTXISD::TCGEN05_LD_RED_32x32b_X64_F32;
+  case Intrinsic::nvvm_tcgen05_ld_red_32x32b_x8_f32:
+    return NVPTXISD::TCGEN05_LD_RED_32x32b_X8_F32;
+  case Intrinsic::nvvm_tcgen05_ld_red_16x32bx2_x128_i32:
+    return NVPTXISD::TCGEN05_LD_RED_16x32bx2_X128_I32;
+  case Intrinsic::nvvm_tcgen05_ld_red_16x32bx2_x16_i32:
+    return NVPTXISD::TCGEN05_LD_RED_16x32bx2_X16_I32;
+  case Intrinsic::nvvm_tcgen05_ld_red_16x32bx2_x2_i32:
+    return NVPTXISD::TCGEN05_LD_RED_16x32bx2_X2_I32;
+  case Intrinsic::nvvm_tcgen05_ld_red_16x32bx2_x32_i32:
+    return NVPTXISD::TCGEN05_LD_RED_16x32bx2_X32_I32;
+  case Intrinsic::nvvm_tcgen05_ld_red_16x32bx2_x4_i32:
+    return NVPTXISD::TCGEN05_LD_RED_16x32bx2_X4_I32;
+  case Intrinsic::nvvm_tcgen05_ld_red_16x32bx2_x64_i32:
+    return NVPTXISD::TCGEN05_LD_RED_16x32bx2_X64_I32;
+  case Intrinsic::nvvm_tcgen05_ld_red_16x32bx2_x8_i32:
+    return NVPTXISD::TCGEN05_LD_RED_16x32bx2_X8_I32;
+  case Intrinsic::nvvm_tcgen05_ld_red_32x32b_x128_i32:
+    return NVPTXISD::TCGEN05_LD_RED_32x32b_X128_I32;
+  case Intrinsic::nvvm_tcgen05_ld_red_32x32b_x16_i32:
+    return NVPTXISD::TCGEN05_LD_RED_32x32b_X16_I32;
+  case Intrinsic::nvvm_tcgen05_ld_red_32x32b_x2_i32:
+    return NVPTXISD::TCGEN05_LD_RED_32x32b_X2_I32;
+  case Intrinsic::nvvm_tcgen05_ld_red_32x32b_x32_i32:
+    return NVPTXISD::TCGEN05_LD_RED_32x32b_X32_I32;
+  case Intrinsic::nvvm_tcgen05_ld_red_32x32b_x4_i32:
+    return NVPTXISD::TCGEN05_LD_RED_32x32b_X4_I32;
+  case Intrinsic::nvvm_tcgen05_ld_red_32x32b_x64_i32:
+    return NVPTXISD::TCGEN05_LD_RED_32x32b_X64_I32;
+  case Intrinsic::nvvm_tcgen05_ld_red_32x32b_x8_i32:
+    return NVPTXISD::TCGEN05_LD_RED_32x32b_X8_I32;
+  default:
+    llvm_unreachable("Invalid tcgen05.ld.red intrinsic ID");
+  }
+}
+
+// Lower vector return type of tcgen05.ld intrinsics
+static std::optional<std::tuple<SDValue, SDValue, SDValue>>
+lowerTcgen05LdRed(SDNode *N, SelectionDAG &DAG, bool IsFloat = false,
+                  bool HasOffset = false) {
+  SDLoc DL(N);
+  EVT ResVT = N->getValueType(0);
+  if (!ResVT.isVector())
+    return {}; // already legalized.
+
+  const unsigned NumElts = ResVT.getVectorNumElements();
+
+  // Create the return type of the instructions
+  SmallVector<EVT, 132> ListVTs{NumElts, IsFloat ? MVT::f32 : MVT::i32};
+
+  ListVTs.push_back(N->getValueType(1)); // RedVal
+  ListVTs.push_back(N->getValueType(2)); // Chain
+
+  SDVTList ResVTs = DAG.getVTList(ListVTs);
+
+  // Prepare the Operands
+  SmallVector<SDValue, 8> Ops{N->getOperand(0),  // Chain
+                              N->getOperand(2)}; // ptr
+  unsigned OpIdx = 2;
+
+  if (HasOffset)
+    Ops.push_back(N->getOperand(++OpIdx)); // offset
+
+  Ops.push_back(N->getOperand(++OpIdx)); // redOp
+
+  if (IsFloat) {
+    Ops.push_back(N->getOperand(++OpIdx)); // abs
+    Ops.push_back(N->getOperand(++OpIdx)); // nan
+  }
+
----------------
schwarzschild-radius wrote:

Thanks Artem! Simplified the above conditional code into a loop

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