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

[Bastian Hagedorn via llvm-commits]

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

>From e2c3a769fff39ebbc53b119f830496640b48bf4b Mon Sep 17 00:00:00 2001
From: Bastian Hagedorn <bhagedorn at nvidia.com>
Date: Tue, 15 Sep 2026 07:16:19 +0000
Subject: [PATCH 1/2] [SelectionDAG] Scalarize one-lane intrinsic results

---
 llvm/lib/CodeGen/SelectionDAG/LegalizeTypes.h |  1 +
 .../SelectionDAG/LegalizeVectorTypes.cpp      | 29 +++++++++++++++++++
 llvm/test/CodeGen/NVPTX/f32-ex2.ll            | 20 +++++++++++++
 3 files changed, 50 insertions(+)

diff --git a/llvm/lib/CodeGen/SelectionDAG/LegalizeTypes.h b/llvm/lib/CodeGen/SelectionDAG/LegalizeTypes.h
index 6fc6d61c6a38d..ccff9a07a4d29 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 f6e7feab57ee9..71d1f7c632cee 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 db3dd4a9e6011..523e681c9ea9e 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
+}

>From 14fa1393b2e645a6aae1185d43317018a397814c Mon Sep 17 00:00:00 2001
From: Bastian Hagedorn <bhagedorn at nvidia.com>
Date: Tue, 15 Sep 2026 08:33:20 +0000
Subject: [PATCH 2/2] [SelectionDAG] Fix formatting

---
 llvm/lib/CodeGen/SelectionDAG/LegalizeVectorTypes.cpp | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/llvm/lib/CodeGen/SelectionDAG/LegalizeVectorTypes.cpp b/llvm/lib/CodeGen/SelectionDAG/LegalizeVectorTypes.cpp
index 71d1f7c632cee..487854fa3f4fd 100644
--- a/llvm/lib/CodeGen/SelectionDAG/LegalizeVectorTypes.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/LegalizeVectorTypes.cpp
@@ -84,7 +84,8 @@ void DAGTypeLegalizer::ScalarizeVectorResult(SDNode *N, unsigned ResNo) {
     break;
   case ISD::LOAD:           R = ScalarizeVecRes_LOAD(cast<LoadSDNode>(N));break;
   case ISD::INTRINSIC_WO_CHAIN:
-    R = ScalarizeVecRes_INTRINSIC_WO_CHAIN(N); break;
+    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: