[Mlir-commits] [mlir] 5e5c4ac - [mlir][OpenMP] Introduce 'omp.iterators' for OpenMP iterator modifiers (#182218)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Wed Feb 25 09:16:19 PST 2026
Author: Chi-Chun, Chen
Date: 2026-02-25T11:16:10-06:00
New Revision: 5e5c4acd4e2b039a352ba0ddd34e1da0c8bf6876
URL: https://github.com/llvm/llvm-project/commit/5e5c4acd4e2b039a352ba0ddd34e1da0c8bf6876
DIFF: https://github.com/llvm/llvm-project/commit/5e5c4acd4e2b039a352ba0ddd34e1da0c8bf6876.diff
LOG: [mlir][OpenMP] Introduce 'omp.iterators' for OpenMP iterator modifiers (#182218)
`omp.iterator` provides information of induction variables and iterator
range in OpenMP iterator modifier.
Example:
```
%it = omp.iterator(%i0: index, %i1: index) =
(%lb0 to %ub0 step %st0,
%lb1 to %ub1 step %st1) {
omp.yield(%i0, %i1 : index, index)
} -> !omp.iterated<!llvm.struct<(!llvm.ptr, i64)>>
```
Here's how we can use the omp.iterater to generate multi-dimensional
loop in llvm ir:
```
// Induction variables can be translated from the block arguments
// in omp.iterator.
// lbs, ubs, and steps is encoded in omp.iterator
for (int i0 = lbs[0]; i0 < ubs[0]; i0 += steps[0]) {
for (int i0 = lbs[1]; i1 < ubs[1]; i1 += steps[1]) {
// the result of iterated is <ptr, i64> from the example
iterated = omp.iterators(i, j);
}
}
```
1/3 in stack for implementing affinity clause with iterator modifier
1/3 #182218
2/3 #182222
3/3 #182223
Added:
Modified:
mlir/include/mlir/Dialect/OpenMP/OpenMPOpBase.td
mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
mlir/test/Dialect/OpenMP/invalid.mlir
mlir/test/Dialect/OpenMP/ops.mlir
Removed:
################################################################################
diff --git a/mlir/include/mlir/Dialect/OpenMP/OpenMPOpBase.td b/mlir/include/mlir/Dialect/OpenMP/OpenMPOpBase.td
index 5ad4e4b5b61d1..4dd8e91585a66 100644
--- a/mlir/include/mlir/Dialect/OpenMP/OpenMPOpBase.td
+++ b/mlir/include/mlir/Dialect/OpenMP/OpenMPOpBase.td
@@ -38,6 +38,13 @@ def OpenMP_MapBoundsType : OpenMP_Type<"MapBounds", "map_bounds_ty"> {
let summary = "Type for representing omp map clause bounds information";
}
+def OpenMP_IteratedType : OpenMP_Type<"Iterated", "iterated"> {
+ let summary = "OpenMP iterator-produced list handle";
+
+ let parameters = (ins "Type":$elementType);
+ let assemblyFormat = "`<` $elementType `>`";
+}
+
//===---------------------------------------------------------------------===//
// OpenMP Canonical Loop Info Type
//===---------------------------------------------------------------------===//
diff --git a/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td b/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
index 0f51b08f87dc5..3cce660b43a33 100644
--- a/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
+++ b/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
@@ -845,11 +845,11 @@ def SimdOp : OpenMP_Op<"simd", traits = [
let hasRegionVerifier = 1;
}
-
-def YieldOp : OpenMP_Op<"yield",
- [Pure, ReturnLike, Terminator,
- ParentOneOf<["AtomicUpdateOp", "DeclareReductionOp", "LoopNestOp",
- "PrivateClauseOp"]>]> {
+def YieldOp
+ : OpenMP_Op<"yield", [Pure, ReturnLike, Terminator,
+ ParentOneOf<["AtomicUpdateOp", "DeclareReductionOp",
+ "LoopNestOp", "PrivateClauseOp",
+ "IteratorOp"]>]> {
let summary = "loop yield and termination operation";
let description = [{
"omp.yield" yields SSA values from the OpenMP dialect op region and
@@ -2300,4 +2300,40 @@ def DeclareSimdOp
let hasVerifier = 1;
}
+//===----------------------------------------------------------------------===//
+// Iterator Op
+//===----------------------------------------------------------------------===//
+
+def IteratorOp : OpenMP_Op<"iterator",
+ [AttrSizedOperandSegments,
+ SingleBlockImplicitTerminator<"YieldOp">],
+ clauses = [OpenMP_LoopRelatedClause]> {
+ let summary = "OpenMP iterator modifier";
+ let description = [{
+ The result of `omp.iterator` is an abstract handle of type
+ `!omp.iterated<T>`, representing the list of yielded values. This handle
+ can be directly consumed by OpenMP clauses that accept iterator modifiers,
+ such as `affinity`, `map`, `to`, `from`, or `depend`.
+
+ Example:
+ %it = omp.iterator(%i, %j) =
+ (%lb_i to %ub_i step %st_i, %lb_j to %ub_j step %st_j) {
+ %addr = ...
+ omp.yield(%addr)
+ } -> !omp.iterated<ptr>
+ }];
+
+ let regions = (region SizedRegion<1>:$region);
+ let results = (outs OpenMP_IteratedType:$iterated);
+
+ let assemblyFormat = [{
+ `(` custom<IteratorHeader>($region,
+ $loop_lower_bounds, $loop_upper_bounds, $loop_steps,
+ type($loop_lower_bounds), type($loop_upper_bounds), type($loop_steps))
+ `->` qualified(type($iterated)) attr-dict
+ }];
+
+ let hasVerifier = 1;
+}
+
#endif // OPENMP_OPS
diff --git a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
index 8ffbe99756ff2..4b85d2bbf6db5 100644
--- a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
+++ b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
@@ -4627,6 +4627,126 @@ static void printAffinityClause(OpAsmPrinter &p, Operation *op,
}
}
+//===----------------------------------------------------------------------===//
+// Parser, printer, and verifier for Iterator modifier
+//===----------------------------------------------------------------------===//
+
+static ParseResult
+parseIteratorHeader(OpAsmParser &parser, Region ®ion,
+ SmallVectorImpl<OpAsmParser::UnresolvedOperand> &lbs,
+ SmallVectorImpl<OpAsmParser::UnresolvedOperand> &ubs,
+ SmallVectorImpl<OpAsmParser::UnresolvedOperand> &steps,
+ SmallVectorImpl<Type> &lbTypes,
+ SmallVectorImpl<Type> &ubTypes,
+ SmallVectorImpl<Type> &stepTypes) {
+
+ llvm::SMLoc ivLoc = parser.getCurrentLocation();
+ SmallVector<OpAsmParser::Argument> ivArgs;
+
+ // Parse induction variables: %i : i32, %j : i32
+ if (parser.parseCommaSeparatedList([&]() -> ParseResult {
+ OpAsmParser::Argument &arg = ivArgs.emplace_back();
+ if (parser.parseArgument(arg))
+ return failure();
+
+ // Optional type, default to Index if not provided
+ if (succeeded(parser.parseOptionalColon())) {
+ if (parser.parseType(arg.type))
+ return failure();
+ } else {
+ arg.type = parser.getBuilder().getIndexType();
+ }
+ return success();
+ }))
+ return failure();
+
+ // ) = (
+ if (parser.parseRParen() || parser.parseEqual() || parser.parseLParen())
+ return failure();
+
+ // Parse Ranges: (%lb to %ub step %st, ...)
+ if (parser.parseCommaSeparatedList([&]() -> ParseResult {
+ OpAsmParser::UnresolvedOperand lb, ub, st;
+ if (parser.parseOperand(lb) || parser.parseKeyword("to") ||
+ parser.parseOperand(ub) || parser.parseKeyword("step") ||
+ parser.parseOperand(st))
+ return failure();
+
+ lbs.push_back(lb);
+ ubs.push_back(ub);
+ steps.push_back(st);
+ return success();
+ }))
+ return failure();
+
+ if (parser.parseRParen())
+ return failure();
+
+ if (ivArgs.size() != lbs.size())
+ return parser.emitError(ivLoc)
+ << "mismatch: " << ivArgs.size() << " variables but " << lbs.size()
+ << " ranges";
+
+ for (auto &arg : ivArgs) {
+ lbTypes.push_back(arg.type);
+ ubTypes.push_back(arg.type);
+ stepTypes.push_back(arg.type);
+ }
+
+ return parser.parseRegion(region, ivArgs);
+}
+
+static void printIteratorHeader(OpAsmPrinter &p, Operation *op, Region ®ion,
+ ValueRange lbs, ValueRange ubs,
+ ValueRange steps, TypeRange, TypeRange,
+ TypeRange) {
+ Block &entry = region.front();
+
+ for (unsigned i = 0, e = entry.getNumArguments(); i < e; ++i) {
+ if (i != 0)
+ p << ", ";
+ p.printRegionArgument(entry.getArgument(i));
+ }
+ p << ") = (";
+
+ // (%lb0 to %ub0 step %step0, %lb1 to %ub1 step %step1, ...)
+ for (unsigned i = 0, e = lbs.size(); i < e; ++i) {
+ if (i)
+ p << ", ";
+ p << lbs[i] << " to " << ubs[i] << " step " << steps[i];
+ }
+ p << ") ";
+
+ p.printRegion(region, /*printEntryBlockArgs=*/false,
+ /*printBlockTerminators=*/true);
+}
+
+LogicalResult IteratorOp::verify() {
+ auto iteratedTy = llvm::dyn_cast<omp::IteratedType>(getIterated().getType());
+ if (!iteratedTy)
+ return emitOpError() << "result must be omp.iterated<entry_ty>";
+
+ Block &b = getRegion().front();
+ auto yield = llvm::dyn_cast<omp::YieldOp>(b.getTerminator());
+
+ if (!yield)
+ return emitOpError() << "region must be terminated by omp.yield";
+
+ if (yield.getNumOperands() != 1)
+ return emitOpError()
+ << "omp.yield in omp.iterator region must yield exactly one value";
+
+ mlir::Type yieldedTy = yield.getOperand(0).getType();
+ mlir::Type elemTy = iteratedTy.getElementType();
+
+ if (yieldedTy != elemTy)
+ return emitOpError() << "omp.iterated element type (" << elemTy
+ << ") does not match omp.yield operand type ("
+ << yieldedTy << ")";
+
+ return success();
+}
+
#define GET_ATTRDEF_CLASSES
#include "mlir/Dialect/OpenMP/OpenMPOpsAttributes.cpp.inc"
diff --git a/mlir/test/Dialect/OpenMP/invalid.mlir b/mlir/test/Dialect/OpenMP/invalid.mlir
index 4ee9b2c58a5ef..bc508d66fbd5f 100644
--- a/mlir/test/Dialect/OpenMP/invalid.mlir
+++ b/mlir/test/Dialect/OpenMP/invalid.mlir
@@ -3167,3 +3167,43 @@ func.func @omp_declare_simd_branch() -> () {
omp.declare_simd inbranch notinbranch
return
}
+
+// -----
+
+func.func @iterator_bad_result_type(%lb : index, %ub : index, %st : index) {
+ // expected-error at +1 {{result #0 must be OpenMP iterator-produced list handle, but got 'index'}}
+ %0 = omp.iterator(%i: index) = (%lb to %ub step %st) {
+ omp.yield(%i : index)
+ } -> index
+ return
+}
+
+// -----
+
+func.func @iterator_missing_yield(%lb : index, %ub : index, %st : index) {
+ // expected-error at +1 {{region must be terminated by omp.yield}}
+ %0 = omp.iterator(%i: index) = (%lb to %ub step %st) {
+ func.return
+ } -> !omp.iterated<index>
+ return
+}
+
+// -----
+
+func.func @iterator_yield_wrong_num_operands(%lb : index, %ub : index, %st : index) {
+ // expected-error at +1 {{omp.yield in omp.iterator region must yield exactly one value}}
+ %0 = omp.iterator(%i: index) = (%lb to %ub step %st) {
+ omp.yield(%i, %i : index, index)
+ } -> !omp.iterated<index>
+ return
+}
+
+// -----
+
+func.func @iterator_yield_type_mismatch(%lb : index, %ub : index, %st : index) {
+ // expected-error at +1 {{omp.iterated element type ('i64') does not match omp.yield operand type ('index')}}
+ %0 = omp.iterator(%i: index) = (%lb to %ub step %st) {
+ omp.yield(%i : index)
+ } -> !omp.iterated<i64>
+ return
+}
diff --git a/mlir/test/Dialect/OpenMP/ops.mlir b/mlir/test/Dialect/OpenMP/ops.mlir
index d73f09ae29cd0..c02c134b008a0 100644
--- a/mlir/test/Dialect/OpenMP/ops.mlir
+++ b/mlir/test/Dialect/OpenMP/ops.mlir
@@ -3550,3 +3550,41 @@ func.func @task_affinity_multi() {
}
return
}
+
+// CHECK-LABEL: func.func @omp_iterator
+func.func @omp_iterator(%s2 : !llvm.struct<(ptr, i64)>) -> () {
+ // CHECK: %[[IT:.*]] = omp.iterator(%[[IV:.*]]: index) = ({{.*}} to {{.*}} step {{.*}}) {
+ // CHECK: omp.yield(%{{.*}} : !llvm.struct<(ptr, i64)>)
+ // CHECK: } -> !omp.iterated<!llvm.struct<(ptr, i64)>>
+ %lb = arith.constant 1 : index
+ %ub = arith.constant 4 : index
+ %st = arith.constant 1 : index
+
+ %0 = omp.iterator(%iv: index) = (%lb to %ub step %st) {
+ omp.yield(%s2 : !llvm.struct<(ptr, i64)>)
+ } -> !omp.iterated<!llvm.struct<(ptr, i64)>>
+ return
+}
+
+// CHECK-LABEL: func.func @omp_iterator_2d
+func.func @omp_iterator_2d(%s2 : !llvm.struct<(ptr, i64)>) -> () {
+ // CHECK: %[[IT:.*]] = omp.iterator(%[[IV0:.*]]: index, %[[IV1:.*]]: index) =
+ // CHECK-SAME: ({{.*}} to {{.*}} step {{.*}}, {{.*}} to {{.*}} step {{.*}}) {
+ // CHECK: omp.yield(%{{.*}} : !llvm.struct<(ptr, i64)>)
+ // CHECK: } -> !omp.iterated<!llvm.struct<(ptr, i64)>>
+
+ %lb0 = arith.constant 1 : index
+ %ub0 = arith.constant 4 : index
+ %st0 = arith.constant 1 : index
+ %lb1 = arith.constant 2 : index
+ %ub1 = arith.constant 10 : index
+ %st1 = arith.constant 2 : index
+
+ %0 = omp.iterator(%iv0: index, %iv1: index) =
+ (%lb0 to %ub0 step %st0,
+ %lb1 to %ub1 step %st1) {
+ omp.yield(%s2 : !llvm.struct<(ptr, i64)>)
+ } -> !omp.iterated<!llvm.struct<(ptr, i64)>>
+
+ return
+}
More information about the Mlir-commits
mailing list