About ConvTransposeOpPattern2 matchAndRewrite Impl ==> error "Channel Out (M) must be a multiple of the number of groups"

[aijyo]

When running onnx-mlir.exe --EmitLib I get the error "Channel Out (M) must be a multiple of the number of groups". After inspecting the source I found that the onnx-mlir I built (commit/branch: 72b796c) performs the following in ConvTransposeOpPattern2::matchAndRewrite for ConvTranspose's weight W:

auto reversed_w = onnx_mlir::reverseWeightTensor(rewriter, odsLoc, (*w.begin())); and then passes reversed_w to ::mlir::ONNXConvOp::create.

From documentation I understand that ConvTranspose2d weight shape = (C_in, C_out / g, K_h, K_w) while Conv2d weight shape = (C_out, C_in / g, K_h, K_w). The way ConvTransposeOpPattern2::matchAndRewrite handles reversed_w looks like it may not correctly implement the corresponding transformation between these weight layouts.

Could you please review this logic and confirm whether the implementation of ConvTransposeOpPattern2 is correct? Thank you.

[aijyo]

I modified two functions in Decompose.cpp to handle the case where group > 1 in my ONNX model.
// -----------------------------------------------------------------------------
// Principle:
// - ONNX ConvTranspose weight layout: [Cin, CoutPerG, kH, kW, ...]
// - ONNX Conv weight layout: [Cout, CinPerG, kH, kW, ...]
// where:
// CoutPerG = Cout / group
// CinPerG = Cin / group
//
// The original reverseWeightTensor() assumes weight is [N, C, ...] and blindly
// swaps the first two dims at the end. This is incorrect for ConvTranspose,
// especially when group > 1 (e.g. depthwise group=Cin).
//
// Fix:
// - Keep existing spatial reversal (reverse kernel dims) unchanged.
// - Replace the final "swap non-spatial dims" with a group-aware channel
// transform:
// if group == 1: transpose [1,0,2..]
// else:
// reshape [G, CinPerG, CoutPerG, k...]
// transpose[0,2,1,3..]
// reshape [Cout, CinPerG, k...]
// -----------------------------------------------------------------------------

