From 448b727d14058247aba5ca55c5dd0d38d74c94ce Mon Sep 17 00:00:00 2001
From: Nuno Lopes <nuno.lopes@tecnico.ulisboa.pt>
Date: Sun, 8 Oct 2023 11:10:35 +0100
Subject: [PATCH] preprocessor: improve alignment inference

---
 ir/instr.cpp        | 17 -----------------
 ir/instr.h          |  1 -
 tools/transform.cpp | 19 +++++++++++++++----
 3 files changed, 15 insertions(+), 22 deletions(-)

diff --git a/ir/instr.cpp b/ir/instr.cpp
index c7edb6096..1290b03f1 100644
--- a/ir/instr.cpp
+++ b/ir/instr.cpp
@@ -3619,23 +3619,6 @@ uint64_t GEP::getMaxGEPOffset() const {
   return off;
 }
 
-optional<uint64_t> GEP::getExactOffset() const {
-  uint64_t off = 0;
-  for (auto &[mul, v] : getIdxs()) {
-    if (mul == 0)
-      continue;
-    if (mul >= INT64_MAX)
-      return {};
-
-    if (auto n = getInt(*v)) {
-      off += (int64_t)mul * *n;
-      continue;
-    }
-    return {};
-  }
-  return off;
-}
-
 vector<Value*> GEP::operands() const {
   vector<Value*> v = { ptr };
   for (auto &[sz, idx] : idxs) {
diff --git a/ir/instr.h b/ir/instr.h
index 011bc1724..f7666e81b 100644
--- a/ir/instr.h
+++ b/ir/instr.h
@@ -767,7 +767,6 @@ class GEP final : public MemInstr {
   Value& getPtr() const { return *ptr; }
   auto& getIdxs() const { return idxs; }
   bool isInBounds() const { return inbounds; }
-  std::optional<uint64_t> getExactOffset() const;
 
   std::pair<uint64_t, uint64_t> getMaxAllocSize() const override;
   uint64_t getMaxAccessSize() const override;
diff --git a/tools/transform.cpp b/tools/transform.cpp
index 709121a4f..2e4aed920 100644
--- a/tools/transform.cpp
+++ b/tools/transform.cpp
@@ -1689,10 +1689,21 @@ void Transform::preprocess() {
           continue;
         }
       } else if (auto *gep = dynamic_cast<const GEP*>(i)) {
-        auto off = gep->getExactOffset();
-        if (!off || !is_power2(*off))
-          continue;
-        align = min(aligns[&gep->getPtr()], *off);
+        int mul_bits = sizeof(uint64_t) * 8 - 1;
+        uint64_t const_offset = 0;
+        for (auto &[mul, v] : gep->getIdxs()) {
+          if (mul == 0)
+            continue;
+          if (auto n = getInt(*v))
+            const_offset += mul * (uint64_t)*n;
+          else
+            mul_bits = min(mul_bits, countr_zero(mul));
+        }
+        align = min(aligns[&gep->getPtr()], ((uint64_t)1) << mul_bits);
+        if (const_offset) {
+          align = min(align, ((uint64_t)1) << countr_zero(const_offset));
+        }
+
       } else if (auto *phi = dynamic_cast<const Phi*>(i)) {
         // optimistic handling of phis: unreachable predecessors don't
         // contribute to the result. This is revisited once they become reach