Repository: llvm/llvm-project
File: mlir/lib/Dialect/Vector/Transforms/VectorUnroll.cpp (C++)

Code before:
private:
  vector::UnrollVectorOptions options;
};

/// Checks whether extractShape is a contiguous slice of shape.
/// For extractShape to be contiguous in shape:
/// 1) All but the leading dimension of extractShape and shape must match
/// exactly. 2) The total number of elements in shape must be evenly divisible
/// by
///    the total number of elements in extractShape.
/// Examples:
///   isContiguous([4, 4], [8, 4]) == true
///   isContiguous([2, 4], [8, 4]) == true
///   isContiguous([2, 2], [8, 4]) == false
/// Removes leading unit dimensions to handle cases like:
///   isContiguous([1, 16], [1, 32]) == true
static bool isContiguous(ArrayRef<int64_t> extractShape,
                         ArrayRef<int64_t> shape) {

  if (extractShape.empty() || extractShape.size() > shape.size())
    return false;

  while (extractShape.size() > 1 && extractShape.front() == 1)
    extractShape = extractShape.drop_front();

  while (!shape.empty() && shape.front() == 1) {
    shape = shape.drop_front();
  }

  size_t rankDiff = shape.size() - extractShape.size();
  if (!llvm::equal(extractShape.drop_front(), shape.drop_front(rankDiff + 1)))
    return false;

  int64_t extractElements = ShapedType::getNumElements(extractShape);
  int64_t shapeElements = ShapedType::getNumElements(shape);
  return shapeElements % extractElements == 0;
}

/// Determines what shape to use with `vector.extract_strided_slice` to extract
/// a contiguous memory region from a source vector. The extraction must be
/// contiguous and contain exactly the specified number of elements. If such an
/// extraction shape cannot be determined, returns std::nullopt.
/// EXAMPLE 1:
///   sourceShape = [16], targetElements = 8
///   Working right-to-left:
///   - Take min(8, 16) = 8 from only dim → extractShape = [8],
///     remaining = 8/8 = 1
///     Result: [8]
///
///  EXAMPLE 2:
///   sourceShape = [4, 4], targetElements = 8

Review comment (refactor):
Why do we allow `shape` to become empty if `extractShape`  is now not allowed to be empty? Wouldn't that create issues for `size_t rankDiff` (e.g., `0 - 1`)?

Code after:
private:
  vector::UnrollVectorOptions options;
};

/// Checks whether extractShape is a contiguous slice of shape.
/// For extractShape to be contiguous in shape:
/// 1) All but the leading dimension of extractShape and shape must match
/// exactly. 2) The total number of elements in shape must be evenly divisible
/// by
///    the total number of elements in extractShape.
/// Examples:
///   isContiguous([4, 4], [8, 4]) == true
///   isContiguous([2, 4], [8, 4]) == true
///   isContiguous([2, 2], [8, 4]) == false
/// Removes leading unit dimensions to handle cases like:
///   isContiguous([1, 16], [1, 32]) == true
static bool isContiguous(ArrayRef<int64_t> extractShape,
                         ArrayRef<int64_t> shape) {

  if (extractShape.empty() || shape.empty() ||
      extractShape.size() > shape.size())
    return false;

  while (extractShape.size() > 1 && extractShape.front() == 1)
    extractShape = extractShape.drop_front();

  while (shape.size() > 1 && shape.front() == 1) {
    shape = shape.drop_front();
  }

  size_t rankDiff = shape.size() - extractShape.size();
  if (!llvm::equal(extractShape.drop_front(), shape.drop_front(rankDiff + 1)))
    return false;

  int64_t extractElements = ShapedType::getNumElements(extractShape);
  int64_t shapeElements = ShapedType::getNumElements(shape);
  return shapeElements % extractElements == 0;
}

/// Determines what shape to use with `vector.extract_strided_slice` to extract
/// a contiguous memory region from a source vector. The extraction must be
/// contiguous and contain exactly the specified number of elements. If such an
/// extraction shape cannot be determined, returns std::nullopt.
/// EXAMPLE 1:
///   sourceShape = [16], targetElements = 8
///   Working right-to-left:
///   - Take min(8, 16) = 8 from only dim → extractShape = [8],
///     remaining = 8/8 = 1
///     Result: [8]
///
///  EXAMPLE 2: