Repository: llvm/llvm-project
File: llvm/include/llvm/IR/MatrixBuilder.h (C/C++)

Code before:

  /// Divide matrix \p LHS by scalar \p RHS. If the operands are integers, \p
  /// IsUnsigned indicates whether UDiv or SDiv should be used.
  Value *CreateScalarDiv(Value *LHS, Value *RHS, bool IsUnsigned) {
    assert(LHS->getType()->isVectorTy() && !RHS->getType()->isVectorTy());
    assert(!isa<ScalableVectorType>(LHS->getType()) &&
           "LHS Assumed to be fixed width");
    RHS =
        B.CreateVectorSplat(cast<VectorType>(LHS->getType())->getElementCount(),
                            RHS, "scalar.splat");
    return cast<VectorType>(LHS->getType())
                   ->getElementType()
                   ->isFloatingPointTy()
               ? B.CreateFDiv(LHS, RHS)
               : (IsUnsigned ? B.CreateUDiv(LHS, RHS) : B.CreateSDiv(LHS, RHS));
  }

  /// Create a shuffle vector that reorders elements from row-major to
  /// column-major order for a matrix with \p NumRows rows and \p NumColumns
  /// columns.
  Value *CreateRowMajorToColumnMajorShuffle(Value *Vec, unsigned NumRows,
                                            unsigned NumColumns,
                                            const Twine &Name = "") {
    SmallVector<int, 16> Mask;
    for (unsigned I = 0, N = NumRows * NumColumns; I < N; ++I)
      Mask.push_back((I % NumRows) * NumColumns + (I / NumRows));
    return B.CreateShuffleVector(Vec, Mask, Name);
  }

  /// Create an assumption that \p Idx is less than \p NumElements.
  void CreateIndexAssumption(Value *Idx, unsigned NumElements,
                             Twine const &Name = "") {
    Value *NumElts =
        B.getIntN(Idx->getType()->getScalarSizeInBits(), NumElements);
    auto *Cmp = B.CreateICmpULT(Idx, NumElts);
    if (isa<ConstantInt>(Cmp))
      assert(cast<ConstantInt>(Cmp)->isOne() && "Index must be valid!");
    else
      B.CreateAssumption(Cmp);
  }
  /// Compute the index to access the element at (\p RowIdx, \p ColumnIdx) from
  /// a matrix with \p NumRows or \p NumCols embedded in a vector depending
  /// on matrix major ordering.
  Value *CreateIndex(Value *RowIdx, Value *ColumnIdx, unsigned NumRows,
                     unsigned NumCols, bool IsMatrixRowMajor = false,
                     Twine const &Name = "") {
    unsigned MaxWidth = std::max(RowIdx->getType()->getScalarSizeInBits(),
                                 ColumnIdx->getType()->getScalarSizeInBits());
    Type *IntTy = IntegerType::get(RowIdx->getType()->getContext(), MaxWidth);
    RowIdx = B.CreateZExt(RowIdx, IntTy);
    ColumnIdx = B.CreateZExt(ColumnIdx, IntTy);

Review comment (refactor):
Can you make a private helper for this called `CreateConstantColumnMajorIndex` put it near where the  current `CreateColumnMajorIndex` is. I will need it for a futrue pr to fix this https://github.com/llvm/llvm-project/blob/a56993a694ed02775285b9fe0e23fce8346491c9/clang/lib/CodeGen/CGExprScalar.cpp#L2557-L2559

Code after:

  /// Divide matrix \p LHS by scalar \p RHS. If the operands are integers, \p
  /// IsUnsigned indicates whether UDiv or SDiv should be used.
  Value *CreateScalarDiv(Value *LHS, Value *RHS, bool IsUnsigned) {
    assert(LHS->getType()->isVectorTy() && !RHS->getType()->isVectorTy());
    assert(!isa<ScalableVectorType>(LHS->getType()) &&
           "LHS Assumed to be fixed width");
    RHS =
        B.CreateVectorSplat(cast<VectorType>(LHS->getType())->getElementCount(),
                            RHS, "scalar.splat");
    return cast<VectorType>(LHS->getType())
                   ->getElementType()
                   ->isFloatingPointTy()
               ? B.CreateFDiv(LHS, RHS)
               : (IsUnsigned ? B.CreateUDiv(LHS, RHS) : B.CreateSDiv(LHS, RHS));
  }

  /// Create an assumption that \p Idx is less than \p NumElements.
  void CreateIndexAssumption(Value *Idx, unsigned NumElements,
                             Twine const &Name = "") {
    Value *NumElts =
        B.getIntN(Idx->getType()->getScalarSizeInBits(), NumElements);
    auto *Cmp = B.CreateICmpULT(Idx, NumElts);
    if (isa<ConstantInt>(Cmp))
      assert(cast<ConstantInt>(Cmp)->isOne() && "Index must be valid!");
    else
      B.CreateAssumption(Cmp);
  }
  /// Compute the index to access the element at (\p RowIdx, \p ColumnIdx) from
  /// a matrix with \p NumRows or \p NumCols embedded in a vector depending
  /// on matrix major ordering.
  Value *CreateIndex(Value *RowIdx, Value *ColumnIdx, unsigned NumRows,
                     unsigned NumCols, bool IsMatrixRowMajor = false,
                     Twine const &Name = "") {
    unsigned MaxWidth = std::max(RowIdx->getType()->getScalarSizeInBits(),
                                 ColumnIdx->getType()->getScalarSizeInBits());
    Type *IntTy = IntegerType::get(RowIdx->getType()->getContext(), MaxWidth);
    RowIdx = B.CreateZExt(RowIdx, IntTy);
    ColumnIdx = B.CreateZExt(ColumnIdx, IntTy);
    if (IsMatrixRowMajor) {
      Value *NumColsV = B.getIntN(MaxWidth, NumCols);
      return CreateRowMajorIndex(RowIdx, ColumnIdx, NumColsV, Name);
    }
    Value *NumRowsV = B.getIntN(MaxWidth, NumRows);
    return CreateColumnMajorIndex(RowIdx, ColumnIdx, NumRowsV, Name);
  }

private:
  /// Compute the index to access the element at (\p RowIdx, \p ColumnIdx) from
  /// a matrix with \p NumRows embedded in a vector.
  Value *CreateColumnMajorIndex(Value *RowIdx, Value *ColumnIdx,