[APInt] Implement average functions without sign/zero-extension. NFC. (#85212)

Removing the extension to FullWidth should make them much more efficient
in the 64-bit case, because 65-bit APInts use a separate allocation for
their bits.

1 files changed, 8 insertions, 24 deletions

diff --git a/llvm/lib/Support/APInt.cpp b/llvm/lib/Support/APInt.cpp
index 7053f3b..c530f7b 100644
--- a/llvm/lib/Support/APInt.cpp
+++ b/llvm/lib/Support/APInt.cpp
@@ -3096,37 +3096,21 @@ void llvm::LoadIntFromMemory(APInt &IntVal, const uint8_t *Src,
 }
 
 APInt APIntOps::avgFloorS(const APInt &C1, const APInt &C2) {
-  // Return floor((C1 + C2)/2)
-  assert(C1.getBitWidth() == C2.getBitWidth() && "Unequal bitwidths");
-  unsigned FullWidth = C1.getBitWidth() + 1;
-  APInt C1Ext = C1.sext(FullWidth);
-  APInt C2Ext = C2.sext(FullWidth);
-  return (C1Ext + C2Ext).extractBits(C1.getBitWidth(), 1);
+  // Return floor((C1 + C2) / 2)
+  return (C1 & C2) + (C1 ^ C2).ashr(1);
 }
 
 APInt APIntOps::avgFloorU(const APInt &C1, const APInt &C2) {
-  // Return floor((C1 + C2)/2)
-  assert(C1.getBitWidth() == C2.getBitWidth() && "Unequal bitwidths");
-  unsigned FullWidth = C1.getBitWidth() + 1;
-  APInt C1Ext = C1.zext(FullWidth);
-  APInt C2Ext = C2.zext(FullWidth);
-  return (C1Ext + C2Ext).extractBits(C1.getBitWidth(), 1);
+  // Return floor((C1 + C2) / 2)
+  return (C1 & C2) + (C1 ^ C2).lshr(1);
 }
 
 APInt APIntOps::avgCeilS(const APInt &C1, const APInt &C2) {
-  // Return ceil((C1 + C2)/2)
-  assert(C1.getBitWidth() == C2.getBitWidth() && "Unequal bitwidths");
-  unsigned FullWidth = C1.getBitWidth() + 1;
-  APInt C1Ext = C1.sext(FullWidth);
-  APInt C2Ext = C2.sext(FullWidth);
-  return (C1Ext + C2Ext + 1).extractBits(C1.getBitWidth(), 1);
+  // Return ceil((C1 + C2) / 2)
+  return (C1 | C2) - (C1 ^ C2).ashr(1);
 }
 
 APInt APIntOps::avgCeilU(const APInt &C1, const APInt &C2) {
-  // Return ceil((C1 + C2)/2)
-  assert(C1.getBitWidth() == C2.getBitWidth() && "Unequal bitwidths");
-  unsigned FullWidth = C1.getBitWidth() + 1;
-  APInt C1Ext = C1.zext(FullWidth);
-  APInt C2Ext = C2.zext(FullWidth);
-  return (C1Ext + C2Ext + 1).extractBits(C1.getBitWidth(), 1);
+  // Return ceil((C1 + C2) / 2)
+  return (C1 | C2) - (C1 ^ C2).lshr(1);
 }