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

Code before:
  const unsigned Size = DL.getTypeStoreSize(AI->getAllocatedType());
  return I->getOperand(0) == AI &&
         match(I->getOperand(2), m_SpecificInt(Size)) && !I->isVolatile();
}

static Value *calculateVectorIndex(Value *Ptr, AllocaAnalysis &AA) {
  IRBuilder<> B(Ptr->getContext());

  Ptr = Ptr->stripPointerCasts();
  if (Ptr == AA.Alloca)
    return B.getInt32(0);

  auto *GEP = cast<GetElementPtrInst>(Ptr);
  auto I = AA.Vector.GEPVectorIdx.find(GEP);
  assert(I != AA.Vector.GEPVectorIdx.end() && "Must have entry for GEP!");

  if (!I->second.Full) {
    Value *Result = nullptr;
    B.SetInsertPoint(GEP);

    if (I->second.VarIndex) {
      Result = I->second.VarIndex;
      Result = B.CreateSExtOrTrunc(Result, B.getInt32Ty());

      if (I->second.VarDiv)
        Result = B.CreateExactSDiv(Result, I->second.VarDiv);

      if (I->second.VarMul)
        Result = B.CreateMul(Result, I->second.VarMul);
    }

    if (I->second.ConstIndex) {
      if (Result)
        Result = B.CreateAdd(Result, I->second.ConstIndex);
      else
        Result = I->second.ConstIndex;
    }

    if (!Result)
      Result = B.getInt32(0);

    I->second.Full = Result;
  }

  return I->second.Full;
}

static std::optional<GEPToVectorIndex>
computeGEPToVectorIndex(GetElementPtrInst *GEP, AllocaInst *Alloca,
                        Type *VecElemTy, const DataLayout &DL) {
  // TODO: Extracting a "multiple of X" from a GEP might be a useful generic

Review comment (suggestion):
I have to double check, but if your `getelementptr` has the `nuw` flag; you should be able to handle all your indices as `unsigned` (a negative would wrap).

Then you can use `zext` and specially `udiv exact`. The advantage is that `udiv` doesn't have to consider that the negatives round-up instead of down.

For example, for `a/4`:

On signed we get
```
        v_ashrrev_i32_e32 v1, 31, v0
        v_lshrrev_b32_e32 v1, 30, v1
        v_add_i32_e32 v0, vcc, v0, v1
        v_ashrrev_i32_e32 v0, 2, v0
```
vs
```
        v_lshrrev_b32_e32 v0, 2, v0
```

https://godbolt.org/z/cez9hzG7r

Code after:
  // only affects a subset of the vector.
  const unsigned Size = DL.getTypeStoreSize(AI->getAllocatedType());
  return I->getOperand(0) == AI &&
         match(I->getOperand(2), m_SpecificInt(Size)) && !I->isVolatile();
}

static Value *calculateVectorIndex(Value *Ptr, AllocaAnalysis &AA) {
  IRBuilder<> B(Ptr->getContext());

  Ptr = Ptr->stripPointerCasts();
  if (Ptr == AA.Alloca)
    return B.getInt32(0);

  auto *GEP = cast<GetElementPtrInst>(Ptr);
  auto I = AA.Vector.GEPVectorIdx.find(GEP);
  assert(I != AA.Vector.GEPVectorIdx.end() && "Must have entry for GEP!");

  if (!I->second.Full) {
    Value *Result = nullptr;
    B.SetInsertPoint(GEP);

    if (I->second.VarIndex) {
      Result = I->second.VarIndex;
      Result = B.CreateSExtOrTrunc(Result, B.getInt32Ty());

      if (I->second.VarMul)
        Result = B.CreateMul(Result, I->second.VarMul);

      if (I->second.VarShift)
        Result = B.CreateAShr(Result, I->second.VarShift, "", /*isExact*/ true);
    }

    if (I->second.ConstIndex) {
      if (Result)
        Result = B.CreateAdd(Result, I->second.ConstIndex);
      else
        Result = I->second.ConstIndex;
    }

    if (!Result)
      Result = B.getInt32(0);

    I->second.Full = Result;
  }

  return I->second.Full;
}

static std::optional<GEPToVectorIndex>
computeGEPToVectorIndex(GetElementPtrInst *GEP, AllocaInst *Alloca,
                        Type *VecElemTy, const DataLayout &DL) {