diff --git a/mlir/include/mlir/Dialect/SCF/IR/ValueBoundsOpInterfaceImpl.h b/mlir/include/mlir/Dialect/SCF/IR/ValueBoundsOpInterfaceImpl.h new file mode 100644 --- /dev/null +++ b/mlir/include/mlir/Dialect/SCF/IR/ValueBoundsOpInterfaceImpl.h @@ -0,0 +1,20 @@ +//===- ValueBoundsOpInterfaceImpl.h - Impl. of ValueBoundsOpInterface -----===// +// +// 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 +// +//===----------------------------------------------------------------------===// + +#ifndef MLIR_DIALECT_SCF_IR_VALUEBOUNDSOPINTERFACEIMPL_H +#define MLIR_DIALECT_SCF_IR_VALUEBOUNDSOPINTERFACEIMPL_H + +namespace mlir { +class DialectRegistry; + +namespace scf { +void registerValueBoundsOpInterfaceExternalModels(DialectRegistry ®istry); +} // namespace scf +} // namespace mlir + +#endif // MLIR_DIALECT_SCF_IR_VALUEBOUNDSOPINTERFACEIMPL_H diff --git a/mlir/include/mlir/InitAllDialects.h b/mlir/include/mlir/InitAllDialects.h --- a/mlir/include/mlir/InitAllDialects.h +++ b/mlir/include/mlir/InitAllDialects.h @@ -58,6 +58,7 @@ #include "mlir/Dialect/PDLInterp/IR/PDLInterp.h" #include "mlir/Dialect/Quant/QuantOps.h" #include "mlir/Dialect/SCF/IR/SCF.h" +#include "mlir/Dialect/SCF/IR/ValueBoundsOpInterfaceImpl.h" #include "mlir/Dialect/SCF/TransformOps/SCFTransformOps.h" #include "mlir/Dialect/SCF/Transforms/BufferizableOpInterfaceImpl.h" #include "mlir/Dialect/SPIRV/IR/SPIRVDialect.h" @@ -144,6 +145,7 @@ memref::registerRuntimeVerifiableOpInterfaceExternalModels(registry); memref::registerValueBoundsOpInterfaceExternalModels(registry); scf::registerBufferizableOpInterfaceExternalModels(registry); + scf::registerValueBoundsOpInterfaceExternalModels(registry); shape::registerBufferizableOpInterfaceExternalModels(registry); sparse_tensor::registerBufferizableOpInterfaceExternalModels(registry); tensor::registerBufferizableOpInterfaceExternalModels(registry); diff --git a/mlir/lib/Dialect/SCF/IR/CMakeLists.txt b/mlir/lib/Dialect/SCF/IR/CMakeLists.txt --- a/mlir/lib/Dialect/SCF/IR/CMakeLists.txt +++ b/mlir/lib/Dialect/SCF/IR/CMakeLists.txt @@ -1,6 +1,7 @@ add_mlir_dialect_library(MLIRSCFDialect SCF.cpp DeviceMappingInterface.cpp + ValueBoundsOpInterfaceImpl.cpp ADDITIONAL_HEADER_DIRS ${MLIR_MAIN_INCLUDE_DIR}/mlir/SCF @@ -15,5 +16,6 @@ MLIRIR MLIRLoopLikeInterface MLIRSideEffectInterfaces + MLIRValueBoundsOpInterface ) diff --git a/mlir/lib/Dialect/SCF/IR/ValueBoundsOpInterfaceImpl.cpp b/mlir/lib/Dialect/SCF/IR/ValueBoundsOpInterfaceImpl.cpp new file mode 100644 --- /dev/null +++ b/mlir/lib/Dialect/SCF/IR/ValueBoundsOpInterfaceImpl.cpp @@ -0,0 +1,50 @@ +//===- ValueBoundsOpInterfaceImpl.cpp - Impl. of ValueBoundsOpInterface ---===// +// +// 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 +// +//===----------------------------------------------------------------------===// + +#include "mlir/Dialect/SCF/IR/ValueBoundsOpInterfaceImpl.h" + +#include "mlir/Dialect/SCF/IR/SCF.h" +#include "mlir/Interfaces/ValueBoundsOpInterface.h" + +using namespace mlir; + +namespace mlir { +namespace scf { +namespace { + +struct ForOpInterface + : public ValueBoundsOpInterface::ExternalModel { + void populateBoundsForIndexValue(Operation *op, Value value, + ValueBoundsConstraintSet &cstr) const { + auto forOp = cast(op); + // Only IV is supported at the moment. + if (value != forOp.getInductionVar()) + return; + + // TODO: Take into account step size. + cstr.bound(value) >= forOp.getLowerBound(); + cstr.bound(value) < forOp.getUpperBound(); + } + + void populateBoundsForShapedValueDim(Operation *op, Value value, int64_t dim, + ValueBoundsConstraintSet &cstr) const { + // iter_arg / return value not supported. + return; + } +}; + +} // namespace +} // namespace scf +} // namespace mlir + +void mlir::scf::registerValueBoundsOpInterfaceExternalModels( + DialectRegistry ®istry) { + registry.addExtension(+[](MLIRContext *ctx, scf::SCFDialect *dialect) { + scf::ForOp::attachInterface(*ctx); + }); +} diff --git a/mlir/test/Dialect/SCF/value-bounds-op-interface-impl.mlir b/mlir/test/Dialect/SCF/value-bounds-op-interface-impl.mlir new file mode 100644 --- /dev/null +++ b/mlir/test/Dialect/SCF/value-bounds-op-interface-impl.mlir @@ -0,0 +1,14 @@ +// RUN: mlir-opt %s -test-affine-reify-value-bounds -verify-diagnostics \ +// RUN: -split-input-file | FileCheck %s + +// CHECK-LABEL: func @scf_for( +// CHECK-SAME: %[[a:.*]]: index, %[[b:.*]]: index, %[[c:.*]]: index +// CHECK: "test.some_use"(%[[a]], %[[b]]) +func.func @scf_for(%a: index, %b: index, %c: index) { + scf.for %iv = %a to %b step %c { + %0 = "test.reify_bound"(%iv) {type = "LB"} : (index) -> (index) + %1 = "test.reify_bound"(%iv) {type = "UB"} : (index) -> (index) + "test.some_use"(%0, %1) : (index, index) -> () + } + return +} diff --git a/utils/bazel/llvm-project-overlay/mlir/BUILD.bazel b/utils/bazel/llvm-project-overlay/mlir/BUILD.bazel --- a/utils/bazel/llvm-project-overlay/mlir/BUILD.bazel +++ b/utils/bazel/llvm-project-overlay/mlir/BUILD.bazel @@ -3038,6 +3038,7 @@ ":SCFPassIncGen", ":Support", ":TensorDialect", + ":ValueBoundsOpInterface", ":ViewLikeInterface", "//llvm:Support", ],