Skip to content

[mlir] Stop restricting linalg::ReduceOp to have same number of inputs and outputs #107005

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
wants to merge 1 commit into from
Closed

[mlir] Stop restricting linalg::ReduceOp to have same number of inputs and outputs #107005

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -331,7 +331,7 @@ def MapOp : LinalgStructuredBase_Op<"map", [
def ReduceOp : LinalgStructuredBase_Op<"reduce", [
DeclareOpInterfaceMethods<OpAsmOpInterface, ["getAsmResultNames"]>,
DeclareOpInterfaceMethods<OpAsmOpInterface, ["getAsmBlockArgumentNames"]>,
SameVariadicOperandSize,
AttrSizedOperandSegments,
SingleBlockImplicitTerminator<"YieldOp">]> {
let summary = "Reduce operator";
let description = [{
Expand Down
16 changes: 11 additions & 5 deletions mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1339,11 +1339,12 @@ LogicalResult GenericOp::fold(FoldAdaptor, SmallVectorImpl<OpFoldResult> &) {
static ParseResult parseDstStyleOp(
OpAsmParser &parser, OperationState &result,
function_ref<ParseResult(OpAsmParser &, NamedAttrList &)> parseAttrsFn =
nullptr) {
nullptr,
bool addOperandSegmentSizes = false) {
// Parse `ins` and `outs`.
SmallVector<Type, 4> inputTypes, outputTypes;
if (parseCommonStructuredOpParts(parser, result, inputTypes, outputTypes,
/*addOperandSegmentSizes=*/false))
addOperandSegmentSizes))
return failure();

// Add result types.
Expand Down Expand Up @@ -1694,9 +1695,12 @@ ParseResult ReduceOp::parse(OpAsmParser &parser, OperationState &result) {
}

if (parseDstStyleOp(
parser, result, [&](OpAsmParser &parser, NamedAttrList &attributes) {
parser, result,
[&](OpAsmParser &parser, NamedAttrList &attributes) {
return parseDenseI64ArrayAttr(parser, attributes, "dimensions");
}))
},
/*addOperandSegmentSizes=*/true))

return failure();

if (payloadOpName.has_value()) {
Expand Down Expand Up @@ -1731,7 +1735,9 @@ void ReduceOp::print(OpAsmPrinter &p) {

printCommonStructuredOpParts(p, getDpsInputs(), getDpsInits());
printDenseI64ArrayAttr(p, getDimensionsAttrName(), getDimensions());
p.printOptionalAttrDict((*this)->getAttrs(), {getDimensionsAttrName()});
p.printOptionalAttrDict(
(*this)->getAttrs(),
{getDimensionsAttrName(), getOperandSegmentSizesAttrName()});
if (!payloadOp) {
// Print region if the payload op was not detected.
p.increaseIndent();
Expand Down
42 changes: 42 additions & 0 deletions mlir/test/Dialect/Linalg/roundtrip.mlir
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -485,6 +485,48 @@ func.func @variadic_reduce_memref(%input1: memref<16x32x64xf32>,

// -----

func.func @reduce_asymmetric(%input: tensor<16x32x64xi32>, %input2: tensor<16x32x64xi32>,
%init: tensor<16x64xi32>) -> tensor<16x64xi32> {
%reduce = linalg.reduce
Copy link
Contributor

@MaheshRavishankar MaheshRavishankar Nov 2, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Like I dont understand what this is trying to do. You are effectively doing a multiply and a reduce? How is this different from just using a linalg.generic. As in why use a linalg.reduce better than linalg.generic?

Copy link
Contributor Author

@oowekyala oowekyala Nov 4, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes this is a multiply-accumulate type operation. The difference is just that it's higher-level than a linalg generic. That is nice when you want to generate this op in a compiler pass, or when you want to pattern-match it. All linalg ops are just sugar over a generic, so I don't think this use case should be disqualified because it can be implemented with a generic.

ins(%input, %input2:tensor<16x32x64xi32>, tensor<16x32x64xi32>)
outs(%init:tensor<16x64xi32>)
dimensions = [1]
(%in: i32, %in2: i32, %out: i32) {
%0 = arith.muli %in, %in2: i32
%1 = arith.addi %out, %0: i32
linalg.yield %1: i32
}
func.return %reduce : tensor<16x64xi32>
}
// CHECK-LABEL: func @reduce_asymmetric
// CHECK: linalg.reduce ins(%{{.*}}, %{{.*}}: tensor<16x32x64xi32>, tensor<16x32x64xi32>)
// CHECK-NOT: operandSegmentSize
// CHECK-SAME: outs(%{{.*}}: tensor<16x64xi32>)
// CHECK-SAME: dimensions = [1]

// -----

func.func @reduce_asymmetric_memref(%input: memref<16x32x64xi32>, %input2: memref<16x32x64xi32>,
%init: memref<16x64xi32>) {
linalg.reduce
ins(%input, %input2:memref<16x32x64xi32>, memref<16x32x64xi32>)
outs(%init:memref<16x64xi32>)
dimensions = [1]
(%in: i32, %in2: i32, %out: i32) {
%0 = arith.muli %in, %in2: i32
%1 = arith.addi %out, %0: i32
linalg.yield %1: i32
}
func.return
}
// CHECK-LABEL: func @reduce_asymmetric_memref
// CHECK: linalg.reduce ins(%{{.*}}, %{{.*}}: memref<16x32x64xi32>, memref<16x32x64xi32>)
// CHECK-NOT: operandSegmentSize
// CHECK-SAME: outs(%{{.*}}: memref<16x64xi32>)
// CHECK-SAME: dimensions = [1]

// -----

func.func @transpose(%input: tensor<16x32x64xf32>,
%init: tensor<32x64x16xf32>) -> tensor<32x64x16xf32> {
%transpose = linalg.transpose
Expand Down