[llvm] [SelectionDAG] Scalarize one-lane intrinsic results (PR #223622)

[via llvm-commits]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-backend-nvptx

Author: Bastian Hagedorn (bastianhagedorn)

<details>
<summary>Changes</summary>

Scalarize `INTRINSIC_WO_CHAIN` results of type `vector<1 x T>` by scalarizing their one-lane vector operands and emitting the scalar intrinsic overload.

This fixes the NVVM EX2 `vector<1xf32>` overloads, which previously reached SelectionDAG vector legalization without an intrinsic scalarization path and crashed NVPTX instruction selection. Both normal and FTZ forms now emit the existing scalar PTX instructions.

Test: `llc -mcpu=sm_50 < llvm/test/CodeGen/NVPTX/f32-ex2.ll | FileCheck llvm/test/CodeGen/NVPTX/f32-ex2.ll`.

---
Full diff: https://github.com/llvm/llvm-project/pull/223622.diff


3 Files Affected:

- (modified) llvm/lib/CodeGen/SelectionDAG/LegalizeTypes.h (+1) 
- (modified) llvm/lib/CodeGen/SelectionDAG/LegalizeVectorTypes.cpp (+29) 
- (modified) llvm/test/CodeGen/NVPTX/f32-ex2.ll (+20) 


``````````diff
diff --git a/llvm/lib/CodeGen/SelectionDAG/LegalizeTypes.h b/llvm/lib/CodeGen/SelectionDAG/LegalizeTypes.h
index 6fc6d61c6a38df..ccff9a07a4d298 100644
--- a/llvm/lib/CodeGen/SelectionDAG/LegalizeTypes.h
+++ b/llvm/lib/CodeGen/SelectionDAG/LegalizeTypes.h
@@ -826,6 +826,7 @@ class LLVM_LIBRARY_VISIBILITY DAGTypeLegalizer {
   SDValue ScalarizeVecRes_TernaryOp(SDNode *N);
   SDValue ScalarizeVecRes_UnaryOp(SDNode *N);
   SDValue ScalarizeVecRes_StrictFPOp(SDNode *N);
+  SDValue ScalarizeVecRes_INTRINSIC_WO_CHAIN(SDNode *N);
   SDValue ScalarizeVecRes_OverflowOp(SDNode *N, unsigned ResNo);
   SDValue ScalarizeVecRes_InregOp(SDNode *N);
   SDValue ScalarizeVecRes_VecInregOp(SDNode *N);
diff --git a/llvm/lib/CodeGen/SelectionDAG/LegalizeVectorTypes.cpp b/llvm/lib/CodeGen/SelectionDAG/LegalizeVectorTypes.cpp
index f6e7feab57ee94..71d1f7c632ceea 100644
--- a/llvm/lib/CodeGen/SelectionDAG/LegalizeVectorTypes.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/LegalizeVectorTypes.cpp
@@ -83,6 +83,8 @@ void DAGTypeLegalizer::ScalarizeVectorResult(SDNode *N, unsigned ResNo) {
     R = ScalarizeVecRes_ATOMIC_LOAD(cast<AtomicSDNode>(N));
     break;
   case ISD::LOAD:           R = ScalarizeVecRes_LOAD(cast<LoadSDNode>(N));break;
+  case ISD::INTRINSIC_WO_CHAIN:
+    R = ScalarizeVecRes_INTRINSIC_WO_CHAIN(N); break;
   case ISD::SCALAR_TO_VECTOR:  R = ScalarizeVecRes_SCALAR_TO_VECTOR(N); break;
   case ISD::VECTOR_DEINTERLEAVE:
   case ISD::VECTOR_INTERLEAVE:
@@ -586,6 +588,33 @@ SDValue DAGTypeLegalizer::ScalarizeVecRes_LOAD(LoadSDNode *N) {
   return Result;
 }
 
+SDValue DAGTypeLegalizer::ScalarizeVecRes_INTRINSIC_WO_CHAIN(SDNode *N) {
+  assert(N->getValueType(0).getVectorNumElements() == 1 &&
+         "Unexpected vector type");
+
+  SDLoc DL(N);
+  SmallVector<SDValue, 4> Ops{N->getOperand(0)};
+  for (unsigned I = 1; I < N->getNumOperands(); ++I) {
+    SDValue Operand = N->getOperand(I);
+    EVT OperandVT = Operand.getValueType();
+    if (OperandVT.isVector()) {
+      assert(OperandVT.getVectorNumElements() == 1 &&
+             "Unexpected vector operand type");
+      if (getTypeAction(OperandVT) == TargetLowering::TypeScalarizeVector) {
+        Operand = GetScalarizedVector(Operand);
+      } else {
+        Operand = DAG.getExtractVectorElt(DL, OperandVT.getVectorElementType(),
+                                          Operand, 0);
+      }
+    }
+    Ops.push_back(Operand);
+  }
+
+  return DAG.getNode(ISD::INTRINSIC_WO_CHAIN, DL,
+                     N->getValueType(0).getVectorElementType(), Ops,
+                     N->getFlags());
+}
+
 SDValue DAGTypeLegalizer::ScalarizeVecRes_UnaryOp(SDNode *N) {
   // Get the dest type - it doesn't always match the input type, e.g. int_to_fp.
   EVT DestVT = N->getValueType(0).getVectorElementType();
diff --git a/llvm/test/CodeGen/NVPTX/f32-ex2.ll b/llvm/test/CodeGen/NVPTX/f32-ex2.ll
index db3dd4a9e6011b..523e681c9ea9e3 100644
--- a/llvm/test/CodeGen/NVPTX/f32-ex2.ll
+++ b/llvm/test/CodeGen/NVPTX/f32-ex2.ll
@@ -5,6 +5,8 @@ target triple = "nvptx-nvidia-cuda"
 
 declare float @llvm.nvvm.ex2.approx.f32(float)
 declare float @llvm.nvvm.ex2.approx.ftz.f32(float)
+declare <1 x float> @llvm.nvvm.ex2.approx.v1f32(<1 x float>)
+declare <1 x float> @llvm.nvvm.ex2.approx.ftz.v1f32(<1 x float>)
 
 ; CHECK-LABEL: ex2_float
 define float @ex2_float(float %0) {
@@ -35,3 +37,21 @@ define float @ex2_float_ftz(float %0) {
   %res = call float @llvm.nvvm.ex2.approx.ftz.f32(float %0)
   ret float %res
 }
+
+; CHECK-LABEL: ex2_float_v1
+define <1 x float> @ex2_float_v1(<1 x float> %0) {
+; CHECK-LABEL: ex2_float_v1(
+; CHECK:       {
+; CHECK:         ex2.approx.f32
+  %res = call <1 x float> @llvm.nvvm.ex2.approx.v1f32(<1 x float> %0)
+  ret <1 x float> %res
+}
+
+; CHECK-LABEL: ex2_float_v1_ftz
+define <1 x float> @ex2_float_v1_ftz(<1 x float> %0) {
+; CHECK-LABEL: ex2_float_v1_ftz(
+; CHECK:       {
+; CHECK:         ex2.approx.ftz.f32
+  %res = call <1 x float> @llvm.nvvm.ex2.approx.ftz.v1f32(<1 x float> %0)
+  ret <1 x float> %res
+}

``````````

</details>


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