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

[Vedant Paranjape 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)
----------------
VedantParanjape wrote:

Um, yeah. I didn't think of that. I guess because getOptimizableIVOf jumps through the trunc->add->iv chain anyways?

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