[llvm] [VectorCombine] Fix crash in scalarizeVPIntrinsic (PR #72039)

[Michael Maitland via llvm-commits]

https://github.com/michaelmaitland updated https://github.com/llvm/llvm-project/pull/72039

>From ea6077076d66468bcdcd5a5a97be4838901ec55f Mon Sep 17 00:00:00 2001
From: Michael Maitland <michaeltmaitland at gmail.com>
Date: Sat, 11 Nov 2023 11:52:17 -0800
Subject: [PATCH 1/3] [VectorCombine] Fix crash in scalarizeVPIntrinsic

When getSplatOp returns nullptr, the intrinsic cannot be scalarized.
This patch includes a test case that fixes a crash from trying to
scalarize the VPIntrinsic when getSplatOp returns nullptr. This case can
be scalarized in the future by improving getSplatOp, but the checks
added in this patch should remain.

This fixes https://github.com/llvm/llvm-project/issues/72034.
---
 llvm/lib/Transforms/Vectorize/VectorCombine.cpp  |  2 ++
 ...vpintrin-scalarization-shufflevector-splat.ll | 16 ++++++++++++++++
 2 files changed, 18 insertions(+)
 create mode 100644 llvm/test/Transforms/VectorCombine/RISCV/vpintrin-scalarization-shufflevector-splat.ll

diff --git a/llvm/lib/Transforms/Vectorize/VectorCombine.cpp b/llvm/lib/Transforms/Vectorize/VectorCombine.cpp
index 9a91ce207bb04c9..6d79c9b13add067 100644
--- a/llvm/lib/Transforms/Vectorize/VectorCombine.cpp
+++ b/llvm/lib/Transforms/Vectorize/VectorCombine.cpp
@@ -843,6 +843,8 @@ bool VectorCombine::scalarizeVPIntrinsic(Instruction &I) {
 
   Value *ScalarOp0 = getSplatValue(Op0);
   Value *ScalarOp1 = getSplatValue(Op1);
+  if (!ScalarOp0 || !ScalarOp1)
+    return false;
   Value *ScalarVal =
       ScalarIntrID
           ? Builder.CreateIntrinsic(VecTy->getScalarType(), *ScalarIntrID,
diff --git a/llvm/test/Transforms/VectorCombine/RISCV/vpintrin-scalarization-shufflevector-splat.ll b/llvm/test/Transforms/VectorCombine/RISCV/vpintrin-scalarization-shufflevector-splat.ll
new file mode 100644
index 000000000000000..d3c477d180aa54f
--- /dev/null
+++ b/llvm/test/Transforms/VectorCombine/RISCV/vpintrin-scalarization-shufflevector-splat.ll
@@ -0,0 +1,16 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 3
+; RUN: opt -S -mtriple=riscv64 -mattr=+v %s -passes=vector-combine | FileCheck %s --check-prefix
+
+declare <4 x i64> @llvm.vp.add.v4i64(<4 x i64>, <4 x i64>, <4 x i1>, i32)
+
+define <4 x i64> @add_v4i64_allonesmask(<4 x i64> %x) {
+; CHECK-LABEL: define <4 x i64> @add_v4i64_allonesmask(
+; CHECK-SAME: <4 x i64> [[X:%.*]]) #[[ATTR1:[0-9]+]] {
+; CHECK-NEXT:    [[TMP1:%.*]] = shufflevector <4 x i64> [[X]], <4 x i64> zeroinitializer, <4 x i32> zeroinitializer
+; CHECK-NEXT:    [[TMP2:%.*]] = call <4 x i64> @llvm.vp.add.v4i64(<4 x i64> [[TMP1]], <4 x i64> zeroinitializer, <4 x i1> <i1 true, i1 true, i1 true, i1 true>, i32 0)
+; CHECK-NEXT:    ret <4 x i64> [[TMP2]]
+;
+  %1 = shufflevector <4 x i64> %x, <4 x i64> zeroinitializer, <4 x i32> zeroinitializer
+  %2 = call <4 x i64> @llvm.vp.add.v4i64(<4 x i64> %1, <4 x i64> zeroinitializer, <4 x i1> <i1 true, i1 true, i1 true, i1 true>, i32 0)
+  ret <4 x i64> %2
+}

>From 2965b5583f349de9646dbde65c42b911fb3fe07b Mon Sep 17 00:00:00 2001
From: Michael Maitland <michaeltmaitland at gmail.com>
Date: Sat, 11 Nov 2023 16:11:57 -0800
Subject: [PATCH 2/3] move getSplatValue check earlier in the function

---
 llvm/lib/Transforms/Vectorize/VectorCombine.cpp | 11 +++++++----
 1 file changed, 7 insertions(+), 4 deletions(-)

diff --git a/llvm/lib/Transforms/Vectorize/VectorCombine.cpp b/llvm/lib/Transforms/Vectorize/VectorCombine.cpp
index 6d79c9b13add067..71ff21082c732d0 100644
--- a/llvm/lib/Transforms/Vectorize/VectorCombine.cpp
+++ b/llvm/lib/Transforms/Vectorize/VectorCombine.cpp
@@ -753,6 +753,13 @@ bool VectorCombine::scalarizeVPIntrinsic(Instruction &I) {
   if (!isSplatValue(Op0) || !isSplatValue(Op1))
     return false;
 
+  // Check getSplatValue early in this function, to avoid doing unnecessary
+  // work.
+  Value *ScalarOp0 = getSplatValue(Op0);
+  Value *ScalarOp1 = getSplatValue(Op1);
+  if (!ScalarOp0 || !ScalarOp1)
+    return false;
+
   // For the binary VP intrinsics supported here, the result on disabled lanes
   // is a poison value. For now, only do this simplification if all lanes
   // are active.
@@ -841,10 +848,6 @@ bool VectorCombine::scalarizeVPIntrinsic(Instruction &I) {
   if (!SafeToSpeculate && !isKnownNonZero(EVL, DL, 0, &AC, &VPI, &DT))
     return false;
 
-  Value *ScalarOp0 = getSplatValue(Op0);
-  Value *ScalarOp1 = getSplatValue(Op1);
-  if (!ScalarOp0 || !ScalarOp1)
-    return false;
   Value *ScalarVal =
       ScalarIntrID
           ? Builder.CreateIntrinsic(VecTy->getScalarType(), *ScalarIntrID,

>From 81910b20ebf7209a249956cbb051597dde60f9f1 Mon Sep 17 00:00:00 2001
From: Michael Maitland <michaeltmaitland at gmail.com>
Date: Sat, 11 Nov 2023 16:53:46 -0800
Subject: [PATCH 3/3] fix check-prefix option

---
 .../RISCV/vpintrin-scalarization-shufflevector-splat.ll         | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/llvm/test/Transforms/VectorCombine/RISCV/vpintrin-scalarization-shufflevector-splat.ll b/llvm/test/Transforms/VectorCombine/RISCV/vpintrin-scalarization-shufflevector-splat.ll
index d3c477d180aa54f..80602d12e26f8a8 100644
--- a/llvm/test/Transforms/VectorCombine/RISCV/vpintrin-scalarization-shufflevector-splat.ll
+++ b/llvm/test/Transforms/VectorCombine/RISCV/vpintrin-scalarization-shufflevector-splat.ll
@@ -1,5 +1,5 @@
 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 3
-; RUN: opt -S -mtriple=riscv64 -mattr=+v %s -passes=vector-combine | FileCheck %s --check-prefix
+; RUN: opt -S -mtriple=riscv64 -mattr=+v %s -passes=vector-combine | FileCheck %s --check-prefix=CHECK
 
 declare <4 x i64> @llvm.vp.add.v4i64(<4 x i64>, <4 x i64>, <4 x i1>, i32)