From 4be86d18c0fc97a5eab26628e2a830f914c591eb Mon Sep 17 00:00:00 2001
From: Andrew Litteken <andrew.litteken@gmail.com>
Date: Wed, 22 Dec 2021 17:42:10 -0600
Subject: [PATCH] [IROutliner] Disallow outlining calls that return twice.

Functions that return twice can cause the IR Outliner to miscompile the given program. These function rely on information about the stack to be the same, and this may not necessarily be the case if called from an outlined function. So, we simply call these instructions illegal for the outliner to remove.

Reviewers: paquette

Differential Revision: https://reviews.llvm.org/D110007
---
 llvm/include/llvm/Transforms/IPO/IROutliner.h      |  5 ++
 .../Transforms/IROutliner/illegal-returns-twice.ll | 66 ++++++++++++++++++++++
 2 files changed, 71 insertions(+)
 create mode 100644 llvm/test/Transforms/IROutliner/illegal-returns-twice.ll

diff --git a/llvm/include/llvm/Transforms/IPO/IROutliner.h b/llvm/include/llvm/Transforms/IPO/IROutliner.h
index dcae445..16f597a 100644
--- a/llvm/include/llvm/Transforms/IPO/IROutliner.h
+++ b/llvm/include/llvm/Transforms/IPO/IROutliner.h
@@ -367,6 +367,11 @@ private:
       Function *F = CI.getCalledFunction();
       if (!F || CI.isIndirectCall() || !F->hasName())
         return false;
+      // Returning twice can cause issues with the state of the function call
+      // that were not expected when the function was used, so we do not include
+      // the call in outlined functions.
+      if (CI.canReturnTwice())
+        return false;
       return true;
     }
     // TODO: Handle FreezeInsts.  Since a frozen value could be frozen inside
diff --git a/llvm/test/Transforms/IROutliner/illegal-returns-twice.ll b/llvm/test/Transforms/IROutliner/illegal-returns-twice.ll
new file mode 100644
index 0000000..31113ab
--- /dev/null
+++ b/llvm/test/Transforms/IROutliner/illegal-returns-twice.ll
@@ -0,0 +1,66 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt -S -verify -iroutliner -ir-outlining-no-cost < %s | FileCheck %s
+
+; This test checks that we do not outline functions that are marked as returns
+; twice, since these can alter the frame of the function and affect how the
+; outliner behaves, causing miscompiles.
+
+; Function Attrs: optsize returns_twice
+declare i32 @setjmp(i32*) local_unnamed_addr #1
+@tmp_jmpb = global [37 x i32] zeroinitializer, align 16
+
+define void @function1() {
+; CHECK-LABEL: @function1(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[A:%.*]] = alloca i32, align 4
+; CHECK-NEXT:    [[B:%.*]] = alloca i32, align 4
+; CHECK-NEXT:    [[C:%.*]] = alloca i32, align 4
+; CHECK-NEXT:    store i32 2, i32* [[A]], align 4
+; CHECK-NEXT:    store i32 3, i32* [[B]], align 4
+; CHECK-NEXT:    store i32 4, i32* [[C]], align 4
+; CHECK-NEXT:    [[CALL:%.*]] = call i32 @setjmp(i32* getelementptr inbounds ([37 x i32], [37 x i32]* @tmp_jmpb, i64 0, i64 0))
+; CHECK-NEXT:    call void @outlined_ir_func_0(i32* [[A]], i32* [[B]], i32* [[C]])
+; CHECK-NEXT:    ret void
+;
+entry:
+  %a = alloca i32, align 4
+  %b = alloca i32, align 4
+  %c = alloca i32, align 4
+  store i32 2, i32* %a, align 4
+  store i32 3, i32* %b, align 4
+  store i32 4, i32* %c, align 4
+  %call = call i32 @setjmp(i32* getelementptr inbounds ([37 x i32], [37 x i32]* @tmp_jmpb, i64 0, i64 0))
+  %al = load i32, i32* %a
+  %bl = load i32, i32* %b
+  %cl = load i32, i32* %c
+  ret void
+}
+
+define void @function2() {
+; CHECK-LABEL: @function2(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[A:%.*]] = alloca i32, align 4
+; CHECK-NEXT:    [[B:%.*]] = alloca i32, align 4
+; CHECK-NEXT:    [[C:%.*]] = alloca i32, align 4
+; CHECK-NEXT:    store i32 2, i32* [[A]], align 4
+; CHECK-NEXT:    store i32 3, i32* [[B]], align 4
+; CHECK-NEXT:    store i32 4, i32* [[C]], align 4
+; CHECK-NEXT:    [[CALL:%.*]] = call i32 @setjmp(i32* getelementptr inbounds ([37 x i32], [37 x i32]* @tmp_jmpb, i64 0, i64 0))
+; CHECK-NEXT:    call void @outlined_ir_func_0(i32* [[A]], i32* [[B]], i32* [[C]])
+; CHECK-NEXT:    ret void
+;
+entry:
+  %a = alloca i32, align 4
+  %b = alloca i32, align 4
+  %c = alloca i32, align 4
+  store i32 2, i32* %a, align 4
+  store i32 3, i32* %b, align 4
+  store i32 4, i32* %c, align 4
+  %call = call i32 @setjmp(i32* getelementptr inbounds ([37 x i32], [37 x i32]* @tmp_jmpb, i64 0, i64 0))
+  %al = load i32, i32* %a
+  %bl = load i32, i32* %b
+  %cl = load i32, i32* %c
+  ret void
+}
+
+attributes #1 = { optsize returns_twice }
-- 
2.7.4