Skip to content

Commit 8aba696

Browse files
MacDueMacDue
committed
Fixups
1 parent b685de4 commit 8aba696

File tree

4 files changed

+50
-27
lines changed

4 files changed

+50
-27
lines changed

‎mlir/include/mlir/Dialect/Vector/IR/ScalableValueBoundsConstraintSet.h

Lines changed: 7 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -17,9 +17,9 @@ namespace mlir::vector {
1717

1818
namespace detail {
1919

20-
// Parent class for the value bounds RTTIExtends. Uses protected inheritance to
21-
// hide all ValueBoundsConstraintSet methods by default (as some do not use the
22-
// ScalableValueBoundsConstraintSet, so may produce unexpected results).
20+
/// Parent class for the value bounds RTTIExtends. Uses protected inheritance to
21+
/// hide all ValueBoundsConstraintSet methods by default (as some do not use the
22+
/// ScalableValueBoundsConstraintSet, so may produce unexpected results).
2323
struct ValueBoundsConstraintSet : protected ::mlir::ValueBoundsConstraintSet {
2424
using ::mlir::ValueBoundsConstraintSet::ValueBoundsConstraintSet;
2525
};
@@ -67,7 +67,7 @@ struct ScalableValueBoundsConstraintSet
6767
static FailureOr<ConstantOrScalableBound>
6868
computeScalableBound(Value value, std::optional<int64_t> dim,
6969
unsigned vscaleMin, unsigned vscaleMax,
70-
presburger::BoundType boundType,
70+
presburger::BoundType boundType, bool closedUB = true,
7171
StopConditionFn stopCondition = nullptr);
7272

7373
/// Get the value of vscale. Returns `nullptr` vscale as not been encountered.
@@ -90,6 +90,9 @@ struct ScalableValueBoundsConstraintSet
9090
private:
9191
const unsigned vscaleMin;
9292
const unsigned vscaleMax;
93+
94+
// This will be set when the first `vector.vscale` operation is found within
95+
// the `ValueBoundsOpInterface` implementation then reused from there on.
9396
Value vscale = nullptr;
9497
};
9598

‎mlir/lib/Dialect/Vector/IR/ScalableValueBoundsConstraintSet.cpp

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -44,7 +44,7 @@ char ScalableValueBoundsConstraintSet::ID = 0;
4444
FailureOr<ConstantOrScalableBound>
4545
ScalableValueBoundsConstraintSet::computeScalableBound(
4646
Value value, std::optional<int64_t> dim, unsigned vscaleMin,
47-
unsigned vscaleMax, presburger::BoundType boundType,
47+
unsigned vscaleMax, presburger::BoundType boundType, bool closedUB,
4848
StopConditionFn stopCondition) {
4949
using namespace presburger;
5050

@@ -76,8 +76,7 @@ ScalableValueBoundsConstraintSet::computeScalableBound(
7676

7777
SmallVector<AffineMap, 1> lowerBound(1), upperBound(1);
7878
scalableCstr.cstr.getSliceBounds(pos, 1, value.getContext(), &lowerBound,
79-
&upperBound,
80-
/*closedUB=*/true);
79+
&upperBound, closedUB);
8180

8281
auto invalidBound = [](auto &bound) {
8382
return !bound[0] || bound[0].getNumResults() != 1;

‎mlir/lib/Dialect/Vector/IR/ValueBoundsOpInterfaceImpl.cpp

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -32,8 +32,8 @@ struct VectorScaleOpInterface
3232
scalableCstr->bound(value) == cstr.getExpr(vscale);
3333
} else {
3434
// We know vscale is confined to [vscaleMin, vscaleMax].
35-
cstr.bound(value) >= scalableCstr->getVscaleMin();
36-
cstr.bound(value) <= scalableCstr->getVscaleMax();
35+
scalableCstr->bound(value) >= scalableCstr->getVscaleMin();
36+
scalableCstr->bound(value) <= scalableCstr->getVscaleMax();
3737
scalableCstr->setVscale(vscaleOp);
3838
}
3939
}

‎mlir/test/Dialect/Vector/test-scalable-bounds.mlir

Lines changed: 39 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -1,20 +1,20 @@
11
// RUN: mlir-opt %s -test-affine-reify-value-bounds -cse -verify-diagnostics \
22
// RUN: -verify-diagnostics -split-input-file | FileCheck %s
33

4-
#fixedDim0Map = affine_map<(d0)[s0] -> (-d0 + 32400, s0)>
5-
#fixedDim1Map = affine_map<(d0)[s0] -> (-d0 + 16, s0)>
4+
#map_dim_i = affine_map<(d0)[s0] -> (-d0 + 32400, s0)>
5+
#map_dim_j = affine_map<(d0)[s0] -> (-d0 + 16, s0)>
66

77
// Here the upper bound for min_i is 4 x vscale, as we know 4 x vscale is
8-
// always less than 32400. The bound for min_j is 16 as at vscale > 4,
9-
// 4 x vscale will be > 16, so the value will be clamped at 16.
8+
// always less than 32400. The bound for min_j is 16, as 16 is always less
9+
// 4 x vscale_max (vscale_max is the UB for vscale).
1010

1111
// CHECK: #[[$SCALABLE_BOUND_MAP_0:.*]] = affine_map<()[s0] -> (s0 * 4)>
1212

1313
// CHECK-LABEL: @fixed_size_loop_nest
1414
// CHECK-DAG: %[[VSCALE:.*]] = vector.vscale
15-
// CHECK-DAG: %[[SCALABLE_BOUND:.*]] = affine.apply #[[$SCALABLE_BOUND_MAP_0]]()[%[[VSCALE]]]
16-
// CHECK-DAG: %[[C16:.*]] = arith.constant 16 : index
17-
// CHECK: "test.some_use"(%[[SCALABLE_BOUND]], %[[C16]]) : (index, index) -> ()
15+
// CHECK-DAG: %[[UB_i:.*]] = affine.apply #[[$SCALABLE_BOUND_MAP_0]]()[%[[VSCALE]]]
16+
// CHECK-DAG: %[[UB_j:.*]] = arith.constant 16 : index
17+
// CHECK: "test.some_use"(%[[UB_i]], %[[UB_j]]) : (index, index) -> ()
1818
func.func @fixed_size_loop_nest() {
1919
%c16 = arith.constant 16 : index
2020
%c32400 = arith.constant 32400 : index
@@ -23,9 +23,9 @@ func.func @fixed_size_loop_nest() {
2323
%vscale = vector.vscale
2424
%c4_vscale = arith.muli %vscale, %c4 : index
2525
scf.for %i = %c0 to %c32400 step %c4_vscale {
26-
%min_i = affine.min #fixedDim0Map(%i)[%c4_vscale]
26+
%min_i = affine.min #map_dim_i(%i)[%c4_vscale]
2727
scf.for %j = %c0 to %c16 step %c4_vscale {
28-
%min_j = affine.min #fixedDim1Map(%j)[%c4_vscale]
28+
%min_j = affine.min #map_dim_j(%j)[%c4_vscale]
2929
%bound_i = "test.reify_scalable_bound"(%min_i) {type = "UB", vscale_min = 1, vscale_max = 16} : (index) -> index
3030
%bound_j = "test.reify_scalable_bound"(%min_j) {type = "UB", vscale_min = 1, vscale_max = 16} : (index) -> index
3131
"test.some_use"(%bound_i, %bound_j) : (index, index) -> ()
@@ -36,28 +36,28 @@ func.func @fixed_size_loop_nest() {
3636

3737
// -----
3838

39-
#dynamicDimMap = affine_map<(d0)[s0, s1] -> (-d0 + s1, s0)>
39+
#map_dynamic_dim = affine_map<(d0)[s0, s1] -> (-d0 + s1, s0)>
4040

41-
// Here upper bounds for both min_i and min_j are both 4 x vscale, as we know
42-
// that is always the largest value they could take. As if `dim < 4 x vscale`
43-
// then 4 x vscale is an overestimate, and if `dim > 4 x vscale` then the min
44-
// will be clamped to 4 x vscale.
41+
// Here upper bounds for both min_i and min_j are both (conservatively)
42+
// 4 x vscale, as we know that is always the largest value they could take. As
43+
// if `dim < 4 x vscale` then 4 x vscale is an overestimate, and if
44+
// `dim > 4 x vscale` then the min will be clamped to 4 x vscale.
4545

4646
// CHECK: #[[$SCALABLE_BOUND_MAP_1:.*]] = affine_map<()[s0] -> (s0 * 4)>
4747

4848
// CHECK-LABEL: @dynamic_size_loop_nest
4949
// CHECK: %[[VSCALE:.*]] = vector.vscale
50-
// CHECK: %[[SCALABLE_BOUND:.*]] = affine.apply #[[$SCALABLE_BOUND_MAP_1]]()[%[[VSCALE]]]
51-
// CHECK: "test.some_use"(%[[SCALABLE_BOUND]], %[[SCALABLE_BOUND]]) : (index, index) -> ()
50+
// CHECK: %[[UB_ij:.*]] = affine.apply #[[$SCALABLE_BOUND_MAP_1]]()[%[[VSCALE]]]
51+
// CHECK: "test.some_use"(%[[UB_ij]], %[[UB_ij]]) : (index, index) -> ()
5252
func.func @dynamic_size_loop_nest(%dim0: index, %dim1: index) {
5353
%c4 = arith.constant 4 : index
5454
%c0 = arith.constant 0 : index
5555
%vscale = vector.vscale
5656
%c4_vscale = arith.muli %vscale, %c4 : index
5757
scf.for %i = %c0 to %dim0 step %c4_vscale {
58-
%min_i = affine.min #dynamicDimMap(%i)[%c4_vscale, %dim0]
58+
%min_i = affine.min #map_dynamic_dim(%i)[%c4_vscale, %dim0]
5959
scf.for %j = %c0 to %dim1 step %c4_vscale {
60-
%min_j = affine.min #dynamicDimMap(%j)[%c4_vscale, %dim1]
60+
%min_j = affine.min #map_dynamic_dim(%j)[%c4_vscale, %dim1]
6161
%bound_i = "test.reify_scalable_bound"(%min_i) {type = "UB", vscale_min = 1, vscale_max = 16} : (index) -> index
6262
%bound_j = "test.reify_scalable_bound"(%min_j) {type = "UB", vscale_min = 1, vscale_max = 16} : (index) -> index
6363
"test.some_use"(%bound_i, %bound_j) : (index, index) -> ()
@@ -138,3 +138,24 @@ func.func @duplicate_vscale_values() {
138138
"test.some_use"(%bound) : (index) -> ()
139139
return
140140
}
141+
142+
// -----
143+
144+
// Test some non-scalable code to ensure that works too:
145+
146+
#map_dim_i = affine_map<(d0)[s0] -> (-d0 + 1024, s0)>
147+
148+
// CHECK-LABEL: @non_scalable_code
149+
// CHECK: %[[C4:.*]] = arith.constant 4 : index
150+
// CHECK: "test.some_use"(%[[C4]]) : (index) -> ()
151+
func.func @non_scalable_code() {
152+
%c1024 = arith.constant 1024 : index
153+
%c4 = arith.constant 4 : index
154+
%c0 = arith.constant 0 : index
155+
scf.for %i = %c0 to %c1024 step %c4 {
156+
%min_i = affine.min #map_dim_i(%i)[%c4]
157+
%bound_i = "test.reify_scalable_bound"(%min_i) {type = "UB", vscale_min = 1, vscale_max = 16} : (index) -> index
158+
"test.some_use"(%bound_i) : (index) -> ()
159+
}
160+
return
161+
}

0 commit comments

Comments
 (0)