From 55d04ca26edea2be55bfd9ca60f4232e5e8cf181 Mon Sep 17 00:00:00 2001 From: Pavel Kopyl Date: Wed, 30 Oct 2024 15:07:41 +0100 Subject: [PATCH] [EVM] Support immutable operands in the stackification. --- llvm/lib/Target/EVM/EVMControlFlowGraph.h | 6 ++ .../Target/EVM/EVMControlFlowGraphBuilder.cpp | 22 ++++- .../Target/EVM/EVMControlFlowGraphBuilder.h | 3 +- .../Target/EVM/EVMOptimizedCodeTransform.cpp | 97 ++++++++++++++----- .../Target/EVM/EVMOptimizedCodeTransform.h | 9 +- .../Target/EVM/EVMStackLayoutGenerator.cpp | 91 +++++++---------- llvm/lib/Target/EVM/EVMStackLayoutGenerator.h | 4 + 7 files changed, 143 insertions(+), 89 deletions(-) diff --git a/llvm/lib/Target/EVM/EVMControlFlowGraph.h b/llvm/lib/Target/EVM/EVMControlFlowGraph.h index 86bfb57eff5b..ebebcb194184 100644 --- a/llvm/lib/Target/EVM/EVMControlFlowGraph.h +++ b/llvm/lib/Target/EVM/EVMControlFlowGraph.h @@ -168,6 +168,12 @@ struct CFG { struct BuiltinCall { MachineInstr *Builtin = nullptr; + // This contains commutable operand indexes, if any. + // Operand indexes correspond to a stack layout, i.e., the first use + // operand in the instruction definition is located on the stack top + // and has zero index. + std::optional> CommutableOpIndexes = + std::nullopt; bool TerminatesOrReverts = false; }; diff --git a/llvm/lib/Target/EVM/EVMControlFlowGraphBuilder.cpp b/llvm/lib/Target/EVM/EVMControlFlowGraphBuilder.cpp index f84628db7af5..927b00cf4d3d 100644 --- a/llvm/lib/Target/EVM/EVMControlFlowGraphBuilder.cpp +++ b/llvm/lib/Target/EVM/EVMControlFlowGraphBuilder.cpp @@ -21,6 +21,7 @@ #include "MCTargetDesc/EVMMCTargetDesc.h" #include "llvm/CodeGen/MachineFunction.h" +#include #include #include @@ -204,6 +205,19 @@ void ControlFlowGraphBuilder::collectInstrOperands(const MachineInstr &MI, Output.push_back(TemporarySlot{&MI, MO.getReg(), ArgsNumber++}); } +std::optional> +ControlFlowGraphBuilder::getCommutableOpIndexes(const MachineInstr &MI) const { + unsigned Opc = MI.getOpcode(); + if (MI.isCommutable() || Opc == EVM::ADDMOD || Opc == EVM::MULMOD) + // Operand indexes correspond to the stack layout. In all the cases, + // commutable operands take two top stack slots. Even for ADDMOD and + // MULMOD instructions, because their MOD operand takes third slot from the + // TOS. + return std::pair(0, 1); + + return std::nullopt; +} + void ControlFlowGraphBuilder::handleMachineInstr(MachineInstr &MI) { bool TerminatesOrReverts = false; unsigned Opc = MI.getOpcode(); @@ -252,9 +266,11 @@ void ControlFlowGraphBuilder::handleMachineInstr(MachineInstr &MI) { default: { Stack Input, Output; collectInstrOperands(MI, Input, Output); - CurrentBlock->Operations.emplace_back( - CFG::Operation{std::move(Input), std::move(Output), - CFG::BuiltinCall{&MI, TerminatesOrReverts}}); + std::optional> CommutableOpIndexes = + getCommutableOpIndexes(MI); + CurrentBlock->Operations.emplace_back(CFG::Operation{ + std::move(Input), std::move(Output), + CFG::BuiltinCall{&MI, CommutableOpIndexes, TerminatesOrReverts}}); } break; } diff --git a/llvm/lib/Target/EVM/EVMControlFlowGraphBuilder.h b/llvm/lib/Target/EVM/EVMControlFlowGraphBuilder.h index 151b5253af4b..3c305dedcac5 100644 --- a/llvm/lib/Target/EVM/EVMControlFlowGraphBuilder.h +++ b/llvm/lib/Target/EVM/EVMControlFlowGraphBuilder.h @@ -42,7 +42,8 @@ class ControlFlowGraphBuilder { void handleReturn(const MachineInstr &MI); void handleBasicBlockSuccessors(MachineBasicBlock &MBB); StackSlot getDefiningSlot(const MachineInstr &MI, Register Reg) const; - + std::optional> + getCommutableOpIndexes(const MachineInstr &MI) const; void collectInstrOperands(const MachineInstr &MI, Stack &Input, Stack &Output) const; diff --git a/llvm/lib/Target/EVM/EVMOptimizedCodeTransform.cpp b/llvm/lib/Target/EVM/EVMOptimizedCodeTransform.cpp index 0e38d53f96f0..a90534aff7ba 100644 --- a/llvm/lib/Target/EVM/EVMOptimizedCodeTransform.cpp +++ b/llvm/lib/Target/EVM/EVMOptimizedCodeTransform.cpp @@ -122,12 +122,17 @@ EVMOptimizedCodeTransform::EVMOptimizedCodeTransform(EVMAssembly &Assembly, MachineFunction &MF) : Assembly(Assembly), Layout(Layout), FuncInfo(&Cfg.FuncInfo), MF(MF) {} -void EVMOptimizedCodeTransform::assertLayoutCompatibility( - Stack const &SourceStack, Stack const &TargetStack) { - assert(SourceStack.size() == TargetStack.size()); +bool EVMOptimizedCodeTransform::AreLayoutsCompatible(Stack const &SourceStack, + Stack const &TargetStack) { + if (SourceStack.size() != TargetStack.size()) + return false; + for (unsigned Idx = 0; Idx < SourceStack.size(); ++Idx) - assert(std::holds_alternative(TargetStack[Idx]) || - SourceStack[Idx] == TargetStack[Idx]); + if (!std::holds_alternative(TargetStack[Idx]) && + !(SourceStack[Idx] == TargetStack[Idx])) + return false; + + return true; } void EVMOptimizedCodeTransform::createStackLayout(Stack TargetStack) { @@ -248,7 +253,7 @@ void EVMOptimizedCodeTransform::createStackLayout(Stack TargetStack) { assert(Assembly.getStackHeight() == static_cast(CurrentStack.size())); } -void EVMOptimizedCodeTransform::operator()(CFG::BasicBlock const &Block) { +void EVMOptimizedCodeTransform::operator()(const CFG::BasicBlock &Block) { // Current location for the entry BB was set up in operator()(). if (&Block != FuncInfo->Entry) Assembly.setCurrentLocation(Block.MBB); @@ -260,7 +265,7 @@ void EVMOptimizedCodeTransform::operator()(CFG::BasicBlock const &Block) { auto const &BlockInfo = Layout.blockInfos.at(&Block); // Assert that the stack is valid for entering the block. - assertLayoutCompatibility(CurrentStack, BlockInfo.entryLayout); + assert(AreLayoutsCompatible(CurrentStack, BlockInfo.entryLayout)); // Might set some slots to junk, if not required by the block. CurrentStack = BlockInfo.entryLayout; @@ -270,17 +275,63 @@ void EVMOptimizedCodeTransform::operator()(CFG::BasicBlock const &Block) { if (EVMUtils::valueOrNullptr(BlockLabels, &Block)) Assembly.appendLabel(); - for (auto const &Operation : Block.Operations) { + for (const auto &Operation : Block.Operations) { + // Check if we can choose cheaper stack shuffling if the Operation is an + // instruction with commutable arguments. // Create required layout for entering the Operation. - createStackLayout(Layout.operationEntryLayout.at(&Operation)); - - // Assert that we have the inputs of the Operation on stack top. - assert(static_cast(CurrentStack.size()) == Assembly.getStackHeight()); - assert(CurrentStack.size() >= Operation.Input.size()); + if (const auto *Inst = std::get_if(&Operation.Operation); + Inst && Inst->CommutableOpIndexes) { + // Get the stack layout before the instruction. + const Stack &DefaultTargetStack = + Layout.operationEntryLayout.at(&Operation); + size_t DefaultCost = + EvaluateStackTransform(CurrentStack, DefaultTargetStack); + unsigned OpIdx1 = Inst->CommutableOpIndexes->first; + unsigned OpIdx2 = Inst->CommutableOpIndexes->second; + // Swap the commutable stack items and measure the stack shuffling cost + // again. + assert(DefaultTargetStack.size() > OpIdx1 && + DefaultTargetStack.size() > OpIdx2); + Stack CommutedTargetStack = DefaultTargetStack; + std::swap(CommutedTargetStack[CommutedTargetStack.size() - OpIdx1 - 1], + CommutedTargetStack[CommutedTargetStack.size() - OpIdx2 - 1]); + size_t CommutedCost = + EvaluateStackTransform(CurrentStack, CommutedTargetStack); + // Choose the cheapest transformation. + createStackLayout(CommutedCost < DefaultCost ? CommutedTargetStack + : DefaultTargetStack); + +#ifndef NDEBUG + // Assert that we have the inputs of the Operation on stack top. + assert(static_cast(CurrentStack.size()) == + Assembly.getStackHeight()); + assert(CurrentStack.size() >= Operation.Input.size()); + Stack StackInput = EVMUtils::to_vector( + EVMUtils::take_last(CurrentStack, Operation.Input.size())); + // Adjust the StackInput to match the commuted stack. + if (CommutedCost < DefaultCost) { + std::swap(StackInput[StackInput.size() - OpIdx1 - 1], + StackInput[StackInput.size() - OpIdx2 - 1]); + } + assert(AreLayoutsCompatible(StackInput, Operation.Input)); +#endif // NDEBUG + } else { + createStackLayout(Layout.operationEntryLayout.at(&Operation)); + +#ifndef NDEBUG + // Assert that we have the inputs of the Operation on stack top. + assert(static_cast(CurrentStack.size()) == + Assembly.getStackHeight()); + assert(CurrentStack.size() >= Operation.Input.size()); + const Stack StackInput = EVMUtils::to_vector( + EVMUtils::take_last(CurrentStack, Operation.Input.size())); + assert(AreLayoutsCompatible(StackInput, Operation.Input)); +#endif // NDEBUG + } + +#ifndef NDEBUG size_t BaseHeight = CurrentStack.size() - Operation.Input.size(); - assertLayoutCompatibility(EVMUtils::to_vector(EVMUtils::take_last( - CurrentStack, Operation.Input.size())), - Operation.Input); +#endif // NDEBUG // Perform the Operation. std::visit(*this, Operation.Operation); @@ -289,9 +340,9 @@ void EVMOptimizedCodeTransform::operator()(CFG::BasicBlock const &Block) { assert(static_cast(CurrentStack.size()) == Assembly.getStackHeight()); assert(CurrentStack.size() == BaseHeight + Operation.Output.size()); assert(CurrentStack.size() >= Operation.Output.size()); - assertLayoutCompatibility(EVMUtils::to_vector(EVMUtils::take_last( - CurrentStack, Operation.Output.size())), - Operation.Output); + assert(AreLayoutsCompatible(EVMUtils::to_vector(EVMUtils::take_last( + CurrentStack, Operation.Output.size())), + Operation.Output)); } // Exit the block. @@ -342,11 +393,11 @@ void EVMOptimizedCodeTransform::operator()(CFG::BasicBlock const &Block) { CurrentStack.pop_back(); // Assert that we have a valid stack for both jump targets. - assertLayoutCompatibility( + assert(AreLayoutsCompatible( CurrentStack, - Layout.blockInfos.at(CondJump.NonZero).entryLayout); - assertLayoutCompatibility( - CurrentStack, Layout.blockInfos.at(CondJump.Zero).entryLayout); + Layout.blockInfos.at(CondJump.NonZero).entryLayout)); + assert(AreLayoutsCompatible( + CurrentStack, Layout.blockInfos.at(CondJump.Zero).entryLayout)); { // Restore the stack afterwards for the non-zero case below. diff --git a/llvm/lib/Target/EVM/EVMOptimizedCodeTransform.h b/llvm/lib/Target/EVM/EVMOptimizedCodeTransform.h index dce8819454bc..11abab698068 100644 --- a/llvm/lib/Target/EVM/EVMOptimizedCodeTransform.h +++ b/llvm/lib/Target/EVM/EVMOptimizedCodeTransform.h @@ -20,7 +20,6 @@ #include "llvm/CodeGen/LiveIntervals.h" #include "llvm/CodeGen/MachineLoopInfo.h" -#include #include namespace llvm { @@ -47,11 +46,11 @@ class EVMOptimizedCodeTransform { EVMOptimizedCodeTransform(EVMAssembly &Assembly, const CFG &Cfg, const StackLayout &Layout, MachineFunction &MF); - /// Assert that it is valid to transition from \p SourceStack to \p - /// TargetStack. That is \p SourceStack matches each slot in \p + /// Checks if it's valid to transition from \p SourceStack to \p + /// TargetStack, that is \p SourceStack matches each slot in \p /// TargetStack that is not a JunkSlot exactly. - static void assertLayoutCompatibility(Stack const &SourceStack, - Stack const &TargetStack); + static bool AreLayoutsCompatible(Stack const &SourceStack, + Stack const &TargetStack); /// Shuffles CurrentStack to the desired \p TargetStack while emitting the /// shuffling code to Assembly. diff --git a/llvm/lib/Target/EVM/EVMStackLayoutGenerator.cpp b/llvm/lib/Target/EVM/EVMStackLayoutGenerator.cpp index c67b76c297e7..ad2ac68a0902 100644 --- a/llvm/lib/Target/EVM/EVMStackLayoutGenerator.cpp +++ b/llvm/lib/Target/EVM/EVMStackLayoutGenerator.cpp @@ -698,6 +698,37 @@ Stack StackLayoutGenerator::compressStack(Stack CurStack) { return CurStack; } +/// Returns the number of operations required to transform stack \p Source to +/// \p Target. +size_t llvm::EvaluateStackTransform(Stack Source, Stack const &Target) { + size_t OpGas = 0; + auto Swap = [&](unsigned SwapDepth) { + if (SwapDepth > 16) + OpGas += 1000; + else + OpGas += 3; // SWAP* gas price; + }; + + auto DupOrPush = [&](StackSlot const &Slot) { + if (canBeFreelyGenerated(Slot)) + OpGas += 3; + else { + auto Depth = EVMUtils::findOffset(EVMUtils::get_reverse(Source), Slot); + if (!Depth) + llvm_unreachable("No slot in the stack"); + + if (*Depth < 16) + OpGas += 3; // DUP* gas price + else + OpGas += 1000; + } + }; + auto Pop = [&]() { OpGas += 2; }; + + createStackLayout(Source, Target, Swap, DupOrPush, Pop); + return OpGas; +} + void StackLayoutGenerator::fillInJunk(CFG::BasicBlock const &Block, CFG::FunctionInfo const *FunctionInfo) { /// Recursively adds junk to the subgraph starting on \p Entry. @@ -739,70 +770,16 @@ void StackLayoutGenerator::fillInJunk(CFG::BasicBlock const &Block, }); }; - /// Returns the number of operations required to transform \p Source to \p - /// Target. - auto EvaluateTransform = [&](Stack Source, Stack const &Target) -> size_t { - size_t OpGas = 0; - auto Swap = [&](unsigned SwapDepth) { - if (SwapDepth > 16) - OpGas += 1000; - else - OpGas += 3; // SWAP* gas price; - }; - - auto DupOrPush = [&](StackSlot const &Slot) { - if (canBeFreelyGenerated(Slot)) - OpGas += 3; - else { - if (auto Depth = - EVMUtils::findOffset(EVMUtils::get_reverse(Source), Slot)) { - if (*Depth < 16) - OpGas += 3; // gas price for DUP - else - OpGas += 1000; - } else { - // This has to be a previously unassigned return variable. - // We at least sanity-check that it is among the return variables at - // all. -#ifndef NDEBUG - bool VarExists = false; - assert(std::holds_alternative(Slot)); - for (CFG::BasicBlock *Exit : FunctionInfo->Exits) { - const Stack &RetValues = - std::get(Exit->Exit).RetValues; - - for (const StackSlot &Val : RetValues) { - if (const VariableSlot *VarSlot = std::get_if(&Val)) - if (*VarSlot == std::get(Slot)) - VarExists = true; - } - } - assert(VarExists); -#endif // NDEBUG - // Strictly speaking the cost of the - // PUSH0 depends on the targeted EVM version, but the difference will - // not matter here. - OpGas += 2; - } - } - }; - - auto Pop = [&]() { OpGas += 2; }; - - createStackLayout(Source, Target, Swap, DupOrPush, Pop); - return OpGas; - }; - /// Returns the number of junk slots to be prepended to \p TargetLayout for /// an optimal transition from \p EntryLayout to \p TargetLayout. auto GetBestNumJunk = [&](Stack const &EntryLayout, Stack const &TargetLayout) -> size_t { - size_t BestCost = EvaluateTransform(EntryLayout, TargetLayout); + size_t BestCost = EvaluateStackTransform(EntryLayout, TargetLayout); size_t BestNumJunk = 0; size_t MaxJunk = EntryLayout.size(); for (size_t NumJunk = 1; NumJunk <= MaxJunk; ++NumJunk) { - size_t Cost = EvaluateTransform(EntryLayout, Stack{NumJunk, JunkSlot{}} + - TargetLayout); + size_t Cost = EvaluateStackTransform( + EntryLayout, Stack{NumJunk, JunkSlot{}} + TargetLayout); if (Cost < BestCost) { BestCost = Cost; BestNumJunk = NumJunk; diff --git a/llvm/lib/Target/EVM/EVMStackLayoutGenerator.h b/llvm/lib/Target/EVM/EVMStackLayoutGenerator.h index 75b2beb233c7..e45a23fc6235 100644 --- a/llvm/lib/Target/EVM/EVMStackLayoutGenerator.h +++ b/llvm/lib/Target/EVM/EVMStackLayoutGenerator.h @@ -24,6 +24,10 @@ namespace llvm { +/// Returns the number of operations required to transform stack \p Source to +/// \p Target. +size_t EvaluateStackTransform(Stack Source, Stack const &Target); + struct StackLayout { struct BlockInfo { /// Complete stack layout that is required for entering a block.