[FuncSpec]Fix assertion failure when value is not added to solver

This patch improves the fix in D110529 to prevent from crashing on value
with byval attribute that is not added in SCCP solver.

Authored-by: sinan.lin@linux.alibaba.com
Reviewed By: ChuanqiXu

Differential Revision: https://reviews.llvm.org/D126355

llvm/lib/Transforms/IPO/FunctionSpecialization.cpp

@@ -710,6 +710,11 @@ class FunctionSpecializer {
                             SmallVectorImpl<CallArgBinding> &Constants) {
     Function *F = A->getParent();
 
+    // SCCP solver does not record an argument that will be constructed on
+    // stack.
+    if (A->hasByValAttr() && !F->onlyReadsMemory())
+      return;
+
     // Iterate over all the call sites of the argument's parent function.
     for (User *U : F->users()) {
       if (!isa<CallInst>(U) && !isa<InvokeInst>(U))
@@ -729,12 +734,6 @@ class FunctionSpecializer {
       if (isa<PoisonValue>(V))
         return;
 
-      // For now, constant expressions are fine but only if they are function
-      // calls.
-      if (auto *CE = dyn_cast<ConstantExpr>(V))
-        if (!isa<Function>(CE->getOperand(0)))
-          return;
-
       // TrackValueOfGlobalVariable only tracks scalar global variables.
       if (auto *GV = dyn_cast<GlobalVariable>(V)) {
         // Check if we want to specialize on the address of non-constant

llvm/test/Transforms/FunctionSpecialization/function-specialization-constant-expression.ll

@@ -10,11 +10,6 @@
 @Global = internal constant %struct {i8 0, i16 1, i32 2, i64 3, i64 4}
 
 define internal i64 @func2(i64 *%x) {
-; CHECK-LABEL: @func2(
-; CHECK-NEXT:  entry:
-; CHECK-NEXT:    [[VAL:%.*]] = ptrtoint i64* [[X:%.*]] to i64
-; CHECK-NEXT:    ret i64 [[VAL]]
-;
 entry:
   %val = ptrtoint i64* %x to i64
   ret i64 %val
@@ -36,10 +31,10 @@ define internal i64 @zoo(i1 %flag) {
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    br i1 [[FLAG:%.*]], label [[PLUS:%.*]], label [[MINUS:%.*]]
 ; CHECK:       plus:
-; CHECK-NEXT:    [[TMP0:%.*]] = call i64 @func2(i64* getelementptr inbounds ([[STRUCT:%.*]], %struct* @Global, i32 0, i32 3))
+; CHECK-NEXT:    [[TMP0:%.*]] = call i64 @func2.2(i64* getelementptr inbounds ([[STRUCT:%.*]], %struct* @Global, i32 0, i32 3))
 ; CHECK-NEXT:    br label [[MERGE:%.*]]
 ; CHECK:       minus:
-; CHECK-NEXT:    [[TMP1:%.*]] = call i64 @func2(i64* getelementptr inbounds ([[STRUCT]], %struct* @Global, i32 0, i32 4))
+; CHECK-NEXT:    [[TMP1:%.*]] = call i64 @func2.1(i64* getelementptr inbounds ([[STRUCT]], %struct* @Global, i32 0, i32 4))
 ; CHECK-NEXT:    br label [[MERGE]]
 ; CHECK:       merge:
 ; CHECK-NEXT:    [[TMP2:%.*]] = phi i64 [ [[TMP0]], [[PLUS]] ], [ [[TMP1]], [[MINUS]] ]

llvm/test/Transforms/FunctionSpecialization/function-specialization-constant-expression4.ll

@@ -0,0 +1,32 @@
+; RUN: opt -function-specialization -force-function-specialization -S < %s | FileCheck %s
+
+; Check that we don't crash and specialise on a function call with byval attribute.
+
+; CHECK-NOT: wombat.{{[0-9]+}}
+
+declare i32* @quux()
+declare i32* @eggs()
+
+define i32 @main() {
+; CHECK:       bb:
+; CHECK-NEXT:    tail call void @wombat(i8* undef, i64 undef, i64 undef, i32* byval(i32) bitcast (i32* ()* @quux to i32*))
+; CHECK-NEXT:    tail call void @wombat(i8* undef, i64 undef, i64 undef, i32* byval(i32) bitcast (i32* ()* @eggs to i32*))
+; CHECK-NEXT:    ret i32 undef
+;
+bb:
+  tail call void @wombat(i8* undef, i64 undef, i64 undef, i32* byval(i32) bitcast (i32* ()* @quux to i32*))
+  tail call void @wombat(i8* undef, i64 undef, i64 undef, i32* byval(i32) bitcast (i32* ()* @eggs to i32*))
+  ret i32 undef
+}
+
+define internal void @wombat(i8* %arg, i64 %arg1, i64 %arg2, i32* byval(i32) %func) {
+; CHECK:       bb2:
+; CHECK-NEXT:    [[FUNPTR:%.*]] = bitcast i32* %func to i32* (i8*, i8*)*
+; CHECK-NEXT:    [[TMP:%.*]] = tail call i32* [[FUNPTR]](i8* undef, i8* undef)
+; CHECK-NEXT:    ret void
+;
+bb2:
+  %mycall = bitcast i32* %func to i32* (i8*, i8*)*
+  %tmp = tail call i32* %mycall(i8* undef, i8* undef)
+  ret void
+}

llvm/test/Transforms/FunctionSpecialization/function-specialization-constant-expression5.ll

@@ -0,0 +1,46 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt -function-specialization -force-function-specialization -func-specialization-on-address -S < %s | FileCheck %s
+
+; Check that we don't crash and specialise on a scalar global variable with byval attribute.
+
+; CHECK-NOT: wombat.{{[0-9]+}}
+
+%struct.pluto = type { %struct.spam }
+%struct.quux = type { i16 }
+%struct.spam = type { i16 }
+
+@global.5 = external dso_local global i128
+@global.12 = external global %struct.quux
+
+define internal i16 @wobble(%struct.quux* %arg, i16 %arg1, i128* byval(i128) %arg2, %struct.quux* %arg3) {
+; CHECK-LABEL: @wobble(
+; CHECK-NEXT:  bb:
+; CHECK-NEXT:    [[TMP34:%.*]] = bitcast i128* [[ARG2:%.*]] to i16*
+; CHECK-NEXT:    unreachable
+;
+bb:
+  %tmp34 = bitcast i128* %arg2 to i16*
+  unreachable
+}
+
+define internal i16 @snork() {
+; CHECK-LABEL: @snork(
+; CHECK-NEXT:  bb4:
+; CHECK-NEXT:    [[TMP35:%.*]] = call i16 @wobble(%struct.quux* undef, i16 2, i128* byval(i128) @global.5, %struct.quux* @global.12)
+; CHECK-NEXT:    unreachable
+;
+bb4:
+  %tmp35 = call i16 @wobble(%struct.quux* undef, i16 2, i128* byval(i128) @global.5, %struct.quux* @global.12)
+  unreachable
+}
+
+define i16 @main() {
+; CHECK-LABEL: @main(
+; CHECK-NEXT:  bb:
+; CHECK-NEXT:    [[TMP:%.*]] = call i16 @snork()
+; CHECK-NEXT:    unreachable
+;
+bb:
+  %tmp = call i16 @snork()
+  unreachable
+}