Merge pull request #16767 from argotorg/ssa-cfg-kiss-memspill

SSA CFG: stack to memory spilling

Author: clonker

libyul/CMakeLists.txt

@@ -115,6 +115,10 @@ add_library(yul
 	backends/evm/ssa/io/JSONExporter.h
 	backends/evm/ssa/io/Printer.cpp
 	backends/evm/ssa/io/Printer.h
+	backends/evm/ssa/spill/Emitter.h
+	backends/evm/ssa/spill/MemoryAddressing.cpp
+	backends/evm/ssa/spill/MemoryAddressing.h
+	backends/evm/ssa/spill/SpillSet.cpp
 	backends/evm/ssa/spill/SpillSet.h
 	backends/evm/ssa/transform/IdentityAndNopRemover.cpp
 	backends/evm/ssa/transform/IdentityAndNopRemover.h

libyul/backends/evm/EVMObjectCompiler.cpp

@@ -96,6 +96,7 @@ void EVMObjectCompiler::run(Object const& _object, bool _optimize, bool _viaSSAC
 			ssa::ControlFlowGraphsLiveness const liveness(*controlFlowGraphs);
 			ssa::CodeTransform::run(
 				m_assembly,
+				*controlFlowGraphs,
 				liveness,
 				context
 			);

libyul/backends/evm/ssa/CodeTransform.cpp

@@ -45,36 +45,66 @@ void assertLayoutCompatibility(StackData const& _layout1, StackData const& _layo
 void CodeTransform::run
 (
 	AbstractAssembly& _assembly,
+	ControlFlowGraphs& _controlFlowGraphs,
 	ControlFlowGraphsLiveness const& _controlFlowLiveness,
 	BuiltinContext& _builtinContext
 )
 {
+	yulAssert(&_controlFlowLiveness.controlFlowGraphs.get() == &_controlFlowGraphs);
 	yulAssert(!_controlFlowLiveness.cfgLiveness.empty());
-	ControlFlowGraphs const& controlFlowGraphs = _controlFlowLiveness.controlFlowGraphs.get();
-	yulAssert(controlFlowGraphs.functionGraphs.size() == _controlFlowLiveness.cfgLiveness.size());
-	FunctionLabels const functionLabels = registerFunctionLabels(_assembly, controlFlowGraphs);
-	CallGraph const callGraph(controlFlowGraphs);
-
-	for (std::size_t functionIndex = 0; functionIndex < controlFlowGraphs.functionGraphs.size(); ++functionIndex)
+	yulAssert(_controlFlowGraphs.functionGraphs.size() == _controlFlowLiveness.cfgLiveness.size());
+	FunctionLabels const functionLabels = registerFunctionLabels(_assembly, _controlFlowGraphs);
+	CallGraph const callGraph(_controlFlowGraphs);
+
+	std::size_t const numCFGs = _controlFlowGraphs.functionGraphs.size();
+	std::vector<CallSites> callSitesPerCFG;
+	std::vector<SSACFGStackLayout> layouts;
+	std::vector<spill::SpillSet> spillSetsPerCFG;
+	callSitesPerCFG.reserve(numCFGs);
+	layouts.reserve(numCFGs);
+	spillSetsPerCFG.reserve(numCFGs);
+
+	for (std::size_t functionIndex = 0; functionIndex < numCFGs; ++functionIndex)
 	{
-		std::unique_ptr<SSACFG> const& functionGraphPtr = controlFlowGraphs.functionGraphs[functionIndex];
+		std::unique_ptr<SSACFG> const& functionGraphPtr = _controlFlowGraphs.functionGraphs[functionIndex];
 		yulAssert(functionGraphPtr);
 		SSACFG const& cfg = *functionGraphPtr;
-		auto const callSites = gatherCallSites(cfg);
 		auto const& liveness = _controlFlowLiveness.cfgLiveness[functionIndex];
 		yulAssert(liveness);
 		auto const graphID = static_cast<ControlFlowGraphs::FunctionGraphID>(functionIndex);
+		callSitesPerCFG.push_back(gatherCallSites(cfg));
 		bool const spillingAllowed = !callGraph.isRecursive(graphID);
-		auto const stackLayoutGeneratorResult = StackLayoutGenerator::generate(*liveness, callSites, graphID, spillingAllowed);
+		auto [layout, spillSet] = StackLayoutGenerator::generate(*liveness, callSitesPerCFG.back(), graphID, spillingAllowed);
+		layouts.push_back(std::move(layout));
+		spillSetsPerCFG.push_back(std::move(spillSet));
+	}
+
+	for (std::size_t functionIndex = 0; functionIndex < numCFGs; ++functionIndex)
+		spillSetsPerCFG[functionIndex].closeUnderReachabilityConstraints(
+			*_controlFlowGraphs.functionGraphs[functionIndex],
+			layouts[functionIndex]
+		);
+
+	// build up global addressing based on the spill sets
+	spill::MemoryAddressing const addressing(_controlFlowGraphs, spillSetsPerCFG);
+
+	for (std::size_t functionIndex = 0; functionIndex < _controlFlowGraphs.functionGraphs.size(); ++functionIndex)
+	{
+		SSACFG const& cfg = *_controlFlowGraphs.functionGraphs[functionIndex];
+		auto const graphID = static_cast<ControlFlowGraphs::FunctionGraphID>(functionIndex);
+		auto const& liveness = _controlFlowLiveness.cfgLiveness[functionIndex];
+		yulAssert(liveness);
 		CodeTransform transform(
 			_assembly,
 			_builtinContext,
-			controlFlowGraphs,
+			_controlFlowGraphs,
 			functionLabels,
-			callSites,
+			callSitesPerCFG[functionIndex],
 			cfg,
-			stackLayoutGeneratorResult.layout,
-			graphID
+			layouts[functionIndex],
+			spillSetsPerCFG[functionIndex],
+			graphID,
+			addressing
 		);
 		transform(cfg.entry);
 	}
@@ -120,7 +150,9 @@ CodeTransform::CodeTransform(
 	CallSites const& _callSites,
 	SSACFG const& _cfg,
 	SSACFGStackLayout const& _stackLayout,
-	ControlFlowGraphs::FunctionGraphID _graphID
+	spill::SpillSet const& _spillSet,
+	ControlFlowGraphs::FunctionGraphID _graphID,
+	spill::MemoryAddressing const& _addressing
 ):
 	m_assembly(_assembly),
 	m_builtinContext(_builtinContext),
@@ -129,6 +161,7 @@ CodeTransform::CodeTransform(
 	m_callSites(_callSites),
 	m_cfg(_cfg),
 	m_stackLayout(_stackLayout),
+	m_spillSet(_spillSet),
 	m_graphID(_graphID),
 	m_blockIsTransformed(_cfg.numBlocks(), false),
 	m_blockLabels([this] {
@@ -138,11 +171,17 @@ CodeTransform::CodeTransform(
 			blockLabels.push_back(m_assembly.newLabelId());
 		return blockLabels;
 	}()),
+	m_spillEmitter([&]() -> std::optional<spill::Emitter> {
+		if (_spillSet.numSpilled() > 0)
+			return spill::Emitter(_addressing, _graphID, _assembly);
+		return std::nullopt;
+	}()),
 	m_assemblyCallbacks{
 		.cfg = &_cfg,
 		.assembly = &_assembly,
 		.callSites = &_callSites,
-		.returnLabels = &m_returnLabels
+		.returnLabels = &m_returnLabels,
+		.spillEmitter = m_spillEmitter ? &*m_spillEmitter : nullptr
 	},
 	m_stackData([&]
 	{
@@ -167,6 +206,11 @@ CodeTransform::CodeTransform(
 	for (auto const& arg: m_cfg.arguments | ranges::views::reverse)
 		expectedStackTop.push_back(StackSlot::makeValue(_cfg, arg));
 	assertLayoutCompatibility(m_stack.data(), expectedStackTop);
+
+	// Spilled function args need an `mstore` at function entry so later `mload`s see a populated slot
+	if (m_spillEmitter && isFunctionGraph)
+		for (InstId const argId: m_cfg.arguments)
+			spillStore(argId);
 }
 
 void CodeTransform::operator()(SSACFG::BlockId const _blockId)
@@ -184,18 +228,28 @@ void CodeTransform::operator()(SSACFG::BlockId const _blockId)
 	auto const& block = m_cfg.block(_blockId);
 
 	std::size_t operationIndex = 0;
-	m_cfg.forEachOperation(block, [&](InstId const instId, SSACFG::Inst const&) {
-		yulAssert(operationIndex < blockLayout->operationIn.size());
-		auto const& operationInLayout = blockLayout->operationIn[operationIndex];
-		(*this)(instId, operationInLayout);
-		++operationIndex;
-	});
+
+	// Iterate every Inst in the block in scheduled order. Only Operations advance codegen;
+	// Phis are otherwise pure stack assertions (already materialized on the block's stackIn).
+	for (InstId const instId: block.instructions)
+	{
+		SSACFG::Inst const& inst = m_cfg.inst(instId);
+		if (inst.isPhi())
+			// this is a no-op for not spilled phis
+			spillStore(instId);
+		if (inst.isOperation())
+		{
+			yulAssert(operationIndex < blockLayout->operationIn.size());
+			(*this)(instId, blockLayout->operationIn[operationIndex]);
+			++operationIndex;
+		}
+	}
 	yulAssert(operationIndex == blockLayout->operationIn.size());
 
 	// Shuffle to the block's exit layout before dispatching the exit.
 	// This ensures the condition is on top for ConditionalJump, phi pre-images are
 	// in the right positions for jumps, and return values are accessible for FunctionReturn.
-	auto const shuffleResult = StackShuffler<AssemblyCallbacks>::shuffle(m_stack, blockLayout->exitIn);
+	auto const shuffleResult = StackShuffler<AssemblyCallbacks>::shuffle(m_stack, blockLayout->exitIn, &m_spillSet);
 	yulAssert(shuffleResult.status == StackShufflerResult::Status::Admissible);
 
 	// handle the block exit
@@ -221,7 +275,7 @@ void CodeTransform::operator()(InstId _instId, StackData const& _operationInputL
 
 	// prepare stack for operation
 	{
-		auto const shuffleResult = StackShuffler<AssemblyCallbacks>::shuffle(m_stack, _operationInputLayout);
+		auto const shuffleResult = StackShuffler<AssemblyCallbacks>::shuffle(m_stack, _operationInputLayout, &m_spillSet);
 		yulAssert(shuffleResult.status == StackShufflerResult::Status::Admissible);
 	}
 
@@ -322,6 +376,11 @@ void CodeTransform::operator()(InstId _instId, StackData const& _operationInputL
 	for (InstId const id: m_cfg.outputsOf(_instId))
 		m_stack.push<false>(StackSlot::makeValue(m_cfg, id));
 
+	// Each output the layout decided to spill gets its `mstore` here
+	if (m_spillEmitter)
+		for (InstId const outputId: m_cfg.outputsOf(_instId))
+			spillStore(outputId);
+
 	yulAssert(m_stack.size() == baseHeight + numOutputs);
 	for (auto const& [stackEntry, output]: ranges::views::zip(
 		m_stack.data() | ranges::views::take_last(numOutputs),
@@ -334,6 +393,32 @@ void CodeTransform::operator()(InstId _instId, StackData const& _operationInputL
 	);
 }
 
+void CodeTransform::spillStore(InstId const _value)
+{
+	if (!m_spillEmitter || !m_spillSet.isSpilled(_value))
+		return;
+
+	// Bring `_value` to the stack top so that the `mstore` below consumes it
+	StackData target = m_stack.data();
+	target.push_back(StackSlot::makeValue(m_cfg, _value));
+	// addr(_value) is the slot THIS `mstore` populates, so we have to exclude it from the spill set to
+	// prevent the shuffler from just `mload`ing it back
+	spill::SpillSet const spillSetExcludingSpillee = m_spillSet.without(_value);
+	auto const shuffleResult = StackShuffler<AssemblyCallbacks>::shuffle(m_stack, target, &spillSetExcludingSpillee);
+	yulAssert(
+		shuffleResult.status == StackShufflerResult::Status::Admissible,
+		fmt::format(
+			"shuffler failed to bring spilled value {} to top (status={})",
+			_value,
+			static_cast<int>(shuffleResult.status)
+		)
+	);
+
+	m_spillEmitter->emitStore(_value);
+	// `mstore` consumed the value; the address it pushed never persists on the symbolic stack
+	m_stack.pop<false>();
+}
+
 void CodeTransform::operator()(SSACFG::BlockId const&, SSACFG::BasicBlock::MainExit const&)
 {
 	yulAssert(static_cast<int>(m_stack.size()) == m_assembly.stackHeight());
@@ -446,7 +531,7 @@ void CodeTransform::prepareBlockExitStack(StackData const& _target, PhiInverse c
 	auto const pulledBackTarget = stackPreImage(m_cfg, _target, _phiInverse);
 	// shuffle to target
 	{
-		auto const shuffleResult = StackShuffler<AssemblyCallbacks>::shuffle(m_stack, pulledBackTarget);
+		auto const shuffleResult = StackShuffler<AssemblyCallbacks>::shuffle(m_stack, pulledBackTarget, &m_spillSet);
 		yulAssert(shuffleResult.status == StackShufflerResult::Status::Admissible);
 	}
 	// check that shuffling was successful

libyul/backends/evm/ssa/CodeTransform.h

@@ -18,6 +18,8 @@
 
 #pragma once
 
+#include <libyul/backends/evm/ssa/spill/Emitter.h>
+
 #include <libyul/backends/evm/ssa/PhiInverse.h>
 #include <libyul/backends/evm/ssa/Stack.h>
 #include <libyul/backends/evm/ssa/StackLayout.h>
@@ -52,9 +54,17 @@ struct AssemblyCallbacks
 		case StackSlot::Kind::Value:
 		{
 			auto const id = _slot.value();
-			yulAssert(cfg->isLiteral(id), fmt::format("Tried bringing up non-const {}", id));
-			assembly->appendConstant(cfg->literalPayload(id));
-			return;
+			if (cfg->isLiteral(id))
+			{
+				assembly->appendConstant(cfg->literalPayload(id));
+				return;
+			}
+			if (spillEmitter && spillEmitter->hasAddress(id))
+			{
+				spillEmitter->emitLoad(id);
+				return;
+			}
+			yulAssert(false, fmt::format("Tried bringing up non-spilled non-const {}", id));
 		}
 		case StackSlot::Kind::Junk:
 		{
@@ -87,6 +97,7 @@ struct AssemblyCallbacks
 	AbstractAssembly* assembly{};
 	CallSites const* callSites{};
 	std::map<InstId, AbstractAssembly::LabelID> const* returnLabels{};
+	spill::Emitter const* spillEmitter{};
 };
 static_assert(StackManipulationCallbackConcept<AssemblyCallbacks>);
 
@@ -95,6 +106,7 @@ class CodeTransform
 public:
 	static void run(
 		AbstractAssembly& _assembly,
+		ControlFlowGraphs& _controlFlowGraphs,
 		ControlFlowGraphsLiveness const& _liveness,
 		BuiltinContext& _builtinContext
 	);
@@ -115,7 +127,10 @@ class CodeTransform
 		CallSites const& _callSites,
 		SSACFG const& _cfg,
 		SSACFGStackLayout const& _stackLayout,
-		ControlFlowGraphs::FunctionGraphID _graphID);
+		spill::SpillSet const& _spillSet,
+		ControlFlowGraphs::FunctionGraphID _graphID,
+		spill::MemoryAddressing const& _addressing
+	);
 
 	void operator()(SSACFG::BlockId _blockId);
 	void operator()(InstId _instId, StackData const& _operationInputLayout);
@@ -127,17 +142,22 @@ class CodeTransform
 
 	void prepareBlockExitStack(StackData const& _target, PhiInverse const& _phiInverse);
 
+	/// If `_value` is spilled, shuffles it to the stack top and stores it into its memory slot
+	void spillStore(InstId _value);
+
 	AbstractAssembly& m_assembly;
 	BuiltinContext& m_builtinContext;
 	ControlFlowGraphs const& m_controlFlow;
 	FunctionLabels const& m_functionLabels;
 	CallSites const& m_callSites;
 	SSACFG const& m_cfg;
 	SSACFGStackLayout const& m_stackLayout;
+	spill::SpillSet const& m_spillSet;
 	ControlFlowGraphs::FunctionGraphID const m_graphID;
 
 	std::vector<std::uint8_t> m_blockIsTransformed;
 	std::vector<AbstractAssembly::LabelID> m_blockLabels;
+	std::optional<spill::Emitter> m_spillEmitter{std::nullopt};
 	AssemblyCallbacks m_assemblyCallbacks;
 	StackData m_stackData;
 	Stack<AssemblyCallbacks> m_stack;

libyul/backends/evm/ssa/SSACFGTypes.h

@@ -37,7 +37,9 @@ template<typename R, typename T>
 concept InputRangeOf = ranges::input_range<R> && std::same_as<ranges::range_value_t<R>, T>;
 
 class SSACFG;
-
+class SSACFGStackLayout;
+class StackSlot;
+using StackData = std::vector<StackSlot>;
 using FunctionGraphID = std::uint32_t;
 
 struct BlockId