p << " outs(" << outputs << " : " << outputs.getTypes() << ")";
}
-static void printCommonStructuredOpPartsWithNewLine(OpAsmPrinter &p,
- ValueRange inputs,
- ValueRange outputs) {
- if (!inputs.empty()) {
- p << " ins(" << inputs << " : " << inputs.getTypes() << ")";
- }
- if (!outputs.empty()) {
- p << " outs(" << outputs << " : " << outputs.getTypes() << ")";
- }
-}
//===----------------------------------------------------------------------===//
// Specific parsing and printing for named structured ops created by ods-gen.
//===----------------------------------------------------------------------===//
inputs, /*outputs=*/{}, bodyBuild);
}
+static void addBodyWithPayloadOp(OpAsmParser &parser, OperationState &result,
+ const OperationName &payloadOpName,
+ const NamedAttrList &payloadOpAttrs,
+ ArrayRef<Value> operands) {
+ OpBuilder b(parser.getContext());
+ Region *body = result.addRegion();
+ Block &block = body->emplaceBlock();
+ b.setInsertionPointToStart(&block);
+ SmallVector<Value> bbArgs;
+ for (auto &operand : operands) {
+ block.addArgument(operand.getType().cast<ShapedType>().getElementType(),
+ b.getUnknownLoc());
+ }
+
+ Operation *payloadOp = b.create(
+ result.location, b.getStringAttr(payloadOpName.getStringRef()),
+ block.getArguments(),
+ TypeRange{
+ result.operands.back().getType().cast<ShapedType>().getElementType()},
+ payloadOpAttrs);
+
+ b.create<YieldOp>(result.location, payloadOp->getResults());
+}
+
ParseResult MapOp::parse(OpAsmParser &parser, OperationState &result) {
+ std::optional<OperationName> payloadOpName;
+ NamedAttrList payloadOpAttrs;
+ if (succeeded(parser.parseOptionalLBrace())) {
+ FailureOr<OperationName> operationName = parser.parseCustomOperationName();
+ if (failed(operationName))
+ return failure();
+ if (parser.parseOptionalAttrDict(payloadOpAttrs))
+ return failure();
+ payloadOpName = operationName.value();
+ if (parser.parseRBrace())
+ return failure();
+ }
+
if (parseDstStyleOp(parser, result))
return failure();
- SmallVector<OpAsmParser::Argument> regionArgs;
- if (parser.parseArgumentList(regionArgs, OpAsmParser::Delimiter::Paren,
- /*allowType=*/true, /*allowAttrs=*/true)) {
- return failure();
+ if (payloadOpName.has_value()) {
+ addBodyWithPayloadOp(parser, result, payloadOpName.value(), payloadOpAttrs,
+ makeArrayRef(result.operands).drop_back());
+ } else {
+ SmallVector<OpAsmParser::Argument> regionArgs;
+ if (parser.parseArgumentList(regionArgs, OpAsmParser::Delimiter::Paren,
+ /*allowType=*/true, /*allowAttrs=*/true)) {
+ return failure();
+ }
+ Region *body = result.addRegion();
+ if (parser.parseRegion(*body, regionArgs))
+ return failure();
}
+ return success();
+}
- Region *body = result.addRegion();
- if (parser.parseRegion(*body, regionArgs))
- return failure();
+// Retrieve the operation from the body, if it is the only one (except
+// yield) and if it gets the same amount of arguments as the body does.
+static Operation *findPayloadOp(Block *body) {
+ if (body->getOperations().size() != 2)
+ return nullptr;
+ Operation &payload = body->getOperations().front();
+ assert(isa<YieldOp>(body->getOperations().back()));
+
+ if (payload.getNumOperands() == 0 ||
+ payload.getNumOperands() != body->getNumArguments())
+ return nullptr;
+ for (const auto &[bbArg, operand] :
+ llvm::zip(payload.getOperands(), body->getArguments())) {
+ if (bbArg != operand)
+ return nullptr;
+ }
+ return &payload;
+}
- return success();
+void printShortForm(OpAsmPrinter &p, Operation *payloadOp) {
+ SmallVector<StringRef> elidedAttrs;
+ p << " { " << payloadOp->getName().getStringRef();
+ for (const auto &attr : payloadOp->getAttrs()) {
+ auto fastAttr = attr.getValue().dyn_cast<mlir::arith::FastMathFlagsAttr>();
+ if (fastAttr && fastAttr.getValue() == mlir::arith::FastMathFlags::none) {
+ elidedAttrs.push_back(attr.getName().str());
+ break;
+ }
+ }
+ p.printOptionalAttrDict(payloadOp->getAttrs(), elidedAttrs);
+ p << " }";
}
void MapOp::print(OpAsmPrinter &p) {
- printCommonStructuredOpPartsWithNewLine(
- p, SmallVector<Value>(getDpsInputOperands()),
- SmallVector<Value>(getDpsInitOperands()));
- p.printOptionalAttrDict((*this)->getAttrs());
+ Block *mapper = getBody();
+ Operation *payloadOp = findPayloadOp(mapper);
+ if (payloadOp) {
+ printShortForm(p, payloadOp);
+ }
- p.increaseIndent();
- p.printNewline();
- p << "(";
- llvm::interleaveComma(getMapper().getArguments(), p,
- [&](auto arg) { p.printRegionArgument(arg); });
- p << ") ";
+ printCommonStructuredOpParts(p, SmallVector<Value>(getDpsInputOperands()),
+ SmallVector<Value>(getDpsInitOperands()));
+ p.printOptionalAttrDict((*this)->getAttrs());
- p.printRegion(getMapper(), /*printEntryBlockArgs=*/false);
- p.decreaseIndent();
+ if (!payloadOp) {
+ // Print region if the payload op was not detected.
+ p.increaseIndent();
+ p.printNewline();
+ p << "(";
+ llvm::interleaveComma(mapper->getArguments(), p,
+ [&](auto arg) { p.printRegionArgument(arg); });
+ p << ") ";
+
+ p.printRegion(getMapper(), /*printEntryBlockArgs=*/false);
+ p.decreaseIndent();
+ }
}
LogicalResult MapOp::verify() {
"mapper, but got: "
<< getInputs().size() << " and " << blockArgs.size();
- // The parameters of mapper should all match the element type // of inputs.
+ // The parameters of mapper should all match the element type of inputs.
for (const auto &[bbArgType, inputArg] :
llvm::zip(bodyBlock->getArgumentTypes(), getInputs())) {
auto inputElemType = inputArg.getType().cast<ShapedType>().getElementType();
}
ParseResult ReduceOp::parse(OpAsmParser &parser, OperationState &result) {
+ std::optional<OperationName> payloadOpName;
+ NamedAttrList payloadOpAttrs;
+ if (succeeded(parser.parseOptionalLBrace())) {
+ FailureOr<OperationName> operationName = parser.parseCustomOperationName();
+ if (failed(operationName))
+ return failure();
+ if (parser.parseOptionalAttrDict(payloadOpAttrs))
+ return failure();
+ payloadOpName = operationName.value();
+ if (parser.parseRBrace())
+ return failure();
+ }
+
if (parseDstStyleOp(
parser, result, [&](OpAsmParser &parser, NamedAttrList &attributes) {
return parseDenseI64ArrayAttr(parser, attributes, "dimensions");
}))
return failure();
- SmallVector<OpAsmParser::Argument> regionArgs;
- if (parser.parseArgumentList(regionArgs, OpAsmParser::Delimiter::Paren,
- /*allowType=*/true, /*allowAttrs=*/true)) {
- return failure();
- }
+ if (payloadOpName.has_value()) {
+ addBodyWithPayloadOp(parser, result, payloadOpName.value(), payloadOpAttrs,
+ makeArrayRef(result.operands));
+ } else {
+ SmallVector<OpAsmParser::Argument> regionArgs;
+ if (parser.parseArgumentList(regionArgs, OpAsmParser::Delimiter::Paren,
+ /*allowType=*/true, /*allowAttrs=*/true)) {
+ return failure();
+ }
- Region *body = result.addRegion();
- if (parser.parseRegion(*body, regionArgs))
- return failure();
+ Region *body = result.addRegion();
+ if (parser.parseRegion(*body, regionArgs))
+ return failure();
+ }
return success();
}
}
void ReduceOp::print(OpAsmPrinter &p) {
- printCommonStructuredOpPartsWithNewLine(
- p, SmallVector<Value>(getDpsInputOperands()),
- SmallVector<Value>(getDpsInitOperands()));
+ Block *mapper = getBody();
+ Operation *payloadOp = findPayloadOp(mapper);
+ if (payloadOp) {
+ printShortForm(p, payloadOp);
+ }
+ printCommonStructuredOpParts(p, SmallVector<Value>(getDpsInputOperands()),
+ SmallVector<Value>(getDpsInitOperands()));
printDenseI64ArrayAttr(p, getDimensionsAttrName(), getDimensions());
p.printOptionalAttrDict((*this)->getAttrs(), {getDimensionsAttrName()});
-
- p.increaseIndent();
- p.printNewline();
- p << "(";
- llvm::interleaveComma(getCombiner().getArguments(), p,
- [&](auto arg) { p.printRegionArgument(arg); });
- p << ") ";
-
- p.printRegion(getCombiner(), /*printEntryBlockArgs=*/false);
- p.decreaseIndent();
+ if (!payloadOp) {
+ // Print region if the payload op was not detected.
+ p.increaseIndent();
+ p.printNewline();
+ p << "(";
+ llvm::interleaveComma(mapper->getArguments(), p,
+ [&](auto arg) { p.printRegionArgument(arg); });
+ p << ") ";
+
+ p.printRegion(getCombiner(), /*printEntryBlockArgs=*/false);
+ p.decreaseIndent();
+ }
}
LogicalResult ReduceOp::verify() {
}
void TransposeOp::print(OpAsmPrinter &p) {
- printCommonStructuredOpPartsWithNewLine(
- p, SmallVector<Value>(getDpsInputOperands()),
- SmallVector<Value>(getDpsInitOperands()));
+ printCommonStructuredOpParts(p, SmallVector<Value>(getDpsInputOperands()),
+ SmallVector<Value>(getDpsInitOperands()));
printDenseI64ArrayAttr(p, getPermutationAttrName(), getPermutation());
p.printOptionalAttrDict((*this)->getAttrs(), {getPermutationAttrName()});
}
}
void BroadcastOp::print(OpAsmPrinter &p) {
- printCommonStructuredOpPartsWithNewLine(
- p, SmallVector<Value>(getDpsInputOperands()),
- SmallVector<Value>(getDpsInitOperands()));
+ printCommonStructuredOpParts(p, SmallVector<Value>(getDpsInputOperands()),
+ SmallVector<Value>(getDpsInitOperands()));
printDenseI64ArrayAttr(p, getDimensionsAttrName(), getDimensions());
p.printOptionalAttrDict((*this)->getAttrs(), {getDimensionsAttrName()});
}
func.return %add : tensor<64xf32>
}
// CHECK-LABEL: func @map_binary
-// CHECK: linalg.map ins
+// CHECK: linalg.map { arith.addf } ins
// CHECK-SAME: outs
-// CHECK-NEXT: (%{{.*}}: f32, %{{.*}}: f32) {
-// CHECK-NEXT: arith.addf
-// CHECK-NEXT: linalg.yield
-// CHECK-NEXT: }
// -----
func.return %reduce : tensor<16x64xf32>
}
// CHECK-LABEL: func @reduce
-// CHECK: linalg.reduce ins
+// CHECK: linalg.reduce { arith.addf } ins
// CHECK-SAME: outs
// CHECK-SAME: dimensions = [1]
-// CHECK-NEXT: (%{{.*}}: f32, %{{.*}}: f32) {
-// CHECK-NEXT: arith.addf
-// CHECK-NEXT: linalg.yield
-// CHECK-NEXT: }
// -----
}
func.return
}
-// CHECK-LABEL: func @reduce_memref
-// CHECK: linalg.reduce
+// CHECK-LABEL: func @reduce
+// CHECK: linalg.reduce { arith.addf } ins
+// CHECK-SAME: outs
+// CHECK-SAME: dimensions = [1]
// -----
}
// CHECK-LABEL: func @variadic_reduce
// CHECK: linalg.reduce
+// CHECK-NOT: { arith.addf
// -----
}
func.return
}
-// CHECK-LABEL: func @variadic_reduce_memref
+// CHECK-LABEL: func @variadic_reduce_memref
// CHECK: linalg.reduce
+// CHECK-NOT: { arith.addf
// -----
// CHECK: linalg.broadcast ins
// CHECK-SAME: outs
// CHECK-SAME: dimensions
+
+// -----
+
+func.func @map_arith_with_attr(%lhs: tensor<64xf32>, %rhs: tensor<64xf32>,
+ %init: tensor<64xf32>) -> tensor<64xf32> {
+ %add = linalg.map
+ ins(%lhs, %rhs: tensor<64xf32>, tensor<64xf32>)
+ outs(%init:tensor<64xf32>)
+ (%lhs_elem: f32, %rhs_elem: f32) {
+ %0 = arith.addf %lhs_elem, %rhs_elem fastmath<fast> : f32
+ linalg.yield %0: f32
+ }
+ func.return %add : tensor<64xf32>
+}
+
+// CHECK-LABEL: func @map_arith_with_attr
+// CHECK-NEXT: %[[MAPPED:.*]] = linalg.map
+// CHECK-SAME: { arith.addf {fastmath = #arith.fastmath<fast>} }
+// CHECK-SAME: ins
+// CHECK-SAME: outs
+// CHECK-NEXT: return %[[MAPPED]] : tensor<64xf32>
+
+// -----
+
+func.func @reduce_arith_with_attr(%input: tensor<16x32x64xf32>,
+ %init: tensor<16x64xf32>) -> tensor<16x64xf32> {
+ %reduce = linalg.reduce
+ ins(%input:tensor<16x32x64xf32>)
+ outs(%init:tensor<16x64xf32>)
+ dimensions = [1]
+ (%in: f32, %out: f32) {
+ %0 = arith.addf %in, %out fastmath<fast> : f32
+ linalg.yield %0: f32
+ }
+ func.return %reduce : tensor<16x64xf32>
+}
+// CHECK-LABEL: func @reduce_arith_with_attr
+// CHECK-NEXT: %[[REDUCED:.*]] = linalg.reduce
+// CHECK-SAME: { arith.addf {fastmath = #arith.fastmath<fast>} }
+// CHECK-SAME: ins
+// CHECK-SAME: outs
+// CHECK-SAME: dimensions = [1]
+// CHECK-NEXT: return %[[REDUCED]] : tensor<16x64xf32>