Handle negative numbers ext after constant sinking

Handle negative numbers ext after constant sinking

Author: skarczew

IGC/Compiler/CISACodeGen/PushAnalysis.cpp

@@ -172,8 +172,8 @@ bool PushAnalysis::IsStatelessCBLoad(llvm::Instruction *inst, int &pushableAddre
     if (pAdd && pAdd->getOpcode() == llvm::Instruction::Add) {
       GetPotentialPushableAddresses(pAdd->getOperand(0));
       GetPotentialPushableAddresses(pAdd->getOperand(1));
-    } else if (isa<ZExtInst>(pAddress)) {
-      GetPotentialPushableAddresses(cast<ZExtInst>(pAddress)->getOperand(0));
+    } else if (isa<ZExtInst>(pAddress) || isa<SExtInst>(pAddress)) {
+      GetPotentialPushableAddresses(cast<Instruction>(pAddress)->getOperand(0));
     } else if (isa<ConstantInt>(pAddress)) {
       ConstantInt *pConst = cast<ConstantInt>(pAddress);
       offset += pConst->getZExtValue();

IGC/Compiler/Optimizer/SinkPointerConstAdd.cpp

@@ -16,6 +16,8 @@ SPDX-License-Identifier: MIT
 #include "llvm/IR/IRBuilder.h"
 #include "common/LLVMWarningsPop.hpp"
 
+#include "GenISAIntrinsics/GenIntrinsicInst.h"
+
 // Simple pass which sinks constant add operations in pointer calculations
 // It changes following pattern:
 // %addr_part1 = <base1> + const
@@ -41,7 +43,9 @@ class SinkPointerConstAddPass : public llvm::FunctionPass {
   static char ID;
 
 private:
-  bool getConstantOffset(llvm::Value *value, int &offset);
+  bool getConstantOffset(llvm::Value *value, std::vector<llvm::Instruction *> &zexts, int &offset);
+  void zextToSext(std::vector<llvm::Instruction *> &zexts);
+  bool skipZextToSext(llvm::Instruction *op, llvm::BasicBlock *parentBB);
 };
 
 #define PASS_FLAG "igc-sink-ptr-const-add"
@@ -50,7 +54,8 @@ class SinkPointerConstAddPass : public llvm::FunctionPass {
 #define PASS_ANALYSIS false
 IGC_INITIALIZE_PASS(SinkPointerConstAddPass, PASS_FLAG, PASS_DESCRIPTION, PASS_CFG_ONLY, PASS_ANALYSIS)
 
-bool SinkPointerConstAddPass::getConstantOffset(llvm::Value *value, int &offset) {
+bool SinkPointerConstAddPass::getConstantOffset(llvm::Value *value, std::vector<llvm::Instruction *> &zexts,
+                                                int &offset) {
   // Recursively search for constant add operations - this will stop after the first const add found,
   // and should be called repeatedly until no more const adds can be sunk.
 
@@ -65,7 +70,11 @@ bool SinkPointerConstAddPass::getConstantOffset(llvm::Value *value, int &offset)
       llvm::Instruction *op = llvm::dyn_cast<llvm::Instruction>(instr->getOperand(0));
       // This is a simple pass, only sink within the same basic block
       if (op && instr->getParent() == op->getParent()) {
-        return getConstantOffset(instr->getOperand(0), offset);
+        // Collect zext instructions for later processing
+        if (instr->getOpcode() == llvm::Instruction::ZExt) {
+          zexts.push_back(instr);
+        }
+        return getConstantOffset(instr->getOperand(0), zexts, offset);
       } else {
         return false;
       }
@@ -89,8 +98,8 @@ bool SinkPointerConstAddPass::getConstantOffset(llvm::Value *value, int &offset)
         llvm::Instruction *op0 = llvm::dyn_cast<llvm::Instruction>(instr->getOperand(0));
         llvm::Instruction *op1 = llvm::dyn_cast<llvm::Instruction>(instr->getOperand(1));
         // This is a simple pass, only sink within the same basic block
-        return (op0 && instr->getParent() == op0->getParent() && getConstantOffset(op0, offset)) ? true
-               : (op1 && instr->getParent() == op1->getParent()) ? getConstantOffset(op1, offset)
+        return (op0 && instr->getParent() == op0->getParent() && getConstantOffset(op0, zexts, offset)) ? true
+               : (op1 && instr->getParent() == op1->getParent()) ? getConstantOffset(op1, zexts, offset)
                                                                  : false;
       }
     }
@@ -99,6 +108,51 @@ bool SinkPointerConstAddPass::getConstantOffset(llvm::Value *value, int &offset)
   return false;
 }
 
+bool SinkPointerConstAddPass::skipZextToSext(llvm::Instruction *op, llvm::BasicBlock *parentBB) {
+  // This is a simple pass, only sink within the same basic block
+  if (op && parentBB == op->getParent()) {
+    // Do not change zext of pushed constants or loaded values - UMD provides unsigned offsets
+    if (llvm::GenIntrinsicInst *instr = llvm::dyn_cast<llvm::GenIntrinsicInst>(op)) {
+      if (instr->getIntrinsicID() == llvm::GenISAIntrinsic::GenISA_RuntimeValue) {
+        return true;
+      }
+    } else if (llvm::dyn_cast<llvm::Argument>(op) || llvm::dyn_cast<llvm::LoadInst>(op)) {
+      return true;
+    } else if (llvm::BinaryOperator *bo = llvm::dyn_cast<BinaryOperator>(op)) {
+      llvm::ConstantInt *cOp0 = llvm::dyn_cast<llvm::ConstantInt>(bo->getOperand(0));
+      llvm::ConstantInt *cOp1 = llvm::dyn_cast<llvm::ConstantInt>(bo->getOperand(1));
+      if ((cOp0 && cOp0->isNegative()) || (cOp1 && cOp1->isNegative())) {
+        return false;
+      } else if (bo->getOpcode() == llvm::Instruction::Sub) {
+        return false;
+      } else {
+        return (skipZextToSext(llvm::dyn_cast<llvm::Instruction>(bo->getOperand(0)), parentBB) &&
+                skipZextToSext(llvm::dyn_cast<llvm::Instruction>(bo->getOperand(1)), parentBB));
+      }
+    } else {
+      return false;
+    }
+  }
+  return true;
+}
+
+void SinkPointerConstAddPass::zextToSext(std::vector<llvm::Instruction *> &zexts) {
+  // Remove duplicates
+  std::sort(zexts.begin(), zexts.end());
+  zexts.erase(std::unique(zexts.begin(), zexts.end()), zexts.end());
+  // Convert zext instructions to sext instructions
+  for (auto &zext : zexts) {
+    llvm::Instruction *op = llvm::dyn_cast<llvm::Instruction>(zext->getOperand(0));
+    if (skipZextToSext(op, zext->getParent())) {
+      continue;
+    }
+    llvm::IRBuilder<> builder(zext);
+    llvm::Value *sext = builder.CreateSExt(zext->getOperand(0), zext->getType());
+    zext->replaceAllUsesWith(sext);
+    zext->eraseFromParent();
+  }
+}
+
 bool SinkPointerConstAddPass::runOnFunction(llvm::Function &F) {
   bool changed = false;
   std::vector<llvm::IntToPtrInst *> intToPtrInsts;
@@ -115,11 +169,20 @@ bool SinkPointerConstAddPass::runOnFunction(llvm::Function &F) {
 
   for (auto &intrinsic : intToPtrInsts) {
     int offset = 0;
+    bool localChanged = false;
+    std::vector<llvm::Instruction *> zexts;
     // Keep sinking constant adds until no more can be sunk
-    while (getConstantOffset(intrinsic->getOperand(0), offset)) {
+    while (getConstantOffset(intrinsic->getOperand(0), zexts, offset)) {
+      localChanged = true;
       changed = true;
     }
 
+    if (localChanged) {
+      // In some cases, sinking constant add may introduce negative values in pointer calculations.
+      // Convert affected zext instructions to sext instructions to avoid potential issues.
+      zextToSext(zexts);
+    }
+
     // If we found any constant offset, create new pointer calculation
     if (offset != 0) {
       llvm::IRBuilder<> builder(intrinsic);