[EVM] Support immutable operands in the stackification.

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<std::pair<unsigned, unsigned>> CommutableOpIndexes =
+        std::nullopt;
     bool TerminatesOrReverts = false;
   };
 

llvm/lib/Target/EVM/EVMControlFlowGraphBuilder.cpp

@@ -21,6 +21,7 @@
 #include "MCTargetDesc/EVMMCTargetDesc.h"
 #include "llvm/CodeGen/MachineFunction.h"
 
+#include <optional>
 #include <ostream>
 #include <variant>
 
@@ -204,6 +205,19 @@ void ControlFlowGraphBuilder::collectInstrOperands(const MachineInstr &MI,
     Output.push_back(TemporarySlot{&MI, MO.getReg(), ArgsNumber++});
 }
 
+std::optional<std::pair<unsigned, unsigned>>
+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<unsigned, unsigned>(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<std::pair<unsigned, unsigned>> CommutableOpIndexes =
+        getCommutableOpIndexes(MI);
+    CurrentBlock->Operations.emplace_back(CFG::Operation{
+        std::move(Input), std::move(Output),
+        CFG::BuiltinCall{&MI, CommutableOpIndexes, TerminatesOrReverts}});
   } break;
   }
 

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<std::pair<unsigned, unsigned>>
+  getCommutableOpIndexes(const MachineInstr &MI) const;
   void collectInstrOperands(const MachineInstr &MI, Stack &Input,
                             Stack &Output) const;
 

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<JunkSlot>(TargetStack[Idx]) ||
-           SourceStack[Idx] == TargetStack[Idx]);
+    if (!std::holds_alternative<JunkSlot>(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<int>(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<int>(CurrentStack.size()) == Assembly.getStackHeight());
-    assert(CurrentStack.size() >= Operation.Input.size());
+    if (const auto *Inst = std::get_if<CFG::BuiltinCall>(&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<int>(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<int>(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<int>(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.

llvm/lib/Target/EVM/EVMOptimizedCodeTransform.h

@@ -20,7 +20,6 @@
 #include "llvm/CodeGen/LiveIntervals.h"
 #include "llvm/CodeGen/MachineLoopInfo.h"
 
-#include <optional>
 #include <stack>
 
 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.

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<VariableSlot>(Slot));
-          for (CFG::BasicBlock *Exit : FunctionInfo->Exits) {
-            const Stack &RetValues =
-                std::get<CFG::BasicBlock::FunctionReturn>(Exit->Exit).RetValues;
-
-            for (const StackSlot &Val : RetValues) {
-              if (const VariableSlot *VarSlot = std::get_if<VariableSlot>(&Val))
-                if (*VarSlot == std::get<VariableSlot>(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;

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.