linalg.yield %0: f32
}
```
-
- Shortened print form is available. Applies to simple maps with one
- non-yield operation inside the body.
-
- The example above will be printed as:
- ```
- %add = linalg.map { arith.addf }
- ins(%lhs, %rhs : tensor<64xf32>, tensor<64xf32>)
- outs(%init: tensor<64xf32>)
- ```
}];
let arguments = (ins
outs(%init:tensor<16x64xf32>)
dimensions = [1]
(%in: f32, %out: f32) {
- %0 = arith.addf %out, %in: f32
+ %0 = arith.addf %in, %out: f32
linalg.yield %0: f32
}
```
-
- Shortened print form is available. Applies to simple (not variadic) reduces
- with one non-yield operation inside the body. Applies only if the operation
- takes `%out` as the first argument.
-
- The example above will be printed as:
- ```
- %reduce = linalg.reduce { arith.addf }
- ins(%input:tensor<16x32x64xf32>)
- outs(%init:tensor<16x64xf32>)
- dimensions = [1]
- ```
}];
let arguments = (ins
static void addBodyWithPayloadOp(OpAsmParser &parser, OperationState &result,
const OperationName &payloadOpName,
const NamedAttrList &payloadOpAttrs,
- ArrayRef<Value> operands,
- bool initFirst = false) {
+ ArrayRef<Value> operands) {
OpBuilder b(parser.getContext());
Region *body = result.addRegion();
Block &block = body->emplaceBlock();
block.addArgument(operand.getType().cast<ShapedType>().getElementType(),
b.getUnknownLoc());
}
- SmallVector<Value> payloadOpOperands;
- // If initFirst flag is enabled, we consider init as the first position of
- // payload operands.
- if (initFirst) {
- payloadOpOperands.push_back(block.getArguments().back());
- for (const auto& arg : block.getArguments().drop_back())
- payloadOpOperands.push_back(arg);
- } else {
- payloadOpOperands = {block.getArguments().begin(),
- block.getArguments().end()};
- }
Operation *payloadOp = b.create(
result.location, b.getStringAttr(payloadOpName.getStringRef()),
- payloadOpOperands,
+ block.getArguments(),
TypeRange{
result.operands.back().getType().cast<ShapedType>().getElementType()},
payloadOpAttrs);
+
b.create<YieldOp>(result.location, payloadOp->getResults());
}
// 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.
-// If initFirst flag is enabled, we check that init takes the first position in
-// operands of payload.
-static Operation *findPayloadOp(Block *body, bool initFirst = false) {
+static Operation *findPayloadOp(Block *body) {
if (body->getOperations().size() != 2)
return nullptr;
Operation &payload = body->getOperations().front();
if (payload.getNumOperands() == 0 ||
payload.getNumOperands() != body->getNumArguments())
return nullptr;
- if (initFirst) {
- // check init
- if (payload.getOperands().back() != body->getArgument(0))
+ for (const auto &[bbArg, operand] :
+ llvm::zip(payload.getOperands(), body->getArguments())) {
+ if (bbArg != operand)
return nullptr;
- // check rest
- for (int i = 1; i < body->getNumArguments(); ++i) {
- if (payload.getOperand(i - 1) != body->getArgument(i)) {
- return nullptr;
- }
- }
- } else {
- for (const auto &[bbArg, operand] :
- llvm::zip(payload.getOperands(), body->getArguments())) {
- if (bbArg != operand)
- return nullptr;
- }
}
return &payload;
}
if (payloadOpName.has_value()) {
addBodyWithPayloadOp(parser, result, payloadOpName.value(), payloadOpAttrs,
- makeArrayRef(result.operands), /*initFirst=*/true);
+ makeArrayRef(result.operands));
} else {
SmallVector<OpAsmParser::Argument> regionArgs;
if (parser.parseArgumentList(regionArgs, OpAsmParser::Delimiter::Paren,
void ReduceOp::print(OpAsmPrinter &p) {
Block *mapper = getBody();
- Operation *payloadOp = findPayloadOp(mapper, /*initFirst=*/true);
+ Operation *payloadOp = findPayloadOp(mapper);
if (payloadOp) {
printShortForm(p, payloadOp);
}
outs(%init:tensor<16x64xf32>)
dimensions = [1]
(%in: f32, %out: f32) {
- %0 = arith.addf %out, %in: f32
+ %0 = arith.addf %in, %out: f32
linalg.yield %0: f32
}
func.return %reduce : tensor<16x64xf32>
outs(%init:memref<16x64xf32>)
dimensions = [1]
(%in: f32, %out: f32) {
- %0 = arith.addf %out, %in: f32
+ %0 = arith.addf %in, %out: f32
linalg.yield %0: f32
}
func.return
outs(%init:tensor<16x64xf32>)
dimensions = [1]
(%in: f32, %out: f32) {
- %0 = arith.addf %out, %in fastmath<fast> : f32
+ %0 = arith.addf %in, %out fastmath<fast> : f32
linalg.yield %0: f32
}
func.return %reduce : tensor<16x64xf32>