[llvm] [VPlan] Narrow truncates of an induction plus a constant offset (PR #226035)

[Ramkumar Ramachandra via llvm-commits]

================
@@ -5954,16 +5952,36 @@ void VPlanTransforms::narrowInductionTruncates(VPlan &Plan, VFRange &Range,
         continue;
 
       VPValue *Op = VPI.getOperand(0);
-      auto *WideIV = getOptimizableIVOf(Op, PSE);
-      if (!WideIV)
-        continue;
 
-      // getOptimizableIVOf also matches an add of the IV and its step, which
-      // is not handled here.
-      // TODO: Also narrow truncates of the incremented IV.
-      if (Op != WideIV)
+      // Look through a constant offset, which is folded into the start value
+      // of the narrowed induction below. The offset need not be the induction
+      // step, as start + C + i * step stays affine for any constant C.
+      VPValue *IVOp = Op;
+      APInt Offset;
+      VPValue *X;
+      const APInt *C;
+      if (match(Op, m_c_Add(m_VPValue(X), m_APInt(C)))) {
+        IVOp = X;
+        Offset = *C;
+      } else if (match(Op, m_Sub(m_VPValue(X), m_APInt(C)))) {
+        IVOp = X;
+        Offset = -*C;
+      }
+
+      auto *WideIV = getOptimizableIVOf(IVOp, PSE);
+      if (!WideIV || WideIV != IVOp)
         continue;
 
+      // The offset is added in the induction's type, so it can only be folded
+      // into a start value of that same type.
+      VPValue *Start = WideIV->getStartValue();
+      if (IVOp != Op) {
+        const APInt *StartC;
+        if (!match(Start, m_APInt(StartC)))
+          continue;
+        Start = Plan.getConstantInt(*StartC + Offset);
+      }
----------------
artagnon wrote:

Perhaps getOptimizableIV should return a pair of WideIV + new Start? I think there are also other users of getOptimizableIVOf that could benefit?

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