[EVM] Unfold signextend(c, x) -> ashr(shl(x, 256 - (c + 1) * 8), 256 - (c + 1) * 8)

This transform does unfold to canonical InstCombine form,
instead to sext(trunc(x, (c + 1) * 8), 256).
Unfold signextend to LLVM instructions, iff c is constant,
so LLVM can optimize it better. Do the folding back just
before ISel, since we can do cross BB boundaries.

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

Author: vladimirradosavljevic

llvm/lib/Target/EVM/EVMCodegenPrepare.cpp

@@ -20,11 +20,14 @@
 #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 "llvm/Transforms/InstCombine/InstCombiner.h"
 
 #include "EVM.h"
 
 using namespace llvm;
+using namespace llvm::PatternMatch;
 
 #define DEBUG_TYPE "evm-codegen-prepare"
 
@@ -102,14 +105,47 @@ void EVMCodegenPrepare::processMemTransfer(MemTransferInst *M) {
   M->setCalledFunction(Intrinsic::getDeclaration(M->getModule(), IntrID));
 }
 
+static bool optimizeAShrInst(Instruction *I) {
+  auto *Ty = I->getType();
+  unsigned BitWidth = Ty->getIntegerBitWidth();
+  if (BitWidth != 256)
+    return false;
+
+  // Fold ashr(shl(x, c), c) -> signextend(((256 - c) / 8) - 1, x)
+  // where c is a constant and divisible by 8.
+  Value *X = nullptr;
+  ConstantInt *ShiftAmt = nullptr;
+  if (match(I->getOperand(0),
+            m_OneUse(m_Shl(m_Value(X), m_ConstantInt(ShiftAmt)))) &&
+      match(I->getOperand(1), m_Specific(ShiftAmt)) &&
+      ShiftAmt->getZExtValue() % 8 == 0) {
+    IRBuilder<> Builder(I);
+    unsigned ByteIdx = ((BitWidth - ShiftAmt->getZExtValue()) / 8) - 1;
+    auto *B = ConstantInt::get(Ty, ByteIdx);
+    auto *SignExtend =
+        Builder.CreateIntrinsic(Ty, Intrinsic::evm_signextend, {B, X});
+    SignExtend->takeName(I);
+    I->replaceAllUsesWith(SignExtend);
+
+    // Remove shl after ashr. 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) {
-    for (auto &I : BB) {
+    for (auto &I : make_early_inc_range(BB)) {
       if (auto *M = dyn_cast<MemTransferInst>(&I)) {
         processMemTransfer(M);
         Changed = true;
       }
+      if (I.getOpcode() == Instruction::AShr)
+        Changed |= optimizeAShrInst(&I);
     }
   }
 

llvm/lib/Target/EVM/EVMTargetMachine.cpp

@@ -233,8 +233,8 @@ bool EVMPassConfig::addPreISel() {
 }
 
 void EVMPassConfig::addCodeGenPrepare() {
-  addPass(createEVMCodegenPreparePass());
   TargetPassConfig::addCodeGenPrepare();
+  addPass(createEVMCodegenPreparePass());
 }
 
 bool EVMPassConfig::addInstSelector() {

llvm/lib/Target/EVM/EVMTargetTransformInfo.cpp

@@ -21,6 +21,28 @@ using namespace llvm::PatternMatch;
 
 static std::optional<Instruction *> instCombineSignExtend(InstCombiner &IC,
                                                           IntrinsicInst &II) {
+  unsigned BitWidth = II.getType()->getIntegerBitWidth();
+  if (BitWidth != 256)
+    return std::nullopt;
+
+  // Unfold signextend(c, x) ->
+  //        ashr(shl(x, 256 - (c + 1) * 8), 256 - (c + 1) * 8)
+  // where c is a constant integer.
+  ConstantInt *C = nullptr;
+  if (match(II.getArgOperand(0), m_ConstantInt(C))) {
+    const APInt &B = C->getValue();
+
+    // If the signextend is larger than 31 bits, leave constant
+    // folding to handle it.
+    if (B.uge(APInt(BitWidth, (BitWidth / 8) - 1)))
+      return std::nullopt;
+
+    unsigned ShiftAmt = BitWidth - ((B.getZExtValue() + 1) * 8);
+    auto *Shl = IC.Builder.CreateShl(II.getArgOperand(1), ShiftAmt);
+    auto *Ashr = IC.Builder.CreateAShr(Shl, ShiftAmt);
+    return IC.replaceInstUsesWith(II, Ashr);
+  }
+
   // Fold signextend(b, signextend(b, x)) -> signextend(b, x)
   Value *B = nullptr, *X = nullptr;
   if (match(&II, m_Intrinsic<Intrinsic::evm_signextend>(

llvm/test/CodeGen/EVM/O3-pipeline.ll

@@ -55,10 +55,8 @@ target triple = "evm"
 ; CHECK-NEXT:       Expand reduction intrinsics
 ; CHECK-NEXT:       Natural Loop Information
 ; CHECK-NEXT:       TLS Variable Hoist
-; CHECK-NEXT:       Final transformations before code generation
-; CHECK-NEXT:       Dominator Tree Construction
-; CHECK-NEXT:       Natural Loop Information
 ; CHECK-NEXT:       CodeGen Prepare
+; CHECK-NEXT:       Final transformations before code generation
 ; CHECK-NEXT:       Lower invoke and unwind, for unwindless code generators
 ; CHECK-NEXT:       Remove unreachable blocks from the CFG
 ; CHECK-NEXT:     CallGraph Construction

llvm/test/CodeGen/EVM/fold-signextend.ll

@@ -7,7 +7,8 @@ target triple = "evm"
 define i256 @test_const(i256 %x) {
 ; CHECK-LABEL: define i256 @test_const(
 ; CHECK-SAME: i256 [[X:%.*]]) {
-; CHECK-NEXT:    [[SIGNEXT1:%.*]] = call i256 @llvm.evm.signextend(i256 15, i256 [[X]])
+; CHECK-NEXT:    [[TMP1:%.*]] = shl i256 [[X]], 128
+; CHECK-NEXT:    [[SIGNEXT1:%.*]] = ashr exact i256 [[TMP1]], 128
 ; CHECK-NEXT:    ret i256 [[SIGNEXT1]]
 ;
   %signext1 = call i256 @llvm.evm.signextend(i256 15, i256 %x)
@@ -18,8 +19,8 @@ define i256 @test_const(i256 %x) {
 define i256 @test_const_ne(i256 %x) {
 ; CHECK-LABEL: define i256 @test_const_ne(
 ; CHECK-SAME: i256 [[X:%.*]]) {
-; CHECK-NEXT:    [[SIGNEXT1:%.*]] = call i256 @llvm.evm.signextend(i256 15, i256 [[X]])
-; CHECK-NEXT:    [[SIGNEXT2:%.*]] = call i256 @llvm.evm.signextend(i256 10, i256 [[SIGNEXT1]])
+; CHECK-NEXT:    [[TMP1:%.*]] = shl i256 [[X]], 168
+; CHECK-NEXT:    [[SIGNEXT2:%.*]] = ashr exact i256 [[TMP1]], 168
 ; CHECK-NEXT:    ret i256 [[SIGNEXT2]]
 ;
   %signext1 = call i256 @llvm.evm.signextend(i256 15, i256 %x)