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

Code before:
        "mismatch between dynamic and static basis (kDynamic marker but no "
        "corresponding dynamic basis entry) -- this can only happen due to an "
        "incorrect fold/rewrite");

  return success();
}

OpFoldResult AffineLinearizeIndexOp::fold(FoldAdaptor adaptor) {
  std::optional<SmallVector<int64_t>> maybeStaticBasis =
      foldCstValueToCstAttrBasis(getMixedBasis(), getDynamicBasisMutable(),
                                 adaptor.getDynamicBasis());
  if (maybeStaticBasis) {
    setStaticBasis(*maybeStaticBasis);
    return getResult();
  }
  // No indices linearizes to zero.
  if (getMultiIndex().empty())
    return IntegerAttr::get(getResult().getType(), 0);

  // One single index linearizes to itself.
  if (getMultiIndex().size() == 1)
    return getMultiIndex().front();

  // Return nullptr if any multi-index attribute has not been folded to a
  // concrete integer (e.g. it is still a runtime value or has folded to a
  // non-integer such as ub.PoisonAttr).
  if (llvm::any_of(adaptor.getMultiIndex(), [](Attribute a) {
        return !isa_and_nonnull<IntegerAttr>(a);
      }))
    return nullptr;

  if (!adaptor.getDynamicBasis().empty())
    return nullptr;

  int64_t result = 0;
  int64_t stride = 1;
  for (auto [length, indexAttr] :
       llvm::zip_first(llvm::reverse(getStaticBasis()),
                       llvm::reverse(adaptor.getMultiIndex()))) {
    result = result + cast<IntegerAttr>(indexAttr).getInt() * stride;
    stride = stride * length;
  }
  // Handle the index element with no basis element.
  if (!hasOuterBound())
    result =
        result +
        cast<IntegerAttr>(adaptor.getMultiIndex().front()).getInt() * stride;

  return IntegerAttr::get(getResult().getType(), result);
}


Review comment (suggestion):
```suggestion
  // non-integer attribute such as #ub.poison).
```

Code after:
        "mismatch between dynamic and static basis (kDynamic marker but no "
        "corresponding dynamic basis entry) -- this can only happen due to an "
        "incorrect fold/rewrite");

  return success();
}

OpFoldResult AffineLinearizeIndexOp::fold(FoldAdaptor adaptor) {
  std::optional<SmallVector<int64_t>> maybeStaticBasis =
      foldCstValueToCstAttrBasis(getMixedBasis(), getDynamicBasisMutable(),
                                 adaptor.getDynamicBasis());
  if (maybeStaticBasis) {
    setStaticBasis(*maybeStaticBasis);
    return getResult();
  }
  // No indices linearizes to zero.
  if (getMultiIndex().empty())
    return IntegerAttr::get(getResult().getType(), 0);

  // One single index linearizes to itself.
  if (getMultiIndex().size() == 1)
    return getMultiIndex().front();

  // Return nullptr if any multi-index attribute has not been folded to a
  // concrete integer (e.g. it is still a runtime value or has folded to a
  // non-integer attribute such as #ub.poison).
  if (llvm::any_of(adaptor.getMultiIndex(), [](Attribute a) {
        return !isa_and_nonnull<IntegerAttr>(a);
      }))
    return nullptr;

  if (!adaptor.getDynamicBasis().empty())
    return nullptr;

  int64_t result = 0;
  int64_t stride = 1;
  for (auto [length, indexAttr] :
       llvm::zip_first(llvm::reverse(getStaticBasis()),
                       llvm::reverse(adaptor.getMultiIndex()))) {
    result = result + cast<IntegerAttr>(indexAttr).getInt() * stride;
    stride = stride * length;
  }
  // Handle the index element with no basis element.
  if (!hasOuterBound())
    result =
        result +
        cast<IntegerAttr>(adaptor.getMultiIndex().front()).getInt() * stride;

  return IntegerAttr::get(getResult().getType(), result);
}