Value reverseWeightTensorConvTranspose(
PatternRewriter& rewriter, Location loc, Value input, mlir::IntegerAttr groupAttr) {

onnx_mlir::MultiDialectBuilder<onnx_mlir::OnnxBuilder> create(rewriter, loc);

int64_t group = 1;
if (groupAttr)
    group = groupAttr.getSInt();
llvm::outs() << "reverseWeightTensorConvTranspose group[" << group << "]\n";

ShapedType inputType = mlir::cast<ShapedType>(input.getType());
Type elementType = inputType.getElementType();
assert(inputType.hasRank() && "Need rank to reverse weight tensor.");

// ---------- Keep your original steps 1~3 to reverse spatial (kernel) dims ----------
// NOTE: This block is copied from your current reverseWeightTensor().
int64_t spatialOffset = 2; // for the first two dims
int64_t spatialRank = inputType.getRank() - spatialOffset;

// 1) Move spatial dims to the front: [A, B, D0, D1, ...] -> [D0, D1, ..., A, B]
SmallVector<int64_t, 8> permsVal;
for (int i = 0; i < spatialRank; ++i)
    permsVal.emplace_back(spatialOffset + i);
for (int i = 0; i < spatialOffset; ++i)
    permsVal.emplace_back(i);
Value transposedInput = create.onnx.transposeInt64(input, permsVal);

// 2) Reverse spatial dims pairwise (your existing logic)
ShapedType tInputType = mlir::cast<ShapedType>(transposedInput.getType());
for (int i = 0; i < spatialRank / 2; i += 2) {
    assert((!tInputType.isDynamicDim(0) && !tInputType.isDynamicDim(1)) &&
        "Spatial dimensions for weight tensor need to be static.");
    Value reverse0 = reverseAllElements(rewriter, loc, transposedInput, 0);
    Value reverse1 = reverseAllElements(rewriter, loc, reverse0, 1);

    // Move two reversed dims to the last for next reverse.
    SmallVector<int64_t, 8> permsVal0;
    for (int j = 0; j < inputType.getRank() - 2; ++j)
        permsVal0.emplace_back(j + 2);
    for (int j = 0; j < 2; ++j)
        permsVal0.emplace_back(j);
    transposedInput = create.onnx.transposeInt64(reverse1, permsVal0);
}

// 3) Reverse remaining dim if spatialRank is odd (your existing logic)
if (spatialRank % 2 != 0) {
    ShapedType tInType = mlir::cast<ShapedType>(transposedInput.getType());
    ArrayRef<int64_t> tInShape = tInType.getShape();
    Value reverse0;
    if (tInShape[1] == ShapedType::kDynamic) {
        // Keep your original dynamic workaround.
        int64_t tInRank = (int64_t)tInShape.size();
        Type tInShapeType =
            RankedTensorType::get({ tInRank }, rewriter.getI64Type());
        Value tInShapeVals = create.onnx.shape(tInShapeType, transposedInput);

        SmallVector<int64_t, 8> reshapedShapeVec;
        reshapedShapeVec.emplace_back(tInShape[0]);
        reshapedShapeVec.emplace_back(1);
        for (int i = 1; i < tInType.getRank(); ++i)
            reshapedShapeVec.emplace_back(tInShape[i]);
        Type reshapedType = RankedTensorType::get(reshapedShapeVec, elementType);

        Type firstShapeType = RankedTensorType::get({ 1 }, rewriter.getI64Type());
        Type otherShapeType =
            RankedTensorType::get({ tInRank - 1 }, rewriter.getI64Type());

        Value oneVal = create.onnx.constantInt64(ArrayRef<int64_t>({ 1 }));
        Value firstShapeVal =
            create.onnx.slice(firstShapeType, tInShapeVals, 0, 1);
        Value otherShapeVals =
            create.onnx.slice(otherShapeType, tInShapeVals, 1, tInRank);

        Type reshapeShapeType =
            RankedTensorType::get({ tInRank + 1 }, rewriter.getI64Type());
        Value shape = create.onnx.concat(
            reshapeShapeType, ValueRange{ firstShapeVal, oneVal, otherShapeVals }, 0);

        transposedInput = create.onnx.reshape(reshapedType, transposedInput, shape);
        reverse0 = reverseAllElements(rewriter, loc, transposedInput, 0);
        reverse0 = create.onnx.reshape(tInType, reverse0, tInShapeVals);
    }
    else {
        reverse0 = reverseAllElements(rewriter, loc, transposedInput, 0);
    }

    // Move reversed one dim to the last.
    SmallVector<int64_t, 8> permsVal1;
    for (int j = 0; j < inputType.getRank() - 1; ++j)
        permsVal1.emplace_back(j + 1);
    permsVal1.emplace_back(0);
    transposedInput = create.onnx.transposeInt64(reverse0, permsVal1);
}

// At this point, spatial dims are reversed, and the tensor is expected to be
// back in the original dim order of the input: [A, B, D0, D1, ...] with
// reversed spatial content.

// ---------- FIXED step 4: group-aware channel transform for ConvTranspose ----------
// For ConvTranspose weight, the first two dims are [Cin, CoutPerG], not [N, C].
// We need Conv weight [Cout, CinPerG, ...].

auto outType = mlir::cast<ShapedType>(transposedInput.getType());
assert(outType.hasRank());
int64_t rank = outType.getRank();

// Require static Cin/CoutPerG for now (matches your existing static constraints).
assert(!outType.isDynamicDim(0) && !outType.isDynamicDim(1) &&
    "Cin/CoutPerG must be static for grouped ConvTranspose decomposition.");
int64_t Cin = outType.getDimSize(0);
int64_t CoutPerG = outType.getDimSize(1);

assert(group > 0);
assert(Cin % group == 0 && "Cin must be divisible by group.");
int64_t CinPerG = Cin / group;
int64_t Cout = CoutPerG * group;

if (group == 1) {
    // Simple case: [Cin, Cout, k...] -> [Cout, Cin, k...]
    SmallVector<int64_t, 8> perm;
    perm.reserve(rank);
    perm.emplace_back(1);
    perm.emplace_back(0);
    for (int64_t i = 2; i < rank; ++i) perm.emplace_back(i);
    return create.onnx.transposeInt64(transposedInput, perm);
}

// Build reshape shape: [G, CinPerG, CoutPerG, k...]
SmallVector<int64_t, 8> shape1;
shape1.reserve(rank + 1);
shape1.push_back(group);
shape1.push_back(CinPerG);
shape1.push_back(CoutPerG);
for (int64_t i = 2; i < rank; ++i) shape1.push_back(outType.getDimSize(i));

// Type after reshape
Type reshapedType1 = RankedTensorType::get(shape1, elementType);
Value shape1Const = create.onnx.constantInt64(ArrayRef<int64_t>(shape1));
Value w5 = create.onnx.reshape(reshapedType1, transposedInput, shape1Const);

// Transpose: [G, CinPerG, CoutPerG, k...] -> [G, CoutPerG, CinPerG, k...]
SmallVector<int64_t, 8> perm5;
perm5.reserve(rank + 1);
perm5.push_back(0);
perm5.push_back(2);
perm5.push_back(1);
for (int64_t i = 3; i < rank + 1; ++i) perm5.push_back(i);

auto w5Type = mlir::cast<ShapedType>(w5.getType());
// After transpose, shape becomes:
SmallVector<int64_t, 8> shapeT = shape1;
std::swap(shapeT[1], shapeT[2]);
Type reshapedTypeT = RankedTensorType::get(shapeT, elementType);
Value w5t = create.onnx.transposeInt64(w5, perm5);

// Final reshape to Conv weight: [Cout, CinPerG, k...]
SmallVector<int64_t, 8> shape2;
shape2.reserve(rank);
shape2.push_back(Cout);
shape2.push_back(CinPerG);
for (int64_t i = 2; i < rank; ++i) shape2.push_back(outType.getDimSize(i));

Type reshapedType2 = RankedTensorType::get(shape2, elementType);
Value shape2Const = create.onnx.constantInt64(ArrayRef<int64_t>(shape2));
Value w4 = create.onnx.reshape(reshapedType2, w5t, shape2Const);

return w4;

}

