From a7683110d71a15bfc823688191476d4c822565cf Mon Sep 17 00:00:00 2001
From: Arseny Kapoulkine <arseny.kapoulkine@gmail.com>
Date: Fri, 15 Mar 2024 09:49:00 -0700
Subject: [PATCH] CodeGen: Preserve known tags for LOAD_TVALUE synthesized from
 LOADK (#1201)

When lowering LOADK for booleans/numbers/nils, we deconstruct the
operation using STORE_TAG which informs the rest of the optimization
pipeline about the tag of the value. This is helpful to remove various
tag checks.

When the constant is a string or a vector, we just use
LOAD_TVALUE/STORE_TVALUE. For strings, this could be replaced by pointer
load/store, but for vectors there's no great alternative using current
IR ops; in either case, the optimization needs to be carefully examined
for profitability as simply copying constants into registers for
function calls could become more expensive.

However, there are cases where it's still valuable to preserve the tag.
For vectors, doing any math with vector constants contains tag checks
that could be removed. For both strings and vectors, storing them into a
table has a barrier that for vectors could be elided, and for strings
could be simplified as there's no need to confirm the tag.

With this change we now carry the optional tag of the value with
LOAD_TVALUE. This has no performance effect on existing benchmarks but
does reduce the generated code for benchmarks by ~0.1%, and it makes
vector code more efficient (~5% lift on X64 log1p approximation).
---
 CodeGen/include/Luau/IrData.h     | 3 ++-
 CodeGen/src/IrTranslation.cpp     | 8 ++++++++
 CodeGen/src/OptimizeConstProp.cpp | 4 ++++
 3 files changed, 14 insertions(+), 1 deletion(-)

diff --git a/CodeGen/include/Luau/IrData.h b/CodeGen/include/Luau/IrData.h
index 79d06e5a4..a950370b9 100644
--- a/CodeGen/include/Luau/IrData.h
+++ b/CodeGen/include/Luau/IrData.h
@@ -59,7 +59,8 @@ enum class IrCmd : uint8_t
 
     // Load a TValue from memory
     // A: Rn or Kn or pointer (TValue)
-    // B: int (optional 'A' pointer offset)
+    // B: int/none (optional 'A' pointer offset)
+    // C: tag/none (tag of the value being loaded)
     LOAD_TVALUE,
 
     // Load current environment table
diff --git a/CodeGen/src/IrTranslation.cpp b/CodeGen/src/IrTranslation.cpp
index 995225a63..5d55c8772 100644
--- a/CodeGen/src/IrTranslation.cpp
+++ b/CodeGen/src/IrTranslation.cpp
@@ -14,6 +14,7 @@
 
 LUAU_FASTFLAGVARIABLE(LuauCodegenVectorTag2, false)
 LUAU_FASTFLAGVARIABLE(LuauCodegenVectorTag, false)
+LUAU_FASTFLAGVARIABLE(LuauCodegenLoadTVTag, false)
 
 namespace Luau
 {
@@ -111,6 +112,13 @@ static void translateInstLoadConstant(IrBuilder& build, int ra, int k)
         build.inst(IrCmd::STORE_DOUBLE, build.vmReg(ra), build.constDouble(protok.value.n));
         build.inst(IrCmd::STORE_TAG, build.vmReg(ra), build.constTag(LUA_TNUMBER));
     }
+    else if (FFlag::LuauCodegenLoadTVTag)
+    {
+        // Tag could be LUA_TSTRING or LUA_TVECTOR; for TSTRING we could generate LOAD_POINTER/STORE_POINTER/STORE_TAG, but it's not profitable;
+        // however, it's still valuable to preserve the tag throughout the optimization pipeline to eliminate tag checks.
+        IrOp load = build.inst(IrCmd::LOAD_TVALUE, build.vmConst(k), build.constInt(0), build.constTag(protok.tt));
+        build.inst(IrCmd::STORE_TVALUE, build.vmReg(ra), load);
+    }
     else
     {
         // Remaining tag here right now is LUA_TSTRING, while it can be transformed to LOAD_POINTER/STORE_POINTER/STORE_TAG, it's not profitable right
diff --git a/CodeGen/src/OptimizeConstProp.cpp b/CodeGen/src/OptimizeConstProp.cpp
index d765b800a..ff4f7bfc0 100644
--- a/CodeGen/src/OptimizeConstProp.cpp
+++ b/CodeGen/src/OptimizeConstProp.cpp
@@ -19,6 +19,7 @@ LUAU_FASTINTVARIABLE(LuauCodeGenReuseSlotLimit, 64)
 LUAU_FASTFLAGVARIABLE(DebugLuauAbortingChecks, false)
 LUAU_FASTFLAG(LuauCodegenVectorTag2)
 LUAU_DYNAMIC_FASTFLAGVARIABLE(LuauCodeGenCoverForgprepEffect, false)
+LUAU_FASTFLAG(LuauCodegenLoadTVTag)
 
 namespace Luau
 {
@@ -726,6 +727,9 @@ static void constPropInInst(ConstPropState& state, IrBuilder& build, IrFunction&
                             arg->cmd == IrCmd::UNM_VEC)
                             tag = LUA_TVECTOR;
                     }
+
+                    if (FFlag::LuauCodegenLoadTVTag && arg->cmd == IrCmd::LOAD_TVALUE && arg->c.kind != IrOpKind::None)
+                        tag = function.tagOp(arg->c);
                 }
             }