xref: /llvm-project/mlir/lib/Dialect/Vector/IR/ValueBoundsOpInterfaceImpl.cpp (revision 2861856baf16e43a5e465e87022c6c2c2d238969)
1 //===- ValueBoundsOpInterfaceImpl.cpp - Impl. of ValueBoundsOpInterface ---===//
2 //
3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4 // See https://llvm.org/LICENSE.txt for license information.
5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6 //
7 //===----------------------------------------------------------------------===//
8 
9 #include "mlir/Dialect/Vector/IR/ValueBoundsOpInterfaceImpl.h"
10 
11 #include "mlir/Dialect/Vector/IR/ScalableValueBoundsConstraintSet.h"
12 #include "mlir/Dialect/Vector/IR/VectorOps.h"
13 #include "mlir/Interfaces/ValueBoundsOpInterface.h"
14 
15 using namespace mlir;
16 
17 namespace mlir::vector {
18 namespace {
19 
20 struct VectorScaleOpInterface
21     : public ValueBoundsOpInterface::ExternalModel<VectorScaleOpInterface,
22                                                    VectorScaleOp> {
populateBoundsForIndexValuemlir::vector::__anonb0ac0ff70111::VectorScaleOpInterface23   void populateBoundsForIndexValue(Operation *op, Value value,
24                                    ValueBoundsConstraintSet &cstr) const {
25     auto *scalableCstr = dyn_cast<ScalableValueBoundsConstraintSet>(&cstr);
26     if (!scalableCstr)
27       return;
28     auto vscaleOp = cast<VectorScaleOp>(op);
29     assert(value == vscaleOp.getResult() && "invalid value");
30     if (auto vscale = scalableCstr->getVscaleValue()) {
31       // All copies of vscale are equivalent.
32       scalableCstr->bound(value) == cstr.getExpr(vscale);
33     } else {
34       // We know vscale is confined to [vscaleMin, vscaleMax].
35       scalableCstr->bound(value) >= scalableCstr->getVscaleMin();
36       scalableCstr->bound(value) <= scalableCstr->getVscaleMax();
37       scalableCstr->setVscale(vscaleOp);
38     }
39   }
40 };
41 
42 } // namespace
43 } // namespace mlir::vector
44 
registerValueBoundsOpInterfaceExternalModels(DialectRegistry & registry)45 void mlir::vector::registerValueBoundsOpInterfaceExternalModels(
46     DialectRegistry &registry) {
47   registry.addExtension(+[](MLIRContext *ctx, vector::VectorDialect *dialect) {
48     vector::VectorScaleOp::attachInterface<vector::VectorScaleOpInterface>(
49         *ctx);
50   });
51 }
52