iree-org / iree

A retargetable MLIR-based machine learning compiler and runtime toolkit.
http://iree.dev/
Apache License 2.0
2.48k stars 553 forks source link

Generic utility for mapping elements regardless of the backing type #15442

Open saienduri opened 8 months ago

saienduri commented 8 months ago

Request description

At the moment, we don't have a generic mapping utility for ElementsAttr. This would be very helpful as utility like this is needed for DenseResource, etc. at the MLIR level. An example of where this can be utilized:

struct ConstConverterTensor final
    : OpConversionPattern<mlir::stablehlo::ConstantOp> {
  using OpConversionPattern::OpConversionPattern;

  LogicalResult
  matchAndRewrite(mlir::stablehlo::ConstantOp constOp, OpAdaptor /*adaptor*/,
                  ConversionPatternRewriter &rewriter) const override {
    auto replacementType =
        getTypeConverter()->convertType<ShapedType>(constOp.getType());
    if (!replacementType)
      return rewriter.notifyMatchFailure(constOp, "type conversion failed");

    ElementsAttr replacementAttr = constOp.getValue();
    if (replacementType == constOp.getType()) {
      rewriter.replaceOpWithNewOp<arith::ConstantOp>(constOp, replacementType,
                                                     replacementAttr);
      return success();
    } else {
      auto denseAttr = dyn_cast<DenseElementsAttr>(constOp.getValue());
      if (!denseAttr) {
        return rewriter.notifyMatchFailure(
            constOp,
            "DenseElementsAttr cast failed (only DenseElementsAttr supported)");
      }
      // Signedness conversion.
      replacementAttr = denseAttr.mapValues(replacementType.getElementType(),
                                            [](const APInt &i) { return i; });
      rewriter.replaceOpWithNewOp<arith::ConstantOp>(constOp, replacementType,
                                                     replacementAttr);
      return success();
    }
  }
};

Here, we currently throw an error as only DenseElementsAttr has this utility.

What component(s) does this issue relate to?

Compiler

Additional context

No response

stellaraccident commented 8 months ago

For the backlog / good first issue.