Nella/slp rework

src/backend/analysis/alias_analysis.cpp

@@ -1,15 +1,18 @@
 #include "analysis/alias_analysis.h"
 
+#include "ir/function.h"
 #include "ir/node.h"
+#include "ir/opcode.h"
 #include "ir/type.h"
 
 namespace rat {
 	namespace detail {
 		B32 idLess(const Node* a, const Node* b) { return a->getId() < b->getId(); }
 	} // namespace detail
 
-	AliasAnalysis::AliasAnalysis(U32 pointerBytes)
-	: ptrBytes(pointerBytes) {}
+	AliasAnalysis::AliasAnalysis(U32 pointerBytes, U64 restrictParams)
+	: ptrBytes(pointerBytes),
+		restrictParams(restrictParams) {}
 
 	const AliasAnalysis::Address& AliasAnalysis::decompose(Node* addr) const {
 		auto it = decomposeCache.find(addr);
@@ -80,6 +83,83 @@ namespace rat {
 		return cast<GlobalNode>(a)->getSymbol() != cast<GlobalNode>(b)->getSymbol();
 	}
 
+	I32 AliasAnalysis::paramIndex(const Node* n) {
+		const ProjNode* p = dyn_cast<ProjNode>(n);
+		U32 first = StartNode::paramProjIndex(0);
+		if(!p || !isa<StartNode>(p->getProducer()) || p->getIndex() < first)
+			return -1;
+		return (I32)(p->getIndex() - first);
+	}
+
+	B32 AliasAnalysis::isRestrictParam(I32 index, U64 restricts) {
+		return index >= 0 && index < (I32)kMaxRestrictParams && (restricts >> index & 1);
+	}
+
+	B32 AliasAnalysis::mayCarryRestrict(List<const Node*> work, U64 restricts, U32 cap) {
+		Set<const Node*> seen;
+		while(!work.empty()) {
+			const Node* c = work.back();
+			work.pop_back();
+			if(!c || !seen.insert(c).second)
+				continue;
+			if(seen.size() > cap)
+				return true;
+			Opcode op = c->getOpcode();
+			I32 index = paramIndex(c);
+			if(isRestrictParam(index, restricts))
+				return true;
+			if(index >= 0 || op == Opcode::Constant || op == Opcode::Load || op == Opcode::Proj ||
+				 isIdentified(c))
+				continue;
+			B32 pure = isArithmeticOpcode(op) || isVectorUtilOpcode(op) || op == Opcode::Select ||
+								 op == Opcode::Phi;
+			if(!pure)
+				return true;
+			for(U32 i = 0; i < c->getInputCount(); ++i)
+				if(const Node* in = c->getInput(i); in && in->getType() && in->getType()->isData())
+					work.push_back(in);
+		}
+		return false;
+	}
+
+	U64 AliasAnalysis::restrictFacts(const Function& fn) {
+		U64 restricts = fn.getAttrs().restrictParams;
+		if(!restricts)
+			return 0;
+		List<const Node*> taken;
+		for(Node* n : fn) {
+			if(StoreNode* s = dyn_cast<StoreNode>(n); s && s->getInputCount())
+				taken.push_back(s->getValue());
+			else if(isa<CallNode>(n) || isa<AsmNode>(n))
+				for(U32 i = 2; i < n->getInputCount(); ++i)
+					taken.push_back(n->getInput(i));
+		}
+		if(mayCarryRestrict(std::move(taken), restricts, ~0u))
+			return 0;
+		return restricts;
+	}
+
+	Node* AliasAnalysis::objectOf(Node* addr) const {
+		const Address& a = decompose(addr);
+		Node* base = a.base;
+		I32 index = paramIndex(base);
+		B32 isRestrict = isRestrictParam(index, restrictParams);
+		if(!isRestrict && !isIdentified(base))
+			return nullptr;
+		for(Node* t : a.symbolic)
+			if(restrictParams && mayCarryRestrict({t}, restrictParams, kMaxCarryCone))
+				return nullptr;
+		if(isRestrict) {
+			if(!paramObjects[index])
+				paramObjects[index] = base;
+			return paramObjects[index];
+		}
+		GlobalNode* g = dyn_cast<GlobalNode>(base);
+		if(!g)
+			return base;
+		return symbolObjects.try_emplace(g->getSymbol(), base).first->second;
+	}
+
 	AliasResult AliasAnalysis::alias(Node* addrA, U32 sizeA, Node* addrB, U32 sizeB) const {
 		// alias query between two addresses with known access sizes (0 = unknown)
 		const Address& a = decompose(addrA);

src/backend/analysis/alias_analysis.h

@@ -11,6 +11,7 @@
 
 #include "core.h"
 #include "hash.h"
+#include "ir/function.h"
 
 namespace rat {
 	struct Node;
@@ -22,7 +23,9 @@ namespace rat {
 	enum class AliasResult { NoAlias, MayAlias, MustAlias };
 
 	struct AliasAnalysis {
-		AliasAnalysis(U32 pointerBytes);
+		static constexpr U32 kMaxCarryCone = 32; // cone nodes mayCarryRestrict walks per address term
+
+		AliasAnalysis(U32 pointerBytes, U64 restrictParams);
 
 		AliasResult alias(Node* addrA, U32 sizeA, Node* addrB, U32 sizeB) const;
 
@@ -57,6 +60,11 @@ namespace rat {
 		// object identity, independent of any offset
 		static B32 isIdentified(const Node* n); // an Alloc or a Global
 		static B32 distinctObjects(const Node* a, const Node* b);
+		static I32 paramIndex(const Node* n); // parameter of a Start projection, else -1
+		static B32 isRestrictParam(I32 index, U64 restricts);
+		static B32 mayCarryRestrict(List<const Node*> work, U64 restricts, U32 cap);
+		static U64 restrictFacts(const Function& fn);
+		Node* objectOf(Node* addr) const;
 	private:
 		struct Address {
 			Node* base;						// pointer
@@ -68,7 +76,11 @@ namespace rat {
 		static Node* accessAddress(Node* n); // pointer operand of a Load/Store (else null)
 	private:
 		U32 ptrBytes;
+		U64 restrictParams;
 		mutable Map<Node*, Address> decomposeCache;
+		mutable Map<String, Node*> symbolObjects; // first global node seen per symbol
+		// first projection seen per restrict parameter
+		mutable Node* paramObjects[kMaxRestrictParams] = {};
 	};
 } // namespace rat
 

src/backend/codegen/schedule.cpp

@@ -2,6 +2,7 @@
 
 #include "analysis/alias_analysis.h"
 #include "ir/function.h"
+#include "ir/module.h"
 #include "ir/node.h"
 
 namespace rat {
@@ -394,6 +395,18 @@ namespace rat {
 		return d == 0 || (isSigned && d == -1); // INT_MIN / -1 traps
 	}
 
+	// a load inside a global this module defines cannot fault, so it may run before a branch
+	B32 Schedule::loadsGlobal(const Node* n) const {
+		const LoadNode* l = dyn_cast<LoadNode>(n);
+		if(!l)
+			return false;
+		const GlobalNode* g = dyn_cast<GlobalNode>(l->getPointer());
+		if(!g)
+			return false;
+		const Global* obj = fn.getModule().getGlobal(g->getSymbol());
+		return obj && !obj->isAlias() && obj->getType()->byteSize(8) >= l->getType()->byteSize(8);
+	}
+
 	B32 Schedule::isGuarded(const Node* n) const {
 		if(isa<LoadNode>(n))
 			return false;
@@ -414,7 +427,7 @@ namespace rat {
 			return late; // nothing above is shallower, so the walk cannot move it
 		Opcode op = n->getOpcode();
 		B32 remat = op == Opcode::Constant || op == Opcode::Global;
-		B32 trapping = mayTrap(n) || (!guarded.empty() && guarded[n->getId()]);
+		B32 trapping = (mayTrap(n) && !loadsGlobal(n)) || (!guarded.empty() && guarded[n->getId()]);
 		I32 cur = late, pick = late;
 		while(true) {
 			if(blocks[cur].loopDepth < blocks[pick].loopDepth) {
@@ -615,21 +628,40 @@ namespace rat {
 			flat[fill[keys[i]]++] = items[i];
 	}
 
+	B32 Schedule::emitsEarly(const Node* n) {
+		NodeSpan users = n->getUsers();
+		if(users.empty() || isa<LoadNode>(n) || !isFloating(n))
+			return true;
+		for(Node* u : users)
+			if(!isa<PhiNode>(u))
+				return true;
+		B32 readsUser = false;
+		for(U32 i = 0, e = n->getInputCount(); i < e; ++i) {
+			Node* in = n->getInput(i);
+			B32 user = std::find(users.begin(), users.end(), in) != users.end();
+			if(!user && !isa<ConstantNode>(in))
+				return true;
+			readsUser |= user;
+		}
+		return !readsUser;
+	}
+
 	void Schedule::buildBlockLists() {
 		U32 nb = (U32)blocks.size();
 		List<I32> keys;
 		for(Node* phi : dataPhis)
 			keys.push_back(headBlock(cast<PhiNode>(phi)->getRegion()));
 		bucket(dataPhis, keys, nb, phiStart, phiFlat);
 
-		// placed nodes in function (id) order
+		// placed nodes in function (id) order, induction steps last
 		List<Node*> listed;
 		for(Node* n : floating)
 			if(nodeBlock[n->getId()] >= 0)
 				listed.push_back(n);
 		auto mid = listed.insert(listed.end(), pinned.begin(), pinned.end());
 		std::sort(mid, listed.end(), detail::earlierId);
 		std::inplace_merge(listed.begin(), mid, listed.end(), detail::earlierId);
+		std::stable_partition(listed.begin(), listed.end(), emitsEarly);
 		keys.clear();
 		for(Node* n : listed)
 			keys.push_back(nodeBlock[n->getId()]);
@@ -638,7 +670,7 @@ namespace rat {
 		TopoScratch scratch;
 		scratch.localOf.assign(fn.idBound(), -1);
 		scratch.stHead.assign(fn.idBound(), -1);
-		AliasAnalysis aa(8);
+		AliasAnalysis aa(8, AliasAnalysis::restrictFacts(fn));
 		for(U32 b = 0; b < nb; ++b)
 			topoOrder(
 					nodeFlat.data() + nodeStart[b], (U32)(nodeStart[b + 1] - nodeStart[b]), aa, scratch);
@@ -670,6 +702,14 @@ namespace rat {
 		}
 
 		B32 earlierId(const Node* a, const Node* b) { return a->getId() < b->getId(); }
+
+		// the object a store writes, null for other writers
+		Node* storedObject(const AliasAnalysis& aa, Node* w) {
+			StoreNode* st = dyn_cast<StoreNode>(w);
+			if(!st)
+				return nullptr;
+			return aa.objectOf(st->getPointer());
+		}
 	} // namespace detail
 
 	I32 Schedule::TopoScratch::local(const Node* n) const { return n ? localOf[n->getId()] : -1; }
@@ -743,24 +783,91 @@ namespace rat {
 			std::copy(s.out.begin(), s.out.end(), nodes);
 	}
 
+	// a load runs before the first writer on its state's chain that may clobber it
 	void Schedule::addAntiDeps(NodeSpan nodes, const AliasAnalysis& aa, TopoScratch& s) {
+		s.chain.clear();
+		for(Node* n : nodes) {
+			LoadNode* ld = dyn_cast<LoadNode>(n);
+			if(!ld || !ld->getMemory())
+				continue;
+			I32 start = s.stHead[ld->getMemory()->getId()];
+			if(start < 0)
+				continue;
+			if(I32 w = firstClobber(nodes, ld, start, aa, s); w >= 0)
+				s.addEdge(ld, nodes[w]);
+		}
+	}
+
+	void Schedule::rankWriters(NodeSpan nodes, TopoScratch& s) {
 		U32 k = (U32)nodes.size();
+		s.rank.assign(k, -1);
 		for(U32 i = 0; i < k; ++i) {
-			LoadNode* ld = dyn_cast<LoadNode>(nodes[i]);
-			if(!ld)
+			Node* mem = memoryInputOf(nodes[i]);
+			if(!mem || !isMemWriter(nodes[i]))
 				continue;
-			Node* state = ld->getMemory();
-			U32 guard = 0; // chain length bound; conservatively pin if exceeded
-			for(I32 wi = state ? s.stHead[state->getId()] : -1; wi >= 0; ++guard) {
-				Node* w = nodes[wi];
-				StoreNode* st = dyn_cast<StoreNode>(w);
-				if(!st || detail::storeMayAliasLoad(aa, st, ld) || guard >= k) {
-					s.addEdge(ld, w); // call, aliasing store, or bound hit -> pin here
+			if(isa<StoreNode>(mem) && s.local(mem) >= 0 && s.stHead[mem->getId()] == (I32)i)
+				continue; // inside the chain of mem
+			for(I32 w = (I32)i; w >= 0; w = s.stHead[nodes[w]->getId()]) {
+				s.rank[w] = (I32)s.chain.size();
+				s.chain.push_back(w);
+				if(!isa<StoreNode>(nodes[w]))
 					break;
-				}
-				wi = s.stHead[w->getId()]; // skip disjoint store, walk to next writer
 			}
+			s.chain.push_back(-1);
+		}
+	}
+
+	void Schedule::indexObjects(NodeSpan nodes, const AliasAnalysis& aa, TopoScratch& s) {
+		s.nextWild.resize(s.chain.size());
+		s.byObject.clear();
+		for(U32 p = (U32)s.chain.size(); p-- > 0;) {
+			I32 w = s.chain[p];
+			Node* obj = nullptr;
+			if(w >= 0)
+				obj = detail::storedObject(aa, nodes[w]);
+			if(obj) {
+				s.nextWild[p] = s.nextWild[p + 1];
+				s.byObject.push_back({obj->getId(), (I32)p});
+			} else {
+				s.nextWild[p] = (I32)p;
+			}
+		}
+		std::sort(s.byObject.begin(), s.byObject.end());
+	}
+
+	// first writer from start on that is not a store or may alias ld, or -1
+	I32 Schedule::walkClobber(
+			NodeSpan nodes, LoadNode* ld, I32 start, const AliasAnalysis& aa, const TopoScratch& s) {
+		I32 w = start;
+		while(w >= 0) {
+			StoreNode* st = dyn_cast<StoreNode>(nodes[w]);
+			if(!st || detail::storeMayAliasLoad(aa, st, ld))
+				break;
+			w = s.stHead[st->getId()];
+		}
+		return w;
+	}
+
+	I32 Schedule::firstClobber(
+			NodeSpan nodes, LoadNode* ld, I32 start, const AliasAnalysis& aa, TopoScratch& s) {
+		Node* obj = aa.objectOf(ld->getPointer());
+		if(!obj)
+			return walkClobber(nodes, ld, start, aa, s); // a store to another base may alias
+		if(s.chain.empty()) {
+			rankWriters(nodes, s);
+			indexObjects(nodes, aa, s);
+		}
+		I32 r = s.rank[start];
+		assert(r >= 0 && "writer outside every chain");
+		I32 bound = s.nextWild[r];
+		Pair<U32, I32> key{obj->getId(), r};
+		auto it = std::lower_bound(s.byObject.begin(), s.byObject.end(), key);
+		for(; it != s.byObject.end() && it->first == key.first && it->second < bound; ++it) {
+			I32 w = s.chain[it->second];
+			if(detail::storeMayAliasLoad(aa, cast<StoreNode>(nodes[w]), ld))
+				return w;
 		}
+		return s.chain[bound];
 	}
 
 	void Schedule::addOrderEdges(NodeSpan nodes, TopoScratch& s) {