[EVM] Transform icmp ult (add x, CAdd), CCmp -> icmp eq (evm.signextend(b, x)), x

Signed-off-by: Vladimir Radosavljevic <[email protected]>

Author: vladimirradosavljevic

llvm/lib/Target/EVM/EVMCodegenPrepare.cpp

@@ -20,11 +20,13 @@
 #include "llvm/IR/IntrinsicInst.h"
 #include "llvm/IR/Intrinsics.h"
 #include "llvm/IR/IntrinsicsEVM.h"
+#include "llvm/IR/PatternMatch.h"
 #include "llvm/Pass.h"
 
 #include "EVM.h"
 
 using namespace llvm;
+using namespace llvm::PatternMatch;
 
 #define DEBUG_TYPE "evm-codegen-prepare"
 
@@ -102,6 +104,50 @@ void EVMCodegenPrepare::processMemTransfer(MemTransferInst *M) {
   M->setCalledFunction(Intrinsic::getDeclaration(M->getModule(), IntrID));
 }
 
+static bool optimizeICmp(ICmpInst *I) {
+  auto *Ty = I->getOperand(0)->getType();
+  if (!Ty->isIntegerTy(256))
+    return false;
+
+  if (I->getPredicate() == CmpInst::ICMP_ULT) {
+    Value *X = nullptr;
+    const APInt *CAdd = nullptr, *CCmp = nullptr;
+
+    // icmp ult (add x, CAdd), CCmp -> icmp eq (evm.signextend(b, x)), x
+    // where CCmp is a power of 2 and CAdd is twice smaller than CCmp.
+    if (match(I->getOperand(0), m_OneUse(m_c_Add(m_Value(X), m_APInt(CAdd)))) &&
+        match(I->getOperand(1), m_APInt(CCmp)) && CCmp->isPowerOf2() &&
+        *CAdd == CCmp->lshr(1)) {
+      unsigned CCmpLog2 = CCmp->logBase2();
+
+      // If CCmpLog2 is not divisible by 8, cannot use signextend.
+      if (CCmpLog2 % 8 != 0)
+        return false;
+
+      IRBuilder<> Builder(I);
+      unsigned ByteIdx = (CCmpLog2 / 8) - 1;
+
+      // ByteIdx should be in [0, 31].
+      if (ByteIdx > 31)
+        return false;
+
+      auto *B = ConstantInt::get(Ty, ByteIdx);
+      auto *SignExtend =
+          Builder.CreateIntrinsic(Ty, Intrinsic::evm_signextend, {B, X});
+      auto *NewCmp = Builder.CreateICmp(CmpInst::ICMP_EQ, SignExtend, X);
+      NewCmp->takeName(I);
+      I->replaceAllUsesWith(NewCmp);
+
+      // Remove add after icmp. If to do otherwise, assert will be triggered.
+      auto *ToRemove = cast<Instruction>(I->getOperand(0));
+      I->eraseFromParent();
+      ToRemove->eraseFromParent();
+      return true;
+    }
+  }
+  return false;
+}
+
 bool EVMCodegenPrepare::runOnFunction(Function &F) {
   bool Changed = false;
   for (auto &BB : F) {
@@ -110,6 +156,8 @@ bool EVMCodegenPrepare::runOnFunction(Function &F) {
         processMemTransfer(M);
         Changed = true;
       }
+      if (I.getOpcode() == Instruction::ICmp)
+        Changed |= optimizeICmp(cast<ICmpInst>(&I));
     }
   }