From 842b6861c01cc6961f170d58332ecf0fb0232441 Mon Sep 17 00:00:00 2001 From: Stanislav Funiak Date: Fri, 26 Nov 2021 17:57:30 +0530 Subject: [PATCH] Defines new PDLInterp operations needed for multi-root matching in PDL. This is commit 1 of 4 for the multi-root matching in PDL, discussed in https://llvm.discourse.group/t/rfc-multi-root-pdl-patterns-for-kernel-matching/4148 (topic flagged for review). These operations are: * pdl.get_accepting_ops: Returns a list of operations accepting the given value or a range of values at the specified position. Thus if there are two operations `%op1 = "foo"(%val)` and `%op2 = "bar"(%val)` accepting a value at position 0, `%ops = pdl_interp.get_accepting_ops of %val : !pdl.value at 0` will return both of them. This allows us to traverse upwards from a value to operations accepting the value. * pdl.choose_op: Iteratively chooses one operation from a range of operations. Therefore, writing `%op = pdl_interp.choose_op from %ops` in the example above will select either `%op1`or `%op2`. Testing: Added the corresponding test cases to mlir/test/Dialect/PDLInterp/ops.mlir. Reviewed By: rriddle Differential Revision: https://reviews.llvm.org/D108543 --- .../mlir/Dialect/PDLInterp/IR/PDLInterpOps.td | 137 +++++++++++++++++++++ mlir/lib/Dialect/PDLInterp/IR/PDLInterp.cpp | 79 ++++++++++++ mlir/test/Dialect/PDLInterp/ops.mlir | 43 +++++++ 3 files changed, 259 insertions(+) diff --git a/mlir/include/mlir/Dialect/PDLInterp/IR/PDLInterpOps.td b/mlir/include/mlir/Dialect/PDLInterp/IR/PDLInterpOps.td index ff9b3dd..87033ed 100644 --- a/mlir/include/mlir/Dialect/PDLInterp/IR/PDLInterpOps.td +++ b/mlir/include/mlir/Dialect/PDLInterp/IR/PDLInterpOps.td @@ -371,6 +371,29 @@ def PDLInterp_CheckTypesOp } //===----------------------------------------------------------------------===// +// pdl_interp::ContinueOp +//===----------------------------------------------------------------------===// + +def PDLInterp_ContinueOp + : PDLInterp_Op<"continue", [NoSideEffect, HasParent<"ForEachOp">, + Terminator]> { + let summary = "Breaks the current iteration"; + let description = [{ + `pdl_interp.continue` operation breaks the current iteration within the + `pdl_interp.foreach` region and continues with the next iteration from + the beginning of the region. + + Example: + + ```mlir + pdl_interp.continue + ``` + }]; + + let assemblyFormat = "attr-dict"; +} + +//===----------------------------------------------------------------------===// // pdl_interp::CreateAttributeOp //===----------------------------------------------------------------------===// @@ -514,6 +537,42 @@ def PDLInterp_EraseOp : PDLInterp_Op<"erase"> { } //===----------------------------------------------------------------------===// +// pdl_interp::ExtractOp +//===----------------------------------------------------------------------===// + +def PDLInterp_ExtractOp + : PDLInterp_Op<"extract", [NoSideEffect, + TypesMatchWith< + "`range` is a PDL range whose element type matches type of `result`", + "result", "range", "pdl::RangeType::get($_self)">]> { + let summary = "Extract the item at the specified index in a range"; + let description = [{ + `pdl_interp.extract` operations are used to extract an item from a range + at the specified index. If the index is out of range, returns null. + + Example: + + ```mlir + // Extract the value at index 1 from a range of values. + %ops = pdl_interp.extract 1 of %values : !pdl.value + ``` + }]; + + let arguments = (ins PDL_RangeOf:$range, + Confined:$index); + let results = (outs PDL_AnyType:$result); + let assemblyFormat = "$index `of` $range `:` type($result) attr-dict"; + + let builders = [ + OpBuilder<(ins "Value":$range, "unsigned":$index), [{ + build($_builder, $_state, + range.getType().cast().getElementType(), + range, index); + }]>, + ]; +} + +//===----------------------------------------------------------------------===// // pdl_interp::FinalizeOp //===----------------------------------------------------------------------===// @@ -534,6 +593,48 @@ def PDLInterp_FinalizeOp } //===----------------------------------------------------------------------===// +// pdl_interp::ForEachOp +//===----------------------------------------------------------------------===// + +def PDLInterp_ForEachOp + : PDLInterp_Op<"foreach", [Terminator]> { + let summary = "Iterates over a range of values or ranges"; + let description = [{ + `pdl_interp.foreach` iteratively selects an element from a range of values + and executes the region until pdl.continue is reached. + + In the bytecode interpreter, this operation is implemented by looping over + the values and, for each selection, running the bytecode until we reach + pdl.continue. This may result in multiple matches being reported. Note + that the input range is mutated (popped from). + + Example: + + ```mlir + pdl_interp.foreach %op : !pdl.operation in %ops { + pdl_interp.continue + } -> ^next + ``` + }]; + + let arguments = (ins PDL_RangeOf:$values); + let regions = (region AnyRegion:$region); + let successors = (successor AnySuccessor:$successor); + + let builders = [ + OpBuilder<(ins "Value":$range, "Block *":$successor, "bool":$initLoop)> + ]; + + let extraClassDeclaration = [{ + /// Returns the loop variable. + BlockArgument getLoopVariable() { return region().getArgument(0); } + }]; + let parser = [{ return ::parseForEachOp(parser, result); }]; + let printer = [{ return ::print(p, *this); }]; + let verifier = [{ return ::verify(*this); }]; +} + +//===----------------------------------------------------------------------===// // pdl_interp::GetAttributeOp //===----------------------------------------------------------------------===// @@ -751,6 +852,42 @@ def PDLInterp_GetResultsOp : PDLInterp_Op<"get_results", [NoSideEffect]> { } //===----------------------------------------------------------------------===// +// pdl_interp::GetUsersOp +//===----------------------------------------------------------------------===// + +def PDLInterp_GetUsersOp + : PDLInterp_Op<"get_users", [NoSideEffect]> { + let summary = "Get the users of a `Value`"; + let description = [{ + `pdl_interp.get_users` extracts the users that accept this value. In the + case of a range, the union of users of the all the values are returned, + similarly to ResultRange::getUsers. + + Example: + + ```mlir + // Get all the users of a single value. + %ops = pdl_interp.get_users of %value : !pdl.value + + // Get all the users of the first value in a range. + %ops = pdl_interp.get_users of %values : !pdl.range + ``` + }]; + + let arguments = (ins PDL_InstOrRangeOf:$value); + let results = (outs PDL_RangeOf:$operations); + let assemblyFormat = "`of` $value `:` type($value) attr-dict"; + + let builders = [ + OpBuilder<(ins "Value":$value), [{ + build($_builder, $_state, + pdl::RangeType::get($_builder.getType()), + value); + }]>, + ]; +} + +//===----------------------------------------------------------------------===// // pdl_interp::GetValueTypeOp //===----------------------------------------------------------------------===// diff --git a/mlir/lib/Dialect/PDLInterp/IR/PDLInterp.cpp b/mlir/lib/Dialect/PDLInterp/IR/PDLInterp.cpp index d149ef5..5a14fff 100644 --- a/mlir/lib/Dialect/PDLInterp/IR/PDLInterp.cpp +++ b/mlir/lib/Dialect/PDLInterp/IR/PDLInterp.cpp @@ -66,6 +66,85 @@ static void printCreateOperationOpAttributes(OpAsmPrinter &p, } //===----------------------------------------------------------------------===// +// pdl_interp::ForEachOp +//===----------------------------------------------------------------------===// + +void ForEachOp::build(::mlir::OpBuilder &builder, ::mlir::OperationState &state, + Value range, Block *successor, bool initLoop) { + build(builder, state, range, successor); + if (initLoop) { + // Create the block and the loop variable. + auto range_type = range.getType().cast(); + state.regions.front()->emplaceBlock(); + state.regions.front()->addArgument(range_type.getElementType()); + } +} + +static ParseResult parseForEachOp(OpAsmParser &parser, OperationState &result) { + // Parse the loop variable followed by type. + OpAsmParser::OperandType loopVariable; + Type loopVariableType; + if (parser.parseRegionArgument(loopVariable) || + parser.parseColonType(loopVariableType)) + return failure(); + + // Parse the "in" keyword. + if (parser.parseKeyword("in", " after loop variable")) + return failure(); + + // Parse the operand (value range). + OpAsmParser::OperandType operandInfo; + if (parser.parseOperand(operandInfo)) + return failure(); + + // Resolve the operand. + Type rangeType = pdl::RangeType::get(loopVariableType); + if (parser.resolveOperand(operandInfo, rangeType, result.operands)) + return failure(); + + // Parse the body region. + Region *body = result.addRegion(); + if (parser.parseRegion(*body, {loopVariable}, {loopVariableType})) + return failure(); + + // Parse the attribute dictionary. + if (parser.parseOptionalAttrDict(result.attributes)) + return failure(); + + // Parse the successor. + Block *successor; + if (parser.parseArrow() || parser.parseSuccessor(successor)) + return failure(); + result.addSuccessors(successor); + + return success(); +} + +static void print(OpAsmPrinter &p, ForEachOp op) { + BlockArgument arg = op.getLoopVariable(); + p << ' ' << arg << " : " << arg.getType() << " in " << op.values(); + p.printRegion(op.region(), /*printEntryBlockArgs=*/false); + p.printOptionalAttrDict(op->getAttrs()); + p << " -> "; + p.printSuccessor(op.successor()); +} + +static LogicalResult verify(ForEachOp op) { + // Verify that the operation has exactly one argument. + if (op.region().getNumArguments() != 1) + return op.emitOpError("requires exactly one argument"); + + // Verify that the loop variable and the operand (value range) + // have compatible types. + BlockArgument arg = op.getLoopVariable(); + Type rangeType = pdl::RangeType::get(arg.getType()); + if (rangeType != op.values().getType()) + return op.emitOpError("operand must be a range of loop variable type"); + + return success(); +} + +//===----------------------------------------------------------------------===// // pdl_interp::GetValueTypeOp //===----------------------------------------------------------------------===// diff --git a/mlir/test/Dialect/PDLInterp/ops.mlir b/mlir/test/Dialect/PDLInterp/ops.mlir index 072dfad..8bf559b 100644 --- a/mlir/test/Dialect/PDLInterp/ops.mlir +++ b/mlir/test/Dialect/PDLInterp/ops.mlir @@ -23,3 +23,46 @@ func @operations(%attribute: !pdl.attribute, pdl_interp.finalize } + +// ----- + +func @extract(%attrs : !pdl.range, %ops : !pdl.range, %types : !pdl.range, %vals: !pdl.range) { + // attribute at index 0 + %attr = pdl_interp.extract 0 of %attrs : !pdl.attribute + + // operation at index 1 + %op = pdl_interp.extract 1 of %ops : !pdl.operation + + // type at index 2 + %type = pdl_interp.extract 2 of %types : !pdl.type + + // value at index 3 + %val = pdl_interp.extract 3 of %vals : !pdl.value + + pdl_interp.finalize +} + +// ----- + +func @foreach(%ops: !pdl.range) { + // iterate over a range of operations + pdl_interp.foreach %op : !pdl.operation in %ops { + %val = pdl_interp.get_result 0 of %op + pdl_interp.continue + } -> ^end + + ^end: + pdl_interp.finalize +} + +// ----- + +func @users(%value: !pdl.value, %values: !pdl.range) { + // all the users of a single value + %ops1 = pdl_interp.get_users of %value : !pdl.value + + // all the users of all the values in a range + %ops2 = pdl_interp.get_users of %values : !pdl.range + + pdl_interp.finalize +} -- 2.7.4