[Mlir-commits] [mlir] [OpenMP][MLIR] Add omp.canonical_loop operation, !omp.cli type, omp.new_cli operation (PR #71712)
Jan Leyonberg
llvmlistbot at llvm.org
Wed Nov 8 10:01:42 PST 2023
https://github.com/jsjodin updated https://github.com/llvm/llvm-project/pull/71712
>From e894603dea5363ee4855348ed7d590b403874896 Mon Sep 17 00:00:00 2001
From: Jan Leyonberg <jan_sjodin at yahoo.com>
Date: Wed, 8 Nov 2023 11:23:00 -0600
Subject: [PATCH] [OpenMP][MLIR] Add omp.canonical_loop operation, !omp.cli
type, omp.new_cli operation
This patch adds the omp.canonical_loop to represent canonical loops in
OpenMP. The operation is important in order to represent loops that
can be modified using loop transformation directives in OpenMP. It
also allows simpler code generation by using utilities in the
OpenMPIRBuilder. In addition, the !omp.cli type and omp.new_cli
operation are added to be able to represent dependencies between
canonical loops and future loop transformation operations.
---
mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td | 166 +++++++++++++++++-
mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp | 82 +++++++++
mlir/test/Dialect/OpenMP/cli.mlir | 56 ++++++
3 files changed, 303 insertions(+), 1 deletion(-)
create mode 100644 mlir/test/Dialect/OpenMP/cli.mlir
diff --git a/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td b/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
index 99ac5cfb7b9e922..a485d8b061894f4 100644
--- a/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
+++ b/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
@@ -405,6 +405,170 @@ def SingleOp : OpenMP_Op<"single", [AttrSizedOperandSegments]> {
let hasVerifier = 1;
}
+//===---------------------------------------------------------------------===//
+// OpenMP Canonical Loop Info Type
+//===---------------------------------------------------------------------===//
+
+def CanonicalLoopInfoType : OpenMP_Type<"CanonicalLoopInfo", "cli"> {
+ let summary = "Type for representing a reference to a canonical loop";
+ let description = [{
+ A variable of type CanonicalLoopInfo refers to an OpenMP-compatible
+ canonical loop in the same function. Variables of this type are not
+ available at runtime and therefore cannot be used by the program itself,
+ i.e. an opaque type. It is similar to the transform dialect's
+ `!transform.interface` type, but instead of implementing an interface
+ for each transformation, the OpenMP dialect itself defines possible
+ operations on this type.
+
+ A CanonicalLoopInfo variable can
+
+ 1. passed to omp.canonical_loop to assiciate the loop to that variable
+ 2. passed to omp operations that take a CanonicalLoopInfo argument,
+ such as `omp.unroll`.
+
+ A CanonicalLoopInfo variable can not
+
+ 1. be returned from a function,
+ 2. passed to operations that are not specifically designed to take a
+ CanonicalLoopInfo, including AnyType.
+
+ A CanonicalLoopInfo variable directly corresponds to an object of
+ OpenMPIRBuilder's CanonicalLoopInfo struct when lowering to LLVM-IR.
+ }];
+}
+
+//===---------------------------------------------------------------------===//
+// OpenMP Canonical Loop Info Operation
+//===---------------------------------------------------------------------===//
+
+def NewCliOp : OpenMP_Op<"new_cli"> {
+ let summary = "Create a new Canonical Loop Info value.";
+ let description = [{
+ Create a new CLI that can be passed as an argument to a CanonicalLoopOp
+ and to loop transformation operations to handle dependencies between
+ loop transformation operations.
+ }];
+ let results = (outs CanonicalLoopInfoType:$result);
+ let assemblyFormat = [{
+ attr-dict `:` type($result)
+ }];
+}
+
+
+//===---------------------------------------------------------------------===//
+// OpenMP Canonical Loop Operation
+//===---------------------------------------------------------------------===//
+def CanonicalLoopOp : OpenMP_Op<"canonical_loop", [SingleBlockImplicitTerminator<"omp::YieldOp">]> {
+ let summary = "OpenMP Canonical Loop Operation";
+ let description = [{
+ All loops that conform to OpenMP's definition of a canonical loop can be
+ simplified to a CanonicalLoopOp. In particular, there are no loop-carried
+ variables and the number of iterations it will execute is know before the
+ operation. This allows e.g. to determine the number of threads and chunks
+ the iterations space is split into before executing any iteration. More
+ restrictions may apply in cases such as (collapsed) loop nests, doacross
+ loops, etc.
+
+ The induction variable is always of the same type as the tripcount argument.
+ Since it can never be negative, tripcount is always interpreted as an
+ unsigned integer. It is the caller's responsbility to ensure the tripcount
+ is not negative when its interpretation is signed, i.e.
+ `%tripcount = max(0,%tripcount)`.
+
+ In contrast to other loop operations such as `scf.for`, the number of
+ iterations is determined by only a single variable, the trip-count. The
+ induction variable value is the logical iteration number of that iteration,
+ which OpenMP defines to be between 0 and the trip-count (exclusive).
+ Loop representation having lower-bound, upper-bound, and step-size operands,
+ require passes to do more work than necessary, including handling special
+ cases such as upper-bound smaller than lower-bound, upper-bound equal to
+ the integer type's maximal value, negative step size, etc. This complexity
+ is better only handled once by the front-end and can apply its semantics
+ for such cases while still being able to represent any kind of loop, which
+ kind of the point of a mid-end intermediate representation. User-defined
+ types such as random-access iterators in C++ could not directly be
+ represented anyway.
+
+ An optional argument to a omp.canonical_loop that can be passed in
+ is a CanonicalLoopInfo variale that can be used to refer to the canonical
+ loop to apply transformations -- such as tiling, unrolling, or
+ work-sharing -- to the loop, similar to the transform dialect but
+ with OpenMP-specific semantics.
+
+ A CanonicalLoopOp can be lowered to LLVM-IR using OpenMPIRBuilder's
+ createCanonicalLoop method.
+
+ #### Examples
+
+ Translation from lower-bound, upper-bount, step-size to trip-count.
+ ```c
+ for (int i = 3; i < 42; i+=2) {
+ B[i] = A[i];
+ }
+ ```
+
+ ```mlir
+ %lb = arith.constant 3 : i32
+ %ub = arith.constant 42 : i32
+ %step = arith.constant 2 : i32
+ %range = arith.sub %ub, %lb : i32
+ %tc = arith.div %range, %step : i32
+ omp.canonical_loop %iv : i32 in [0, %tc) {
+ %offset = arith.mul %iv, %step : i32
+ %i = arith.add %offset, %lb : i32
+ %a = load %arrA[%i] : memref<?xf32>
+ store %a, %arrB[%i] : memref<?xf32>
+ }
+ ```
+
+ Nested canonical loop with transformation.
+ ```mlir
+ %outer = omp.new_cli : !omp.cli
+ %inner = omp.new_cli : !omp.cli
+ %outer,%inner = omp.canonical_loop %iv1 : i32 in [0, %tripcount), %outer : !omp.cli{
+ %inner = omp.canonical_loop %iv2 : i32 in [0, %tc), %inner : !omp.cli {
+ %a = load %arrA[%iv1, %iv2] : memref<?x?xf32>
+ store %a, %arrB[%iv1, %iv2] : memref<?x?xf32>
+ }
+ }
+ omp.tile(%outer, %inner : !omp.cli, !omp.cli)
+ ```
+
+ Nested canonical loop with other constructs. The `omp.distribute`
+ operation has not been added yet, so this is suggested use with other
+ constructs.
+ ```mlir
+ omp.target {
+ omp.teams {
+ omp.distribute {
+ %outer = omp.new_cli : !omp.cli
+ %inner = omp.new_cli : !omp.cli
+ omp.canonical_loop %iv1 : i32 in [0, %tripcount), %outer : !omp.cli {
+ %inner = omp.canonical_loop %iv2 : i32 in [0, %tc), %inner : !omp.cli {
+ %a = load %arrA[%iv1, %iv2] : memref<?x?xf32>
+ store %a, %arrB[%iv1, %iv2] : memref<?x?xf32>
+ }
+ }
+ omp.collapse(%outer, %inner)
+ }
+ }
+ }
+ ```
+
+ }];
+ let hasCustomAssemblyFormat = 1;
+ let hasVerifier = 1;
+
+ let arguments = (ins IntLikeType:$tripCount,
+ Optional<CanonicalLoopInfoType>:$cli);
+
+ let regions = (region AnyRegion:$region);
+
+ let extraClassDeclaration = [{
+ ::mlir::Value getInductionVar();
+ }];
+}
+
//===----------------------------------------------------------------------===//
// 2.9.2 Workshare Loop Construct
//===----------------------------------------------------------------------===//
@@ -619,7 +783,7 @@ def SimdLoopOp : OpenMP_Op<"simdloop", [AttrSizedOperandSegments,
def YieldOp : OpenMP_Op<"yield",
[Pure, ReturnLike, Terminator,
ParentOneOf<["WsLoopOp", "ReductionDeclareOp",
- "AtomicUpdateOp", "SimdLoopOp"]>]> {
+ "AtomicUpdateOp", "SimdLoopOp", "CanonicalLoopOp"]>]> {
let summary = "loop yield and termination operation";
let description = [{
"omp.yield" yields SSA values from the OpenMP dialect op region and
diff --git a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
index 480af0e1307c158..6faff0d531e90cb 100644
--- a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
+++ b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
@@ -1551,6 +1551,88 @@ LogicalResult DataBoundsOp::verify() {
return success();
}
+//===----------------------------------------------------------------------===//
+// CanonicaLoopOp
+//===----------------------------------------------------------------------===//
+
+Value mlir::omp::CanonicalLoopOp::getInductionVar() {
+ return getRegion().getArgument(0);
+}
+
+void mlir::omp::CanonicalLoopOp::print(OpAsmPrinter &p) {
+ p << " " << getInductionVar() << " : " << getInductionVar().getType()
+ << " in [0, " << getTripCount() << ")";
+ if (getCli()) {
+ p << ", " << getCli() << " : " << getCli().getType();
+ }
+ p << " ";
+
+ // omp.yield is implicit if no arguments passed to it.
+ p.printRegion(getRegion(), /*printEntryBlockArgs=*/false,
+ /*printBlockTerminators=*/false);
+
+ p.printOptionalAttrDict((*this)->getAttrs());
+}
+
+mlir::ParseResult
+mlir::omp::CanonicalLoopOp::parse(::mlir::OpAsmParser &parser,
+ ::mlir::OperationState &result) {
+ Builder &builder = parser.getBuilder();
+
+ // We derive the type of tripCount from inductionVariable. Unfortunatelty we
+ // cannot do the other way around because MLIR requires the type of tripCount
+ // to be known when calling resolveOperand.
+ OpAsmParser::Argument inductionVariable;
+ if (parser.parseArgument(inductionVariable, /*allowType*/ true) ||
+ parser.parseKeyword("in") || parser.parseLSquare())
+ return failure();
+
+ int zero = -1;
+ SMLoc zeroLoc = parser.getCurrentLocation();
+ if (parser.parseInteger(zero))
+ return failure();
+ if (zero != 0) {
+ parser.emitError(zeroLoc, "Logical iteration space starts with zero");
+ return failure();
+ }
+
+ OpAsmParser::UnresolvedOperand tripcount;
+ if (parser.parseComma() || parser.parseOperand(tripcount) ||
+ parser.parseRParen() ||
+ parser.resolveOperand(tripcount, inductionVariable.type, result.operands))
+ return failure();
+
+ OpAsmParser::UnresolvedOperand cli;
+ Type type;
+ if (succeeded(parser.parseOptionalComma()))
+ if (parser.parseOperand(cli) || parser.parseColonType(type) ||
+ parser.resolveOperand(cli, type, result.operands))
+ return failure();
+
+ // Parse the loop body.
+ Region *region = result.addRegion();
+ if (parser.parseRegion(*region, {inductionVariable}))
+ return failure();
+ CanonicalLoopOp::ensureTerminator(*region, builder, result.location);
+
+ // Parse the optional attribute list.
+ if (parser.parseOptionalAttrDict(result.attributes))
+ return failure();
+
+ return mlir::success();
+}
+
+LogicalResult CanonicalLoopOp::verify() {
+ Value indVar = getInductionVar();
+ Value tripCount = getTripCount();
+
+ if (indVar.getType() != tripCount.getType())
+ return emitOpError(
+ "Region argument must be the same type as the trip count");
+
+ return success();
+}
+
#define GET_ATTRDEF_CLASSES
#include "mlir/Dialect/OpenMP/OpenMPOpsAttributes.cpp.inc"
diff --git a/mlir/test/Dialect/OpenMP/cli.mlir b/mlir/test/Dialect/OpenMP/cli.mlir
new file mode 100644
index 000000000000000..a397b442fd61d96
--- /dev/null
+++ b/mlir/test/Dialect/OpenMP/cli.mlir
@@ -0,0 +1,56 @@
+// RUN: mlir-opt %s | mlir-opt | FileCheck %s
+
+// CHECK-LABEL: @omp_canonloop_raw
+// CHECK-SAME: (%[[tc:.*]]: i32)
+func.func @omp_canonloop_raw(%tc : i32) -> () {
+ // CHECK: omp.canonical_loop %{{.*}} : i32 in [0, %[[tc]]) {
+ "omp.canonical_loop" (%tc) ({
+ ^bb0(%iv: i32):
+ omp.yield
+ }) : (i32) -> ()
+ return
+}
+
+// CHECK-LABEL: @omp_nested_canonloop_raw
+// CHECK-SAME: (%[[tc_outer:.*]]: i32, %[[tc_inner:.*]]: i32)
+func.func @omp_nested_canonloop_raw(%tc_outer : i32, %tc_inner : i32) -> () {
+ // CHECK: %[[outer_cli:.*]] = omp.new_cli : !omp.cli
+ %outer = "omp.new_cli" () : () -> (!omp.cli)
+ // CHECK: %[[inner_cli:.*]] = omp.new_cli : !omp.cli
+ %inner = "omp.new_cli" () : () -> (!omp.cli)
+ // CHECK: omp.canonical_loop %{{.*}} : i32 in [0, %[[tc_outer]]), %[[outer_cli]] : !omp.cli {
+ "omp.canonical_loop" (%tc_outer, %outer) ({
+ ^bb_outer(%iv_outer: i32):
+ // CHECK: omp.canonical_loop %{{.*}} : i32 in [0, %[[tc_inner]]), %[[inner_cli]] : !omp.cli {
+ "omp.canonical_loop" (%tc_inner, %inner) ({
+ ^bb_inner(%iv_inner: i32):
+ omp.yield
+ }) : (i32, !omp.cli) -> ()
+ omp.yield
+ }) : (i32, !omp.cli) -> ()
+ return
+}
+
+// CHECK-LABEL: @omp_canonloop_pretty
+// CHECK-SAME: (%[[tc:.*]]: i32)
+func.func @omp_canonloop_pretty(%tc : i32) -> () {
+ // CHECK: omp.canonical_loop %[[iv:.*]] : i32 in [0, %[[tc]]) {
+ omp.canonical_loop %iv : i32 in [0, %tc) {
+ // CHECK-NEXT: %{{.*}} = llvm.add %[[iv]], %[[iv]] : i32
+ %newval = llvm.add %iv, %iv: i32
+ }
+ return
+}
+
+// CHECK-LABEL: @omp_canonloop_nested_pretty
+func.func @omp_canonloop_nested_pretty(%tc : i32) -> () {
+ // CHECK: %[[cli:.*]] = omp.new_cli : !omp.cli
+ %cli = omp.new_cli : !omp.cli
+ // CHECK: omp.canonical_loop %{{.*}} : i32 in [0, %{{.*}}), %[[cli]] : !omp.cli {
+ omp.canonical_loop %iv1 : i32 in [0, %tc), %cli : !omp.cli {
+ // CHECK: omp.canonical_loop %{{.*}} : i32 in [0, %{{.*}}) {
+ omp.canonical_loop %iv2 : i32 in [0, %tc) {}
+ }
+ return
+}
+
More information about the Mlir-commits
mailing list