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();
- 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();
+ SmallVector<OpAsmParser::Argument> regionArgs;
+ if (parser.parseArgumentList(regionArgs, OpAsmParser::Delimiter::Paren,
+ /*allowType=*/true, /*allowAttrs=*/true)) {
+ return failure();
}
- return success();
-}
-// 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;
-}
+ Region *body = result.addRegion();
+ if (parser.parseRegion(*body, regionArgs))
+ return failure();
-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 << " }";
+ return success();
}
void MapOp::print(OpAsmPrinter &p) {
- Block *mapper = getBody();
- Operation *payloadOp = findPayloadOp(mapper);
- if (payloadOp) {
- printShortForm(p, payloadOp);
- }
-
- printCommonStructuredOpParts(p, SmallVector<Value>(getDpsInputOperands()),
- SmallVector<Value>(getDpsInitOperands()));
+ printCommonStructuredOpPartsWithNewLine(
+ p, SmallVector<Value>(getDpsInputOperands()),
+ SmallVector<Value>(getDpsInitOperands()));
p.printOptionalAttrDict((*this)->getAttrs());
- 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();
- }
+ p.increaseIndent();
+ p.printNewline();
+ p << "(";
+ llvm::interleaveComma(getMapper().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();
- 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();
+ 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();
}
}
void ReduceOp::print(OpAsmPrinter &p) {
- Block *mapper = getBody();
- Operation *payloadOp = findPayloadOp(mapper);
- if (payloadOp) {
- printShortForm(p, payloadOp);
- }
+ 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()});
- 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();
- }
+
+ p.increaseIndent();
+ p.printNewline();
+ p << "(";
+ llvm::interleaveComma(getCombiner().getArguments(), p,
+ [&](auto arg) { p.printRegionArgument(arg); });
+ p << ") ";
+
+ p.printRegion(getCombiner(), /*printEntryBlockArgs=*/false);
+ p.decreaseIndent();
}
LogicalResult ReduceOp::verify() {
}
void TransposeOp::print(OpAsmPrinter &p) {
- printCommonStructuredOpParts(p, SmallVector<Value>(getDpsInputOperands()),
- SmallVector<Value>(getDpsInitOperands()));
+ printCommonStructuredOpPartsWithNewLine(
+ p, SmallVector<Value>(getDpsInputOperands()),
+ SmallVector<Value>(getDpsInitOperands()));
printDenseI64ArrayAttr(p, getPermutationAttrName(), getPermutation());
p.printOptionalAttrDict((*this)->getAttrs(), {getPermutationAttrName()});
}
}
void BroadcastOp::print(OpAsmPrinter &p) {
- printCommonStructuredOpParts(p, SmallVector<Value>(getDpsInputOperands()),
- SmallVector<Value>(getDpsInitOperands()));
+ printCommonStructuredOpPartsWithNewLine(
+ 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 { arith.addf } ins
+// CHECK: linalg.map 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 { arith.addf } ins
+// CHECK: linalg.reduce 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
-// CHECK: linalg.reduce { arith.addf } ins
-// CHECK-SAME: outs
-// CHECK-SAME: dimensions = [1]
+// CHECK-LABEL: func @reduce_memref
+// CHECK: linalg.reduce
// -----
}
// 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>