Merge pull request #10242 from Nciibi/fix/relinker-inplace-follower-wide-store

fix(relinker): move a following in-place replacement into a 256-bit store's stub

core/relinker/codegen/include/codegen/x86/IAmd64OnlyInstructionMatcher.hpp

@@ -27,7 +27,8 @@ public:
 
     virtual void EmitOutOfLine(
         StubBodyBuilder& body,
-        std::span<const std::uint8_t> instruction
+        std::span<const std::uint8_t> instruction,
+        bool allowInPlace = false
     ) const = 0;
 };
 

core/relinker/codegen/src/Amd64OnlyConverter.cpp

@@ -56,6 +56,7 @@ public:
     ) : _seg(seg), _matches(matches), _ph(ph), _branchTargets(branchTargets), _pending(pending), _claimed(claimed), _matcher(matcher) {}
 
     [[nodiscard]] std::optional<Placement> Window(const std::size_t index) const {
+        std::optional<Placement> fallback;
         for (const bool avoidSegment : {true, false}) {
             if (index + 1 < _matches.size() && (_pending.contains(index + 1) || _movable(index + 1, true, avoidSegment))) {
                 if (auto placement = _relocate(index, index, index + 1))
@@ -65,8 +66,10 @@ public:
                 if (auto placement = _relocate(index, index - 1, index))
                     return placement;
             }
+            if (!fallback && index + 1 < _matches.size() && _pending.contains(index + 1))
+                fallback = _relocate(index, index, index + 1, true);
         }
-        return std::nullopt;
+        return fallback;
     }
 
     [[nodiscard]] std::optional<Placement> Island(const std::size_t index) {
@@ -130,7 +133,7 @@ private:
         };
     }
 
-    [[nodiscard]] std::optional<Placement> _relocate(const std::size_t index, const std::size_t first, const std::size_t last) const {
+    [[nodiscard]] std::optional<Placement> _relocate(const std::size_t index, const std::size_t first, const std::size_t last, const bool allowInPlace = false) const {
         if (_entered(_address(first), _end(last)))
             return std::nullopt;
         StubBodyBuilder body;
@@ -138,7 +141,7 @@ private:
         try {
             for (auto k = first; k <= last; ++k) {
                 if (_pending.contains(k)) {
-                    _matcher.EmitOutOfLine(body, _bytes(k));
+                    _matcher.EmitOutOfLine(body, _bytes(k), allowInPlace);
                     if (k != index)
                         placement.Consumed.push_back(k);
                 } else {

core/relinker/codegen/src/x86/Amd64OnlyInstructionMatcher.cpp

@@ -172,7 +172,8 @@ public:
 
     void EmitOutOfLine(
         StubBodyBuilder& body,
-        std::span<const std::uint8_t> instruction
+        std::span<const std::uint8_t> instruction,
+        bool allowInPlace = false
     ) const override;
 
 private:
@@ -355,8 +356,12 @@ const char* Amd64OnlyInstructionMatcher::_emitInPlace(StubBodyBuilder& body, con
     return nullptr;
 }
 
-void Amd64OnlyInstructionMatcher::EmitOutOfLine(StubBodyBuilder& body, const std::span<const std::uint8_t> instruction) const {
-    if (_emitOutOfLine(body, DecodedInstruction{instruction.data(), instruction.size()}) == nullptr)
+void Amd64OnlyInstructionMatcher::EmitOutOfLine(StubBodyBuilder& body, const std::span<const std::uint8_t> instruction, const bool allowInPlace) const {
+    const DecodedInstruction instr{instruction.data(), instruction.size()};
+    const auto* emitted = _emitOutOfLine(body, instr);
+    if (emitted == nullptr && allowInPlace)
+        emitted = _emitInPlace(body, instr);
+    if (emitted == nullptr)
         throw CodegenException("AMD-only instruction has no out-of-line lowering");
     body.Advance(instruction.size());
 }

core/relinker/codegen/tests/Amd64OnlyConverterTests.cpp

@@ -837,6 +837,21 @@ void converterWideStoreWindows() {
     expected.insert(expected.end(), {0xC5, 0xF8, 0x11, 0x0E, 0xC4, 0xE3, 0x7D, 0x19, 0x4E, 0x10, 0x01});
     require(paired.Trampolines.size() == 1 && paired.Trampolines[0].Length == 8 && body(paired.Trampolines[0]) == expected && paired.Reports.size() == 1, "A 256-bit store followed by another was not lowered with it in one stub");
 
+    const auto withRdseed = convert(withStore({}, {0x0F, 0xC7, 0xF8, 0xC3}));
+    expected = halves;
+    expected.insert(expected.end(), {0x0F, 0xC7, 0xF0});
+    require(withRdseed.Trampolines.size() == 1 && withRdseed.Trampolines[0].Length == 7 && body(withRdseed.Trampolines[0]) == expected && withRdseed.Reports.size() == 1, "A 256-bit store followed by an in-place site did not move its replacement into the stub");
+
+    const auto withClwb = convert(withStore({}, {0x66, 0x0F, 0xAE, 0x37, 0xC3}));
+    expected = halves;
+    expected.insert(expected.end(), {0x0F, 0x1F, 0x40, 0x00});
+    require(withClwb.Trampolines.size() == 1 && withClwb.Trampolines[0].Length == 8 && body(withClwb.Trampolines[0]) == expected, "A 256-bit store followed by CLWB did not move its replacement into the stub");
+
+    const auto withRipStore = convert(withStore({}, {0xF3, 0x0F, 0x2B, 0x05, 0x40, 0x00, 0x00, 0x00, 0xC3}));
+    expected = halves;
+    expected.insert(expected.end(), {0xF3, 0x0F, 0x11, 0x05, 0x00, 0x00, 0x00, 0x00});
+    require(withRipStore.Trampolines.size() == 1 && withRipStore.Trampolines[0].Length == 12 && body(withRipStore.Trampolines[0]) == expected && withRipStore.Trampolines[0].Relocations.size() == 1 && withRipStore.Trampolines[0].Relocations[0].SiteTarget == 0x4C && withRipStore.Trampolines[0].Relocations[0].InstructionEnd == 19, "A 256-bit store followed by a RIP-relative in-place site did not relocate its replacement");
+
     const Bytes push = {0x48, 0x8D, 0x64, 0x24, 0xF8, 0x50, 0x48, 0x8D, 0x05, 0, 0, 0, 0, 0x48, 0x89, 0x44, 0x24, 0x08, 0x58};
     const auto ripCalled = convert(withStore({}, {0xFF, 0x15, 0x10, 0x00, 0x00, 0x00}));
     expected = halves;