Repository: llvm/llvm-project
File: mlir/lib/Dialect/Affine/IR/AffineOps.cpp (C++)

Code before:
        "corresponding dynamic basis entry) -- this can only happen due to an "
        "incorrect fold/rewrite");

  if (!llvm::all_of(staticBasis, [](int64_t v) {
        return v > 0 || ShapedType::isDynamic(v);
      }))
    return emitOpError("no basis element may be statically non-positive");

  return success();
}

/// Given mixed basis of affine.delinearize_index/linearize_index replace
/// constant SSA values with the constant integer value and return the new
/// static basis. In case no such candidate for replacement exists, this utility
/// returns std::nullopt.
static std::optional<SmallVector<int64_t>>
foldCstValueToCstAttrBasis(ArrayRef<OpFoldResult> mixedBasis,
                           MutableOperandRange mutableDynamicBasis,
                           ArrayRef<Attribute> dynamicBasis) {
  uint64_t dynamicBasisIndex = 0;
  for (Attribute basis : dynamicBasis) {
    // Skip poison values: they don't have a concrete integer value, so erasing
    // them from the dynamic operands would create an inconsistency between
    // the static basis (which would still hold kDynamic) and the dynamic
    // operand list (which would be one element shorter).
    if (basis && !isa<ub::PoisonAttr>(basis)) {
      mutableDynamicBasis.erase(dynamicBasisIndex);
    } else {
      ++dynamicBasisIndex;
    }
  }

  // No constant SSA value exists.
  if (dynamicBasisIndex == dynamicBasis.size())
    return std::nullopt;

  SmallVector<int64_t> staticBasis;
  for (OpFoldResult basis : mixedBasis) {
    std::optional<int64_t> basisVal = getConstantIntValue(basis);
    if (!basisVal)
      staticBasis.push_back(ShapedType::kDynamic);
    else
      staticBasis.push_back(*basisVal);
  }

  return staticBasis;
}

LogicalResult
AffineDelinearizeIndexOp::fold(FoldAdaptor adaptor,
                               SmallVectorImpl<OpFoldResult> &result) {

Review comment (refactor):
Should this check be `basis && isa<IntegerAttr, IndexAttr>(basis)`? It's possible that a custom op has an `i64` result, but folds to a custom attribute from which we cannot extract an integer.

Code after:
        "corresponding dynamic basis entry) -- this can only happen due to an "
        "incorrect fold/rewrite");

  if (!llvm::all_of(staticBasis, [](int64_t v) {
        return v > 0 || ShapedType::isDynamic(v);
      }))
    return emitOpError("no basis element may be statically non-positive");

  return success();
}

/// Given mixed basis of affine.delinearize_index/linearize_index replace
/// constant SSA values with the constant integer value and return the new
/// static basis. In case no such candidate for replacement exists, this utility
/// returns std::nullopt.
static std::optional<SmallVector<int64_t>>
foldCstValueToCstAttrBasis(ArrayRef<OpFoldResult> mixedBasis,
                           MutableOperandRange mutableDynamicBasis,
                           ArrayRef<Attribute> dynamicBasis) {
  uint64_t dynamicBasisIndex = 0;
  for (Attribute basis : dynamicBasis) {
    // Skip poison values: they don't have a concrete integer value, so erasing
    // them from the dynamic operands would create an inconsistency between
    // the static basis (which would still hold kDynamic) and the dynamic
    // operand list (which would be one element shorter).
    if (basis && isa<IntegerAttr>(basis)) {
      mutableDynamicBasis.erase(dynamicBasisIndex);
    } else {
      ++dynamicBasisIndex;
    }
  }

  // No constant SSA value exists.
  if (dynamicBasisIndex == dynamicBasis.size())
    return std::nullopt;

  SmallVector<int64_t> staticBasis;
  for (OpFoldResult basis : mixedBasis) {
    std::optional<int64_t> basisVal = getConstantIntValue(basis);
    if (!basisVal)
      staticBasis.push_back(ShapedType::kDynamic);
    else
      staticBasis.push_back(*basisVal);
  }

  return staticBasis;
}

LogicalResult
AffineDelinearizeIndexOp::fold(FoldAdaptor adaptor,
                               SmallVectorImpl<OpFoldResult> &result) {