// Insert pads (zeros) between elements in spatial dimensions for ConvTranspose decomposition.
// - First: interleave zeros with padSize = stride - 1 between adjacent elements.
// - Then: append extra zeros at the END of each spatial axis according to output_padding.
// This is critical because output_padding affects the computed output, and should not be
// implemented by padding the final output tensor with zeros.
Value insertPadsConvTransposeInput(PatternRewriter& rewriter, Location loc,
ONNXConvTransposeOp op, Value input) {

ONNXConvTransposeOpShapeHelper shapeHelper(op.getOperation(), {});
shapeHelper.computeShapeAndAssertOnFailure();

SmallVector<int64_t, 2> strides = shapeHelper.strides;
SmallVector<int64_t, 2> outPad = shapeHelper.outputPadding;

int64_t spatialOffset = 2;
// 1) Interleave zeros between elements: padSize = stride - 1
for (unsigned int i = 0; i < strides.size(); ++i) {
    int64_t padBetween = strides[i] - 1;
    if (padBetween > 0) {
        input = insertPadAxis(rewriter, loc, input,
            /*axis*/ spatialOffset + i,
            /*padSize*/ padBetween);
    }
}

// 2) Append zeros at the END according to output_padding
//    NOTE: This must happen on the upsampled tensor (after interleaving),
//    so the subsequent Conv (stride=1) can compute those output positions.
for (unsigned int i = 0; i < outPad.size(); ++i) {
    int64_t tailZeros = outPad[i];
    if (tailZeros > 0) {
        ArrayRef<int64_t> inShape = mlir::cast<ShapedType>(input.getType()).getShape();
        input = emitPadsAxisEnd(rewriter, loc, input, inShape,
            /*axis*/ spatialOffset + i,
            /*size*/ tailZeros);
    }
}

return input;

}

[AlexandreEichenberger]

@tungld it looks like you may have edited this code last. Do you mind just checking the logic? Thanks

[tungld]

@AlexandreEichenberger which code are you talking about? I haven't touched ConvTranspose recently.

[aijyo]

365076a, it’s this branch. At the time I used GPT to help fix the code and managed to get it to compile successfully, but I don’t clearly remember all the details anymore.

I’m not very familiar with MLIR, but I did keep the diff from that branch. If you need it, you can take a look at it. GitHub doesn’t allow uploading .diff files directly, so I uploaded it as a .txt file — you can download it, change the suffix back to .diff, and apply it to that branch.

ConvTranspose-365076adc56021da9eac6ad34fac928732a326a8.txt

@AlexandreEichenberger which code are you talking about? I haven't touched ConvTranspose recently.