From 108f3d9e4b18adedb9ada213914ccfc7adedc638 Mon Sep 17 00:00:00 2001
From: CGH0S7 <776459475@qq.com>
Date: Tue, 30 Jun 2026 00:25:29 +0800
Subject: [PATCH] strength reduce power-of-two GEP offsets

---
 include/mir/MIR.h           |  1 +
 src/mir/AsmPrinter.cpp      |  5 +++++
 src/mir/Lowering.cpp        | 24 ++++++++++++++++++++++--
 src/mir/passes/Peephole.cpp |  1 +
 4 files changed, 29 insertions(+), 2 deletions(-)

diff --git a/include/mir/MIR.h b/include/mir/MIR.h
index de519aa..6aec624 100644
--- a/include/mir/MIR.h
+++ b/include/mir/MIR.h
@@ -55,6 +55,7 @@ enum class Opcode {
   MovReg,
   Adrp,
   AddRegImm,
+  LslImm,
   LdrRegReg,
   StrRegReg,
   SIToFP,
diff --git a/src/mir/AsmPrinter.cpp b/src/mir/AsmPrinter.cpp
index d69999a..3f730ab 100644
--- a/src/mir/AsmPrinter.cpp
+++ b/src/mir/AsmPrinter.cpp
@@ -250,6 +250,11 @@ void PrintAsm(const MachineFunction& function, std::ostream& os) {
           }
           break;
         }
+        case Opcode::LslImm:
+          os << "  lsl " << PhysRegName(ops.at(0).GetReg()) << ", "
+             << PhysRegName(ops.at(1).GetReg()) << ", #"
+             << ops.at(2).GetImm() << "\n";
+          break;
         case Opcode::LdrRegReg: {
           PhysReg reg = ops.at(0).GetReg();
           const char* ldr_cmd = IsFloatReg(reg) ? "ldr" : "ldr";
diff --git a/src/mir/Lowering.cpp b/src/mir/Lowering.cpp
index 8e705ca..4a5395d 100644
--- a/src/mir/Lowering.cpp
+++ b/src/mir/Lowering.cpp
@@ -112,6 +112,19 @@ uint32_t GetAllocaSize(const ir::Instruction& inst, const std::unordered_set<con
   return GetTypeSize(type.get());
 }
 
+bool IsPowerOfTwo(uint32_t value) {
+  return value != 0 && (value & (value - 1)) == 0;
+}
+
+int Log2(uint32_t value) {
+  int shift = 0;
+  while (value > 1) {
+    value >>= 1;
+    shift++;
+  }
+  return shift;
+}
+
 bool LooksLikeConstantArrayModuloSumLoop(const ir::Function& function,
                                          int* per_iteration_sum,
                                          int* modulo) {
@@ -626,13 +639,20 @@ void LowerInstruction(const ir::Instruction& inst, MachineFunction& function,
         }
 
         EmitValueToReg(idx, PhysReg::W9, slots, block);
-        if (stride > 1) {
+        bool shifted = false;
+        if (stride > 1 && IsPowerOfTwo(stride)) {
+          block.Append(Opcode::ZExt, {Operand::Reg(PhysReg::X9), Operand::Reg(PhysReg::W9)});
+          block.Append(Opcode::LslImm, {Operand::Reg(PhysReg::X9), Operand::Reg(PhysReg::X9), Operand::Imm(Log2(stride))});
+          shifted = true;
+        } else if (stride > 1) {
           block.Append(Opcode::MovImm, {Operand::Reg(PhysReg::W10), Operand::Imm(stride)});
           block.Append(Opcode::MulRR, {Operand::Reg(PhysReg::W9), Operand::Reg(PhysReg::W9), Operand::Reg(PhysReg::W10)});
         }
 
         // Extend W9 to X9 and add to base address X8
-        block.Append(Opcode::ZExt, {Operand::Reg(PhysReg::X9), Operand::Reg(PhysReg::W9)});
+        if (!shifted) {
+          block.Append(Opcode::ZExt, {Operand::Reg(PhysReg::X9), Operand::Reg(PhysReg::W9)});
+        }
         block.Append(Opcode::AddRR, {Operand::Reg(PhysReg::X8), Operand::Reg(PhysReg::X8), Operand::Reg(PhysReg::X9)});
       }
 
diff --git a/src/mir/passes/Peephole.cpp b/src/mir/passes/Peephole.cpp
index 3b8418b..daca461 100644
--- a/src/mir/passes/Peephole.cpp
+++ b/src/mir/passes/Peephole.cpp
@@ -150,6 +150,7 @@ void RunPeephole(MachineFunction& function) {
         case Opcode::MovReg:
         case Opcode::Adrp:
         case Opcode::AddRegImm:
+        case Opcode::LslImm:
         case Opcode::LdrRegReg:
         case Opcode::SIToFP:
         case Opcode::FPToSI: