From d78e6c509cdf231395883c598f94c2871607bec6 Mon Sep 17 00:00:00 2001
From: Dmitri Makarov <dmakarov@alumni.stanford.edu>
Date: Thu, 21 Oct 2021 20:00:40 -0700
Subject: [PATCH] [SOL] Prevent breaking ISelDAG connectivity on replacing
 loads by constants

---
 llvm/lib/Target/BPF/BPFISelDAGToDAG.cpp   | 41 +++++++++++++++++++++++
 llvm/test/CodeGen/BPF/preprocess-loads.ll | 35 +++++++++++++++++++
 2 files changed, 76 insertions(+)
 create mode 100644 llvm/test/CodeGen/BPF/preprocess-loads.ll

diff --git a/llvm/lib/Target/BPF/BPFISelDAGToDAG.cpp b/llvm/lib/Target/BPF/BPFISelDAGToDAG.cpp
index f10a0d4c007729..84159d40fd8658 100644
--- a/llvm/lib/Target/BPF/BPFISelDAGToDAG.cpp
+++ b/llvm/lib/Target/BPF/BPFISelDAGToDAG.cpp
@@ -304,6 +304,47 @@ void BPFDAGToDAGISel::PreprocessLoad(SDNode *Node,
 
   LLVM_DEBUG(dbgs() << "Replacing load of size " << size << " with constant "
                     << val << '\n');
+
+  /* Some load nodes have edges from TokenFactor nodes.  In this case
+     replacing the load with a constant makes the DAG disconnected.
+     The following checks if any of the load operands are TokenFactor
+     nodes, and if another TokenFactor is a user of the load, the
+     operand TokenFactor is connected to the user, so that the DAG
+     remains connected after replacing the load node by a constant.
+  */
+  for (unsigned I = 0, E = Node->getNumOperands(); I != E; ++I) {
+    const SDValue &OpV = Node->getOperand(I);
+    SDNode *Op = OpV.getNode();
+    if (Op->getOpcode() == ISD::TokenFactor) {
+      for (SDNode::use_iterator UI = Node->use_begin(), UE = Node->use_end(); UI != UE; ++UI) {
+        SDUse &Use = UI.getUse();
+        SDNode *User = Use.getUser();
+        if (User->getOpcode() == ISD::TokenFactor) {
+          SmallVector<SDValue, 8> ExtendedOps;
+          bool NotExtended = true;
+          for (unsigned UOI = 0, UOE = User->getNumOperands(); UOI != UOE; ++UOI) {
+            const SDValue &Operand = User->getOperand(UOI);
+            if (OpV == Operand) {
+              NotExtended = false;
+              break;
+            }
+            ExtendedOps.push_back(Operand);
+          }
+          if (NotExtended) {
+            ExtendedOps.push_back(OpV);
+            SDValue ExtendedTokenFactor = CurDAG->getTokenFactor(SDLoc(User), ExtendedOps);
+            I--;
+            SDValue From[] = {SDValue(User, 0)};
+            SDValue To[] = {ExtendedTokenFactor};
+            CurDAG->ReplaceAllUsesOfValuesWith(From, To, 1);
+            I++;
+            CurDAG->DeleteNode(User);
+          }
+        }
+      }
+    }
+  }
+
   SDValue NVal = CurDAG->getConstant(val, DL, LD->getValueType(0));
 
   // After replacement, the current node is dead, we need to
diff --git a/llvm/test/CodeGen/BPF/preprocess-loads.ll b/llvm/test/CodeGen/BPF/preprocess-loads.ll
new file mode 100644
index 00000000000000..08701d1a381b7d
--- /dev/null
+++ b/llvm/test/CodeGen/BPF/preprocess-loads.ll
@@ -0,0 +1,35 @@
+; RUN: llc -O2 -march=bpfel --mattr=+solana -filetype=asm < %s | FileCheck --check-prefix=CHECK %s
+
+%Pool = type <{ [0 x i8], [32 x i8], [0 x i8], i8, [0 x i8], [10 x %Decimal], [0 x i8] }>
+%Decimal = type { [0 x i32], i32, [0 x i32], i32, [0 x i32], i32, [0 x i32], i32, [0 x i32] }
+
+@0 = private unnamed_addr constant <{ [16 x i8] }> <{ [16 x i8] c"\00\00\00\00\FF\FF\FF\FF\FF\FF\FF\FF\FF\FF\FF\FF" }>, align 4
+
+declare void @llvm.memcpy.p0i8.p0i8.i64(i8* noalias nocapture writeonly, i8* noalias nocapture readonly, i64, i1 immarg) #5
+
+; optbug::Pool::initialize
+; Function Attrs: nofree noinline nounwind willreturn
+define internal fastcc void @_ZN6optbug4Pool10initialize17h656ab4e5f05591beE(%Pool* noalias nocapture align 1 dereferenceable(193) %self, [32 x i8]* noalias nocapture readonly align 1 dereferenceable(32) %seed, i8 %bump) unnamed_addr #2 {
+start:
+  %0 = getelementptr inbounds [32 x i8], [32 x i8]* %seed, i64 0, i64 0
+  %self56 = getelementptr inbounds %Pool, %Pool* %self, i64 0, i32 0, i64 0
+  call void @llvm.memcpy.p0i8.p0i8.i64(i8* nonnull align 1 dereferenceable(32) %self56, i8* nonnull align 1 dereferenceable(32) %0, i64 32, i1 false)
+  %1 = getelementptr inbounds %Pool, %Pool* %self, i64 0, i32 3
+  store i8 %bump, i8* %1, align 1
+  %2 = getelementptr inbounds %Pool, %Pool* %self, i64 0, i32 5, i64 1
+  %3 = bitcast %Decimal* %2 to i8*
+  tail call void @llvm.memcpy.p0i8.p0i8.i64(i8* nonnull align 1 dereferenceable(16) %3, i8* nonnull align 4 dereferenceable(16) getelementptr inbounds (<{ [16 x i8] }>, <{ [16 x i8] }>* @0, i64 0, i32 0, i64 0), i64 16, i1 false)
+  ret void
+}
+
+; CHECK: *(u64 *)(r1 + 24) = r4
+; CHECK: *(u64 *)(r1 + 16) = r4
+; CHECK: *(u64 *)(r1 + 8) = r4
+; CHECK: *(u64 *)(r1 + 0) = r2
+; CHECK: *(u8 *)(r1 + 32) = r3
+
+attributes #2 = { nofree noinline nounwind willreturn "target-cpu"="generic" }
+
+!llvm.module.flags = !{!0}
+
+!0 = !{i32 7, !"PIC Level", i32 2}