[llvm] [AMDGPU][CodeGenPrepare] Narrow 64 bit math to 32 bit if profitable (PR #130577)

[Nikita Popov via llvm-commits]

================
@@ -8594,6 +8594,68 @@ static bool optimizeBranch(BranchInst *Branch, const TargetLowering &TLI,
   return false;
 }
 
+static bool tryNarrowMathIfNoOverflow(Instruction *I,
+                                      const TargetTransformInfo *TTI,
+                                      const DataLayout &DL) {
+  unsigned Opc = I->getOpcode();
+  Type *OldType = I->getType();
+
+  if (Opc != Instruction::Add && Opc != Instruction::Mul)
+    return false;
+
+  unsigned OrigBit = OldType->getScalarSizeInBits();
+  unsigned MaxBitsNeed = OrigBit;
+  switch (Opc) {
+  case Instruction::Add:
+    MaxBitsNeed = KnownBits::add(computeKnownBits(I->getOperand(0), DL),
+                                 computeKnownBits(I->getOperand(1), DL))
+                      .countMaxActiveBits();
+    break;
+  case Instruction::Mul:
+    MaxBitsNeed = KnownBits::mul(computeKnownBits(I->getOperand(0), DL),
+                                 computeKnownBits(I->getOperand(1), DL))
+                      .countMaxActiveBits();
+    break;
+  default:
+    break;
+  }
+
+  MaxBitsNeed = std::max<unsigned>(bit_ceil(MaxBitsNeed), 8);
+
+  if (OrigBit <= MaxBitsNeed)
+    return false;
+
+  Type *NewType = I->getType()->getWithNewBitWidth(MaxBitsNeed);
+
+  // Old cost
+  InstructionCost OldCost =
+      TTI->getArithmeticInstrCost(Opc, OldType, TTI::TCK_RecipThroughput);
+  // New cost of new op
+  InstructionCost NewCost =
+      TTI->getArithmeticInstrCost(Opc, NewType, TTI::TCK_RecipThroughput);
+  // New cost of narrowing 2 operands (use trunc)
+  NewCost += TTI->getCastInstrCost(Instruction::Trunc, NewType, OldType,
+                                   TTI->getCastContextHint(I),
+                                   TTI::TCK_RecipThroughput) *
+             2;
+  // New cost of zext narrowed result to original type
+  NewCost += TTI->getCastInstrCost(Instruction::ZExt, OldType, NewType,
+                                   TTI->getCastContextHint(I),
+                                   TTI::TCK_RecipThroughput);
+  if (NewCost >= OldCost) {
+    return false;
+  }
+  IRBuilder<> Builder(I);
+  Value *Trun0 = Builder.CreateTrunc(I->getOperand(0), NewType);
----------------
nikic wrote:

```suggestion
  Value *Trunc0 = Builder.CreateTrunc(I->getOperand(0), NewType);
```

https://github.com/llvm/llvm-project/pull/130577