Wrong argument generation for ttnn::pow op in ttmlir-translate --mlir-to-cpp pipeline

[kmitrovicTT]

Problem description and repro

While working on #2011 I added test/ttmlir/EmitC/TTNN/eltwise_binary/power.mlir

// RUN: ttmlir-opt --ttir-to-ttnn-backend-pipeline="system-desc-path=%system_desc_path%" %s > %t.mlir
// RUN: ttmlir-translate --ttnn-to-flatbuffer %t.mlir > %basename_t.ttnn
// RUN: ttmlir-opt --ttnn-modify-signatures-for-dylib --convert-ttnn-to-emitc %t.mlir > %t2.mlir
// RUN: ttmlir-translate --mlir-to-cpp %t2.mlir > %basename_t.cpp

func.func @power(%arg0: tensor<64x128xf32>, %arg1: tensor<64x128xf32>) -> tensor<64x128xf32> {
  %0 = tensor.empty() : tensor<64x128xf32>
  %1 = "ttir.power"(%arg0, %arg1, %0) <{operandSegmentSizes = array<i32: 2, 1>}> : (tensor<64x128xf32>, tensor<64x128xf32>, tensor<64x128xf32>) -> tensor<64x128xf32>
  return %1 : tensor<64x128xf32>
}

Running ./build/bin/ttmlir-opt --ttir-to-emitc-pipeline test/ttmlir/EmitC/TTNN/eltwise_binary/power.mlir yields IMO a good result

...
%10 = emitc.call_opaque "ttnn::pow"(%3, %6, %9)
...

but piping it to ./build/bin/ttmlir-translate --mlir-to-cpp returns

ttnn::Tensor v13 = ttnn::pow(v6, v9, std::nullopt, std::nullopt, v12);

which has one excessive std::nullopt argument, and thus no matching call in third_party/tt-metal/src/tt-metal/ttnn/cpp/ttnn/operations/eltwise/binary/binary_composite.hpp::ExecutePower.

I tried to fix it briefly but no luck, I am not familiar with that part of the code.

I opened this issue and commented out emitc test to get the PR going, since it was blocking an important model.

Note

ttnn has ttnn::power which is unary and ttnn::pow which is binary composite. The latter is the one TTIR_PowerOp and TTNN_PowerOp map to.

[svuckovicTT]

If I run

./build/bin/ttmlir-opt --ttir-to-emitc-pipeline test/ttmlir/EmitC/TTNN/eltwise_binary/power.mlir

I do get

%10 = emitc.call_opaque "ttnn::pow"(%3, %6, %9) {args = [0 : index, 1 : index, #emitc.opaque<"std::nullopt">, #emitc.opaque<"std::nullopt">, 2 : index]} : (!emitc.opaque<"ttnn::Tensor">, !emitc.opaque<"ttnn::Tensor">, !emitc.opaque<"ttnn::Tensor">) -> !emitc.opaque<"ttnn::Tensor">

which is expected to turn into:

ttnn::Tensor vSomething = ttnn::pow(v6, v9, std::nullopt, std::nullopt, v12);

In order to create an emitc.call_opaque op, one must be able to use both operands and attributes as input to the function being called (think of converting a transpose(dim0, dim1) where both dim0 and dim1 are attrs). So there's a need to intermingle both operands and attributes into a single place. The way it's done in emitc is to create an ArrayAttr object, using IntegerAttr to reference operands, and plugging in Attrs themselves where needed.

Let's look at an example, I'll use concat op as it's succinct:

class ConcatOpConversionPattern
    : public TTNNToEmitCBaseOpConversionPattern<ttnn::ConcatOp> {

public:
  using TTNNToEmitCBaseOpConversionPattern<
      ttnn::ConcatOp>::TTNNToEmitCBaseOpConversionPattern;

  LogicalResult
  matchAndRewrite(ttnn::ConcatOp srcOp, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {

    // ttnn::concat op requires a `std::vector<>` of `Tensor` objects, but we
    // can't really create a `std::vector<>` with `Value` objects without
    // introducing an EmitC op that takes in these `Value` objects. We do this
    // by creating a utility function within the IR that converts a list of
    // `Tensor` objects into a `std::vector<ttnn::Tensor>`.

    ttnn_to_emitc::utils::insertVecCreateFnIfNotExists(rewriter, srcOp);

    mlir::emitc::CallOpaqueOp vectorOp = rewriter.create<emitc::CallOpaqueOp>(
        srcOp.getLoc(),
        emitc::OpaqueType::get(rewriter.getContext(),
                               "std::vector<ttnn::Tensor>"),
        ttnn_to_emitc::utils::kCreateVectorFunctionName, nullptr, nullptr,
        adaptor.getInputs());

    ArrayAttr arrayAttrs = rewriter.getArrayAttr(
        {mlir::IntegerAttr::get(rewriter.getIndexType(), 0),
         srcOp.getDimAttr()});

    rewriter.replaceOpWithNewOp<emitc::CallOpaqueOp>(
        srcOp, this->getTypeConverter()->convertType(srcOp.getType()),
        this->convertOpName(srcOp), arrayAttrs, nullptr,
        ValueRange(vectorOp->getResults()));

    return success();
  }
};

Looking at

    rewriter.replaceOpWithNewOp<emitc::CallOpaqueOp>(
        srcOp, this->getTypeConverter()->convertType(srcOp.getType()),
        this->convertOpName(srcOp), arrayAttrs, nullptr,
        ValueRange(vectorOp->getResults()));

we can see that we're piping both the arrayAttrs and ValueRange(vectorOp->getResults()) to the call_opaque op constructor.

Before that, the arrayAttrs object is created:

    ArrayAttr arrayAttrs = rewriter.getArrayAttr(
        {mlir::IntegerAttr::get(rewriter.getIndexType(), 0),
         srcOp.getDimAttr()});

So what's happening here is that we're piping both the attributes AND operands, but how will the call_opaque op know in which order to mix them in? Well, arrayAttr will contain either indices in form of an IntegetAttr which will represent a kind of "pointer" into the operands arg (in this case ValueRange(vectorOp->getResults())) or it will contain an attribute.

To explain it a bit further, I'll comment the lines:

    ArrayAttr arrayAttrs = rewriter.getArrayAttr(
        {mlir::IntegerAttr::get(rewriter.getIndexType(), 0),  // for the first argument in call_opaque op, use operand at index 0 => use first element of ValueRange(vectorOp->getResults())
         srcOp.getDimAttr()});  // for the second argument, use the dim attribute

[svuckovicTT]

In case of pow op, the EltwiseBinaryOpConversionPattern was used, which will insert the two nullopts we're seeing. However, the pow op doesn't share the signature with other eltwise binary ops, so this will not work - a new converter needs to be added.