Repository: llvm/llvm-project
File: llvm/lib/Target/AMDGPU/AMDGPUPromoteAlloca.cpp (C++)

Code before:
  }

  assert(CurPtr == Alloca && "GEP not based on alloca");

  int64_t VecElemSize = DL.getTypeAllocSize(VecElemTy);
  if (VarOffsets.size() > 1)
    return {};

  APInt IndexQuot;
  int64_t Rem;
  APInt::sdivrem(ConstOffset, VecElemSize, IndexQuot, Rem);
  if (Rem != 0)
    return {};

  GEPToVectorIndex Result;

  if (!ConstOffset.isZero())
    Result.ConstIndex = ConstantInt::get(Ctx, IndexQuot.sextOrTrunc(BW));

  if (VarOffsets.empty())
    return Result;

  const auto &VarOffset = VarOffsets.front();
  APInt OffsetQuot;
  APInt::sdivrem(VarOffset.second, VecElemSize, OffsetQuot, Rem);
  uint64_t DivForVarIndex = 0;
  if (Rem != 0) {
    uint64_t Scale = VarOffset.second.getZExtValue();
    if (Scale == 0 || (VecElemSize % Scale) != 0)
      return {};

    DivForVarIndex = VecElemSize / Scale;
    if (!isPowerOf2_64(DivForVarIndex))
      return {};

    KnownBits KB = computeKnownBits(VarOffset.first, DL);
    if (KB.countMinTrailingZeros() < Log2_64(DivForVarIndex))
      return {};
  } else if (OffsetQuot.isZero()) {
    return {};
  }

  Result.VarIndex = VarOffset.first;
  auto *OffsetType = dyn_cast<IntegerType>(Result.VarIndex->getType());
  if (!OffsetType)
    return {};

  if (Rem != 0)
    Result.VarShift = ConstantInt::get(Ctx, APInt(BW, Log2_64(DivForVarIndex)));
  else if (!OffsetQuot.isOne())
    Result.VarMul = ConstantInt::get(Ctx, OffsetQuot.sextOrTrunc(BW));

Review comment (suggestion):
I think we just have two cases for the index calculation `Offset*Scale/VecElemSize`. Let's make code clearer.
```
if (Scale >= VecElemSize) {
  divide the Scale by VecElemSize to get the OffsetQuot
  we only set VarMul if applicable
  return here
}
here we should have VecElemSize > Scale
divide the VecElemSize by Scale to get the divisor
check if the Offset is divisable by the divisor
only set VarShift here.
```


Code after:
    CurPtr = CurGEP->getPointerOperand();
  }

  assert(CurPtr == Alloca && "GEP not based on alloca");

  int64_t VecElemSize = DL.getTypeAllocSize(VecElemTy);
  if (VarOffsets.size() > 1)
    return {};

  // We support vector indices of the form ((VarIndex * stride) >> shift) + B.
  // IndexQuot represents B. Check that the constant offset is a multiple
  // of the vector element size.
  if (ConstOffset.srem(VecElemSize) != 0)
    return {};
  APInt IndexQuot = ConstOffset.sdiv(VecElemSize);

  GEPToVectorIndex Result;

  if (!ConstOffset.isZero())
    Result.ConstIndex = ConstantInt::get(Ctx, IndexQuot.sextOrTrunc(BW));

  // If there are no variable offsets, only a constant offset, then we're done.
  if (VarOffsets.empty())
    return Result;

  // Scale is the stride in the (A * stride) part. Check that there is only one
  // variable offset and extract the scale factor.
  const auto &VarOffset = VarOffsets.front();
  auto ScaleOpt = VarOffset.second.tryZExtValue();
  if (!ScaleOpt || *ScaleOpt == 0)
    return {};

  uint64_t Scale = *ScaleOpt;
  Result.VarIndex = VarOffset.first;
  auto *OffsetType = dyn_cast<IntegerType>(Result.VarIndex->getType());
  if (!OffsetType)
    return {};

  // The vector index for the variable part is: VarIndex * Scale / VecElemSize.
  if (Scale >= (uint64_t)VecElemSize) {
    if (Scale % VecElemSize != 0)
      return {};

    // Scale is a multiple of VecElemSize, so the index is just: VarIndex *
    // (Scale / VecElemSize).
    uint64_t VarMul = Scale / VecElemSize;
    // Only the multiplier is needed.
    if (VarMul != 1)
      Result.VarMul = ConstantInt::get(Ctx, APInt(BW, VarMul));
  } else {
    if ((uint64_t)VecElemSize % Scale != 0)