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 @@ -109,7 +109,8 @@ def AffineForOp : Affine_Op<"for", [AutomaticAllocationScope, ImplicitAffineTerminator, RecursiveSideEffects, DeclareOpInterfaceMethods]> { + ["getSingleInductionVar", "getSingleLowerBound", "getSingleStep", + "getSingleUpperBound"]>]> { let summary = "for operation"; let description = [{ Syntax: diff --git a/mlir/include/mlir/Dialect/SCF/SCFOps.td b/mlir/include/mlir/Dialect/SCF/SCFOps.td --- a/mlir/include/mlir/Dialect/SCF/SCFOps.td +++ b/mlir/include/mlir/Dialect/SCF/SCFOps.td @@ -111,7 +111,8 @@ def ForOp : SCF_Op<"for", [AutomaticAllocationScope, DeclareOpInterfaceMethods, + ["getSingleInductionVar", "getSingleLowerBound", "getSingleStep", + "getSingleUpperBound"]>, DeclareOpInterfaceMethods, SingleBlockImplicitTerminator<"scf::YieldOp">, RecursiveSideEffects]> { diff --git a/mlir/include/mlir/Interfaces/LoopLikeInterface.td b/mlir/include/mlir/Interfaces/LoopLikeInterface.td --- a/mlir/include/mlir/Interfaces/LoopLikeInterface.td +++ b/mlir/include/mlir/Interfaces/LoopLikeInterface.td @@ -61,7 +61,7 @@ }] >, InterfaceMethod<[{ - Return the single lower bound value or attribute if it exist, otherwise + Return the single lower bound value or attribute if it exists, otherwise return llvm::None. }], /*retTy=*/"::mlir::Optional<::mlir::OpFoldResult>", @@ -73,7 +73,7 @@ }] >, InterfaceMethod<[{ - Return the single step value or attribute if it exist, otherwise + Return the single step value or attribute if it exists, otherwise return llvm::None. }], /*retTy=*/"::mlir::Optional<::mlir::OpFoldResult>", @@ -84,6 +84,18 @@ return llvm::None; }] >, + InterfaceMethod<[{ + Return the single upper bound value or attribute if it exists, otherwise + return llvm::None. + }], + /*retTy=*/"::mlir::Optional<::mlir::OpFoldResult>", + /*methodName=*/"getSingleUpperBound", + /*args=*/(ins), + /*methodBody=*/"", + /*defaultImplementation=*/[{ + return llvm::None; + }] + >, ]; } diff --git a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp --- a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp +++ b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp @@ -1873,6 +1873,13 @@ return OpFoldResult(b.getI64IntegerAttr(getStep())); } +Optional AffineForOp::getSingleUpperBound() { + if (!hasConstantUpperBound()) + return llvm::None; + OpBuilder b(getContext()); + return OpFoldResult(b.getI64IntegerAttr(getConstantUpperBound())); +} + /// Returns true if the provided value is the induction variable of a /// AffineForOp. bool mlir::isForInductionVar(Value val) { diff --git a/mlir/lib/Dialect/SCF/SCF.cpp b/mlir/lib/Dialect/SCF/SCF.cpp --- a/mlir/lib/Dialect/SCF/SCF.cpp +++ b/mlir/lib/Dialect/SCF/SCF.cpp @@ -356,6 +356,10 @@ return OpFoldResult(getStep()); } +Optional ForOp::getSingleUpperBound() { + return OpFoldResult(getUpperBound()); +} + /// Prints the initialization list in the form of /// (%inner = %outer, %inner2 = %outer2, <...>) /// where 'inner' values are assumed to be region arguments and 'outer' values