[Mlir-commits] [mlir] [mlir] Handle arith.const expr in dispatchIndexOpFoldResult func (PR #122432)

[llvmlistbot at llvm.org]

rutkoor wrote:

> @rutkoor Can you post the C++ code that builds the `tensor.expand_shape` op? I suspect that something is wrong there.
> 
> As the name suggests, `dispatchIndexOpFoldResult` should just dispatch based on the type of the `OpFoldResult`. I.e., it's just a switch-case statement and not meant to look at the IR that it's operating on.

This is the code,
void ExpandShapeOp::build(OpBuilder &builder, OperationState &result,
                          Type resultType, Value src,
                          ArrayRef<ReassociationIndices> reassociation,
                          ArrayRef<OpFoldResult> outputShape) {
  auto [staticOutputShape, dynamicOutputShape] =
      decomposeMixedValues(SmallVector<OpFoldResult>(outputShape));
  build(builder, result, cast<RankedTensorType>(resultType), src,
        getReassociationIndicesAttribute(builder, reassociation),
        dynamicOutputShape, staticOutputShape);
}

It is invoking decomposeMixedValues function which is trying to extract integer from the expressions.

/// Decompose a vector of mixed static or dynamic values into the corresponding
/// pair of arrays. This is the inverse function of `getMixedValues`.
std::pair<SmallVector<int64_t>, SmallVector<Value>>
decomposeMixedValues(const SmallVectorImpl<OpFoldResult> &mixedValues) {
  SmallVector<int64_t> staticValues;
  SmallVector<Value> dynamicValues;
  for (const auto &it : mixedValues) {
    if (auto attr = dyn_cast<Attribute>(it)) {
      staticValues.push_back(cast<IntegerAttr>(attr).getInt());
    } else {
      staticValues.push_back(ShapedType::kDynamic);
      dynamicValues.push_back(cast<Value>(it));
    }
  }
  return {staticValues, dynamicValues};
}

Instead of invoking decomposeMixedValues function, it should invoke the dispatchIndexOpFoldResult.

https://github.com/llvm/llvm-project/pull/122432