Limit LLVM routine stack size to 512 KiB

Fuzzing tests generate shaders with large arrays or very high numbers of
local variables, which can cause stack overflow. We need to limit the
allowable stack memory usage of generated routines.

Note this change does not yet gracefully deal with routines which exceed
this limit. They will cause a null pointer dereference instead of a
stack overflow.

The 512 KiB stack size limit is chosen to prevent actual stack overflow
for a 1 MiB stack, assuming some earlier calls might want to use the
stack. Also, our legacy 'ASM' compiler for GLSL allocates 4096
'registers' of 4 components for 128-bit SIMD, which already requires
256 KiB.

Bug: b/157555596
Change-Id: I25c57420f6d2af323ce98faf515feca0aa834a4a
Reviewed-on: https://swiftshader-review.googlesource.com/c/SwiftShader/+/51548
Presubmit-Ready: Nicolas Capens <nicolascapens@google.com>
Kokoro-Result: kokoro <noreply+kokoro@google.com>
Reviewed-by: Antonio Maiorano <amaiorano@google.com>
Tested-by: Nicolas Capens <nicolascapens@google.com>
Commit-Queue: Nicolas Capens <nicolascapens@google.com>
diff --git a/src/Reactor/LLVMAsm.cpp b/src/Reactor/LLVMAsm.cpp
index ab6e093..c4d831a 100644
--- a/src/Reactor/LLVMAsm.cpp
+++ b/src/Reactor/LLVMAsm.cpp
@@ -18,29 +18,12 @@
 
 #	include "Debug.hpp"
 #	include "llvm/IR/LegacyPassManager.h"
-#	include "llvm/Support/CommandLine.h"
+#	include "llvm/Support/FileSystem.h"
 #	include <fstream>
 #	include <iomanip>
 #	include <regex>
 #	include <sstream>
 
-namespace {
-bool initAsmOutputOptionsOnce()
-{
-	// Use a static immediately invoked lambda to make this thread safe
-	static auto initialized = []() {
-		const char *argv[] = {
-			"Reactor",
-			"-x86-asm-syntax", "intel"  // Use Intel syntax rather than the default AT&T
-		};
-		llvm::cl::ParseCommandLineOptions(sizeof(argv) / sizeof(argv[0]), argv);
-		return true;
-	}();
-
-	return initialized;
-}
-}  // namespace
-
 namespace rr {
 namespace AsmFile {
 
@@ -57,8 +40,6 @@
 
 bool emitAsmFile(const std::string &filename, llvm::orc::JITTargetMachineBuilder builder, llvm::Module &module)
 {
-	initAsmOutputOptionsOnce();
-
 	auto targetMachine = builder.createTargetMachine();
 	if(!targetMachine)
 		return false;
diff --git a/src/Reactor/LLVMJIT.cpp b/src/Reactor/LLVMJIT.cpp
index d334f7d..b6c5d22 100644
--- a/src/Reactor/LLVMJIT.cpp
+++ b/src/Reactor/LLVMJIT.cpp
@@ -29,7 +29,9 @@
 #include "llvm/ExecutionEngine/Orc/IRCompileLayer.h"
 #include "llvm/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.h"
 #include "llvm/ExecutionEngine/SectionMemoryManager.h"
+#include "llvm/IR/DiagnosticInfo.h"
 #include "llvm/IR/LegacyPassManager.h"
+#include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Host.h"
 #include "llvm/Support/TargetSelect.h"
 #include "llvm/Transforms/InstCombine/InstCombine.h"
@@ -106,6 +108,19 @@
 
 namespace {
 
+// TODO(b/174587935): Eliminate command-line parsing.
+bool parseCommandLineOptionsOnce(int argc, const char *const *argv)
+{
+	// Use a static immediately invoked lambda to make this thread safe
+	static auto initialized = [=]() {
+		llvm::cl::ParseCommandLineOptions(argc, argv);
+
+		return true;
+	}();
+
+	return initialized;
+}
+
 // JITGlobals is a singleton that holds all the immutable machine specific
 // information for the host device.
 class JITGlobals
@@ -129,6 +144,14 @@
 JITGlobals *JITGlobals::get()
 {
 	static JITGlobals instance = [] {
+		const char *argv[] = {
+			"Reactor",
+			"-x86-asm-syntax", "intel",   // Use Intel syntax rather than the default AT&T
+			"-warn-stack-size", "524288"  // Warn when a function uses more than 512 KiB of stack memory
+		};
+
+		parseCommandLineOptionsOnce(sizeof(argv) / sizeof(argv[0]), argv);
+
 		llvm::InitializeNativeTarget();
 		llvm::InitializeNativeTargetAsmPrinter();
 		llvm::InitializeNativeTargetAsmParser();
@@ -617,6 +640,40 @@
 	return v;
 }
 
+// Sets *fatal to true if a diagnostic is received which makes a routine invalid or unusable.
+struct FatalDiagnosticsHandler : public llvm::DiagnosticHandler
+{
+	FatalDiagnosticsHandler(bool *fatal)
+	    : fatal(fatal)
+	{}
+
+	bool handleDiagnostics(const llvm::DiagnosticInfo &info) override
+	{
+		switch(info.getSeverity())
+		{
+			case llvm::DS_Error:
+				ASSERT_MSG(false, "LLVM JIT compilation failure");
+				*fatal = true;
+				break;
+			case llvm::DS_Warning:
+				if(info.getKind() == llvm::DK_StackSize)
+				{
+					// Stack size limit exceeded
+					*fatal = true;
+				}
+				break;
+			case llvm::DS_Remark:
+				break;
+			case llvm::DS_Note:
+				break;
+		}
+
+		return true;  // Diagnostic handled, don't let LLVM print it.
+	}
+
+	bool *fatal;
+};
+
 // JITRoutine is a rr::Routine that holds a LLVM JIT session, compiler and
 // object layer as each routine may require different target machine
 // settings and no Reactor routine directly links against another.
@@ -637,6 +694,9 @@
 	    })
 	    , addresses(count)
 	{
+		bool fatalCompileIssue = false;
+		context->setDiagnosticHandler(std::make_unique<FatalDiagnosticsHandler>(&fatalCompileIssue), true);
+
 #ifdef ENABLE_RR_DEBUG_INFO
 		// TODO(b/165000222): Update this on next LLVM roll.
 		// https://github.com/llvm/llvm-project/commit/98f2bb4461072347dcca7d2b1b9571b3a6525801
@@ -693,10 +753,22 @@
 		// Resolve the function addresses.
 		for(size_t i = 0; i < count; i++)
 		{
+			fatalCompileIssue = false;  // May be set to true by session.lookup()
+
+			// This is where the actual compilation happens.
 			auto symbol = session.lookup({ &dylib }, functionNames[i]);
+
 			ASSERT_MSG(symbol, "Failed to lookup address of routine function %d: %s",
 			           (int)i, llvm::toString(symbol.takeError()).c_str());
-			addresses[i] = reinterpret_cast<void *>(static_cast<intptr_t>(symbol->getAddress()));
+
+			if(fatalCompileIssue)
+			{
+				addresses[i] = nullptr;
+			}
+			else  // Successful compilation
+			{
+				addresses[i] = reinterpret_cast<void *>(static_cast<intptr_t>(symbol->getAddress()));
+			}
 		}
 
 #ifdef ENABLE_RR_EMIT_ASM_FILE