From 80f07411038b0e0452efc7e1cc6adbd2bb232406 Mon Sep 17 00:00:00 2001
From: Nekotekina <nekotekina@gmail.com>
Date: Mon, 29 Aug 2022 03:55:59 +0300
Subject: [PATCH] simd_builder: fix constant locations

---
 Utilities/JIT.cpp | 42 +++++++++++++++++++++++++++---------------
 Utilities/JIT.h   | 17 +++++++++++++----
 2 files changed, 40 insertions(+), 19 deletions(-)

diff --git a/Utilities/JIT.cpp b/Utilities/JIT.cpp
index 6794d111bd..7f59925f75 100644
--- a/Utilities/JIT.cpp
+++ b/Utilities/JIT.cpp
@@ -8,6 +8,7 @@
 #include "util/vm.hpp"
 #include "util/asm.hpp"
 #include "util/v128.hpp"
+#include "util/simd.hpp"
 #include <charconv>
 #include <zlib.h>
 
@@ -357,6 +358,11 @@ asmjit::simd_builder::simd_builder(CodeHolder* ch) noexcept
 	: native_asm(ch)
 {
 	_init(true);
+	consts[~v128()] = this->newLabel();
+}
+
+asmjit::simd_builder::~simd_builder()
+{
 }
 
 void asmjit::simd_builder::_init(bool full)
@@ -402,6 +408,16 @@ void asmjit::simd_builder::_init(bool full)
 	}
 }
 
+void asmjit::simd_builder::operator()() noexcept
+{
+	for (auto&& [x, y] : consts)
+	{
+		this->align(AlignMode::kData, 16);
+		this->bind(y);
+		this->embed(&x, 16);
+	}
+}
+
 void asmjit::simd_builder::vec_cleanup_ret()
 {
 	if (utils::has_avx() && vsize > 16)
@@ -437,23 +453,19 @@ void asmjit::simd_builder::vec_set_const(const Operand& v, const v128& val)
 		return vec_set_all_zeros(v);
 	if (!~val._u)
 		return vec_set_all_ones(v);
-
-	if (uptr(&val) < 0x8000'0000)
-	{
-		// Assume the constant comes from a code or data segment (unsafe)
-		if (x86::Zmm zr(v.id()); zr == v)
-			this->vbroadcasti32x4(zr, x86::oword_ptr(uptr(&val)));
-		else if (x86::Ymm yr(v.id()); yr == v)
-			this->vbroadcasti128(yr, x86::oword_ptr(uptr(&val)));
-		else if (utils::has_avx())
-			this->vmovaps(x86::Xmm(v.id()), x86::oword_ptr(uptr(&val)));
-		else
-			this->movaps(x86::Xmm(v.id()), x86::oword_ptr(uptr(&val)));
-	}
 	else
 	{
-		// TODO
-		fmt::throw_exception("Unexpected constant location");
+		Label co = consts[val];
+		if (!co.isValid())
+			co = consts[val] = this->newLabel();
+		if (x86::Zmm zr(v.id()); zr == v)
+			this->vbroadcasti32x4(zr, x86::oword_ptr(co));
+		else if (x86::Ymm yr(v.id()); yr == v)
+			this->vbroadcasti128(yr, x86::oword_ptr(co));
+		else if (utils::has_avx())
+			this->vmovaps(x86::Xmm(v.id()), x86::oword_ptr(co));
+		else
+			this->movaps(x86::Xmm(v.id()), x86::oword_ptr(co));
 	}
 }
 
diff --git a/Utilities/JIT.h b/Utilities/JIT.h
index dd925a353b..0f616ab6ab 100644
--- a/Utilities/JIT.h
+++ b/Utilities/JIT.h
@@ -42,6 +42,7 @@
 #include <string>
 #include <string_view>
 #include <unordered_map>
+#include <util/v128.hpp>
 
 #if defined(ARCH_X64)
 using native_asm = asmjit::x86::Assembler;
@@ -51,8 +52,6 @@ using native_asm = asmjit::a64::Assembler;
 using native_args = std::array<asmjit::a64::Gp, 4>;
 #endif
 
-union v128;
-
 void jit_announce(uptr func, usz size, std::string_view name);
 
 void jit_announce(auto* func, usz size, std::string_view name)
@@ -215,12 +214,17 @@ namespace asmjit
 #if defined(ARCH_X64)
 	struct simd_builder : native_asm
 	{
+		std::unordered_map<v128, Label> consts;
+
 		Operand v0, v1, v2, v3, v4, v5;
 
 		uint vsize = 16;
 		uint vmask = 0;
 
 		simd_builder(CodeHolder* ch) noexcept;
+		~simd_builder();
+
+		void operator()() noexcept;
 
 		void _init(bool full);
 		void vec_cleanup_ret();
@@ -312,8 +316,7 @@ namespace asmjit
 			if (vmask)
 			{
 				// Build single last iteration (masked)
-				static constexpr u64 all_ones = -1;
-				this->bzhi(reg_cnt, x86::Mem(uptr(&all_ones)), reg_cnt);
+				this->bzhi(reg_cnt, x86::Mem(consts[~u128()], 0), reg_cnt);
 				this->kmovq(x86::k7, reg_cnt);
 				vmask = 7;
 				build();
@@ -427,6 +430,12 @@ inline FT build_function_asm(std::string_view name, F&& builder, ::jit_runtime*
 		builder(compiler, args);
 	}
 
+	if constexpr (std::is_invocable_r_v<void, Asm>)
+	{
+		// Finalization
+		compiler();
+	}
+
 	const auto result = rt._add(&code);
 	jit_announce(result, code.codeSize(), name);
 	return reinterpret_cast<FT>(uptr(result));