diff --git a/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td b/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td --- a/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td +++ b/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td @@ -337,13 +337,16 @@ /// list of AffineIf is not resizable. void setConditional(IntegerSet set, ValueRange operands); + /// Returns true if an else block exists. + bool hasElse() { return !elseRegion().empty(); } + Block *getThenBlock() { assert(!thenRegion().empty() && "Unexpected empty 'then' region."); return &thenRegion().front(); } Block *getElseBlock() { - assert(!elseRegion().empty() && "Empty 'else' region."); + assert(hasElse() && "Empty 'else' region."); return &elseRegion().front(); } @@ -353,7 +356,7 @@ return OpBuilder(&body, std::prev(body.end())); } OpBuilder getElseBodyBuilder() { - assert(!elseRegion().empty() && "Unexpected empty 'else' region."); + assert(hasElse() && "No 'else' block"); Block &body = elseRegion().front(); return OpBuilder(&body, std::prev(body.end())); } @@ -491,6 +494,9 @@ Block *getBody(); OpBuilder getBodyBuilder(); + MutableArrayRef getIVs() { + return getBody()->getArguments(); + } void setSteps(ArrayRef newSteps); static StringRef getLowerBoundsMapAttrName() { return "lowerBoundsMap"; } diff --git a/mlir/include/mlir/Dialect/Affine/Utils.h b/mlir/include/mlir/Dialect/Affine/Utils.h new file mode 100644 --- /dev/null +++ b/mlir/include/mlir/Dialect/Affine/Utils.h @@ -0,0 +1,28 @@ +//===- Utils.h - Affine dialect utilities -----------------------*- C++ -*-===// +// +// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// +//===----------------------------------------------------------------------===// +// +// This header file declares a set of utilities for the affine dialect ops. +// +//===----------------------------------------------------------------------===// + +#ifndef MLIR_DIALECT_AFFINE_UTILS_H +#define MLIR_DIALECT_AFFINE_UTILS_H + +namespace mlir { + +class AffineIfOp; + +/// Hoists out affine.if/else to as high as possible, i.e., past all invariant +/// affine.fors/parallel's. Returns true if any hoisting happened or the `ifOp` +/// was folded away; in the latter case `folded` is set to true. This +/// hoisting could lead to significant code expansion in some cases. +bool hoistAffineIfOp(AffineIfOp ifOp, bool *folded = nullptr); + +} // namespace mlir + +#endif // MLIR_DIALECT_AFFINE_UTILS_H diff --git a/mlir/include/mlir/IR/OpDefinition.h b/mlir/include/mlir/IR/OpDefinition.h --- a/mlir/include/mlir/IR/OpDefinition.h +++ b/mlir/include/mlir/IR/OpDefinition.h @@ -116,6 +116,12 @@ return getOperation()->getParentOfType(); } + /// Returns the closest surrounding parent operation with trait `Trait`. + template