[llvm] [Scalarizer] Fix to only scalarize if intrinsic was marked as isTriviallyScalarizable (PR #113625)

[Farzon Lotfi via llvm-commits]

https://github.com/farzonl created https://github.com/llvm/llvm-project/pull/113625

fixes #113624

>From d1e3d20af6f59ba478ce2ebd5545a228eb164054 Mon Sep 17 00:00:00 2001
From: Farzon Lotfi <farzonlotfi at microsoft.com>
Date: Thu, 24 Oct 2024 17:59:01 -0400
Subject: [PATCH] [Scalarizer] Fix to only scalarize if intrinsic was marked as
 isTriviallyScalarizable

fixes #113624
---
 llvm/lib/Transforms/Scalar/Scalarizer.cpp        | 10 ++++++++++
 llvm/test/Transforms/Scalarizer/uadd_overflow.ll | 16 ++++++++++++++++
 2 files changed, 26 insertions(+)
 create mode 100644 llvm/test/Transforms/Scalarizer/uadd_overflow.ll

diff --git a/llvm/lib/Transforms/Scalar/Scalarizer.cpp b/llvm/lib/Transforms/Scalar/Scalarizer.cpp
index 772f4c6c35ddec..94e10b707102af 100644
--- a/llvm/lib/Transforms/Scalar/Scalarizer.cpp
+++ b/llvm/lib/Transforms/Scalar/Scalarizer.cpp
@@ -1084,6 +1084,16 @@ bool ScalarizerVisitor::visitExtractValueInst(ExtractValueInst &EVI) {
   ValueVector Res;
   if (!isStructOfMatchingFixedVectors(OpTy))
     return false;
+  if (CallInst *CI = dyn_cast<CallInst>(Op)) {
+    Function *F = CI->getCalledFunction();
+    if (!F)
+      return false;
+    Intrinsic::ID ID = F->getIntrinsicID();
+    if (ID == Intrinsic::not_intrinsic || !isTriviallyScalarizable(ID))
+      return false;
+    // Note: Only proceed if Operand is a`CallInst` and it is defined in `isTriviallyScalarizable`.
+  } else
+    return false;
   Type *VecType = cast<FixedVectorType>(OpTy->getContainedType(0));
   std::optional<VectorSplit> VS = getVectorSplit(VecType);
   if (!VS)
diff --git a/llvm/test/Transforms/Scalarizer/uadd_overflow.ll b/llvm/test/Transforms/Scalarizer/uadd_overflow.ll
new file mode 100644
index 00000000000000..39094451523a5e
--- /dev/null
+++ b/llvm/test/Transforms/Scalarizer/uadd_overflow.ll
@@ -0,0 +1,16 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 5
+; RUN: opt %s -passes='function(scalarizer)' -S | FileCheck %s
+
+; Test to make sure that struct return intrinsics that are not `isTriviallyScalarizable` do not get scalarized.
+
+define <3 x i32> @test_(<3 x i32> %a, <3 x i32> %b) {
+; CHECK-LABEL: define <3 x i32> @test_(
+; CHECK-SAME: <3 x i32> [[A:%.*]], <3 x i32> [[B:%.*]]) {
+; CHECK-NEXT:    [[R:%.*]] = call { <3 x i32>, <3 x i1> } @llvm.uadd.with.overflow.v3i32(<3 x i32> [[B]], <3 x i32> [[B]])
+; CHECK-NEXT:    [[EL:%.*]] = extractvalue { <3 x i32>, <3 x i1> } [[R]], 0
+; CHECK-NEXT:    ret <3 x i32> [[EL]]
+;
+  %r = call { <3 x i32>, <3 x i1> } @llvm.uadd.with.overflow.v3i32(<3 x i32> %b, <3 x i32> %b)
+  %el = extractvalue { <3 x i32>, <3 x i1> } %r, 0
+  ret <3 x i32> %el
+}