@@ -1046,59 +1046,81 @@ simplifyMapWithOperands(AffineMap &map, ArrayRef<Value> operands) {
1046
1046
map.getContext ());
1047
1047
}
1048
1048
1049
- // / Assuming `dimOrSym` is a quantity in `map` that is defined by `minOp`.
1050
- // / Assuming that the quantity is of the form:
1051
- // / `affine_min(f(x, y), symbolic_cst)`.
1052
- // / This function checks that `0 < affine_min(f(x, y), symbolic_cst)` and
1053
- // / proceeds with replacing the patterns:
1049
+ // / Assuming `dimOrSym` is a quantity in the apply op map `map` and defined by
1050
+ // / `minOp = affine_min(x_1, ..., x_n)`. This function checks that:
1051
+ // / `0 < affine_min(x_1, ..., x_n)` and proceeds with replacing the patterns:
1054
1052
// / ```
1055
- // / dimOrSym.ceildiv(symbolic_cst )
1056
- // / (dimOrSym + symbolic_cst - 1).floordiv(symbolic_cst )
1053
+ // / dimOrSym.ceildiv(x_k )
1054
+ // / (dimOrSym + x_k - 1).floordiv(x_k )
1057
1055
// / ```
1058
- // / by `1`.
1056
+ // / by `1` for all `k` in `1, ..., n`. This is possible because `x / x_k <= 1` .
1059
1057
// /
1060
- // / Additionally, allows the caller to pass `affineMinKnownToBeNonNegative` to
1061
- // / inject static information that may not be statically discoverable.
1062
1058
// /
1063
1059
// / Warning: ValueBoundsConstraintSet::computeConstantBound is needed to check
1064
- // / for the nonnegative case, if `affineMinKnownToBeNonNegative` is false.
1065
- static LogicalResult replaceAffineMinBoundingBoxExpression (
1066
- AffineMinOp minOp, AffineExpr dimOrSym, AffineMap *map,
1067
- bool affineMinKnownToBeNonNegative = false ) {
1068
- auto affineMinMap = minOp.getAffineMap ();
1069
- if (!affineMinKnownToBeNonNegative) {
1070
- ValueRange values = minOp->getOperands ();
1071
- for (unsigned i = 0 , e = affineMinMap.getNumResults (); i < e; ++i) {
1072
- AffineMap row = affineMinMap.getSubMap (ArrayRef<unsigned >{i});
1073
- FailureOr<int64_t > lowerBound =
1074
- ValueBoundsConstraintSet::computeConstantBound (
1075
- presburger::BoundType::LB, {row, values},
1076
- /* stopCondition=*/ nullptr ,
1077
- /* closedUB=*/ true );
1078
- if (failed (lowerBound) || lowerBound.value () <= 0 )
1079
- return failure ();
1060
+ // / `minOp` is positive.
1061
+ static LogicalResult replaceAffineMinBoundingBoxExpression (AffineMinOp minOp,
1062
+ AffineExpr dimOrSym,
1063
+ AffineMap *map,
1064
+ ValueRange dims,
1065
+ ValueRange syms) {
1066
+ AffineMap affineMinMap = minOp.getAffineMap ();
1067
+
1068
+ // Check the value is positive.
1069
+ for (unsigned i = 0 , e = affineMinMap.getNumResults (); i < e; ++i) {
1070
+ // Compare each expression in the minimum against 0.
1071
+ if (!ValueBoundsConstraintSet::compare (
1072
+ getAsIndexOpFoldResult (minOp.getContext (), 0 ),
1073
+ ValueBoundsConstraintSet::ComparisonOperator::LT,
1074
+ ValueBoundsConstraintSet::Variable (affineMinMap.getSliceMap (i, 1 ),
1075
+ minOp.getOperands ())))
1076
+ return failure ();
1077
+ }
1078
+
1079
+ // / Convert affine symbols and dimensions in minOp to symbols or dimensions in
1080
+ // / the apply op affine map.
1081
+ DenseMap<AffineExpr, AffineExpr> dimSymConversionTable;
1082
+ SmallVector<unsigned > unmappedDims, unmappedSyms;
1083
+ for (auto [i, dim] : llvm::enumerate (minOp.getDimOperands ())) {
1084
+ auto it = llvm::find (dims, dim);
1085
+ if (it == dims.end ()) {
1086
+ unmappedDims.push_back (i);
1087
+ continue ;
1080
1088
}
1089
+ dimSymConversionTable[getAffineDimExpr (i, minOp.getContext ())] =
1090
+ getAffineDimExpr (it.getIndex (), minOp.getContext ());
1091
+ }
1092
+ for (auto [i, sym] : llvm::enumerate (minOp.getSymbolOperands ())) {
1093
+ auto it = llvm::find (syms, sym);
1094
+ if (it == syms.end ()) {
1095
+ unmappedSyms.push_back (i);
1096
+ continue ;
1097
+ }
1098
+ dimSymConversionTable[getAffineSymbolExpr (i, minOp.getContext ())] =
1099
+ getAffineSymbolExpr (it.getIndex (), minOp.getContext ());
1081
1100
}
1082
1101
1083
- AffineMap initialMap = *map;
1084
- for (unsigned i = 0 , e = affineMinMap.getNumResults (); i != e; ++i) {
1085
- auto m = affineMinMap.getSubMap (ArrayRef<unsigned >{i});
1086
- AffineExpr expr = m.getResult (0 );
1087
- if (!expr.isSymbolicOrConstant ())
1102
+ // Create the replacement map.
1103
+ DenseMap<AffineExpr, AffineExpr> repl;
1104
+ AffineExpr c1 = getAffineConstantExpr (1 , minOp.getContext ());
1105
+ for (AffineExpr expr : affineMinMap.getResults ()) {
1106
+ // If we cannot express the result in terms of the apply map symbols and
1107
+ // sims then continue.
1108
+ if (llvm::any_of (unmappedDims,
1109
+ [&](unsigned i) { return expr.isFunctionOfDim (i); }) ||
1110
+ llvm::any_of (unmappedSyms,
1111
+ [&](unsigned i) { return expr.isFunctionOfSymbol (i); }))
1088
1112
continue ;
1089
1113
1090
- DenseMap<AffineExpr, AffineExpr> repl;
1114
+ AffineExpr convertedExpr = expr.replace (dimSymConversionTable);
1115
+
1091
1116
// dimOrSym.ceilDiv(expr) -> 1
1092
- repl[dimOrSym.ceilDiv (expr )] = getAffineConstantExpr ( 1 , minOp. getContext ()) ;
1117
+ repl[dimOrSym.ceilDiv (convertedExpr )] = c1 ;
1093
1118
// (dimOrSym + expr - 1).floorDiv(expr) -> 1
1094
- repl[(dimOrSym + expr - 1 ).floorDiv (expr)] =
1095
- getAffineConstantExpr (1 , minOp.getContext ());
1096
- auto newMap = map->replace (repl);
1097
- if (newMap == *map)
1098
- continue ;
1099
- *map = newMap;
1119
+ repl[(dimOrSym + convertedExpr - 1 ).floorDiv (convertedExpr)] = c1;
1100
1120
}
1101
-
1121
+ AffineMap initialMap = *map;
1122
+ *map = initialMap.replace (repl, initialMap.getNumDims (),
1123
+ initialMap.getNumSymbols ());
1102
1124
return success (*map != initialMap);
1103
1125
}
1104
1126
@@ -1127,11 +1149,11 @@ static LogicalResult replaceDimOrSym(AffineMap *map,
1127
1149
if (!v)
1128
1150
return failure ();
1129
1151
1130
- auto minOp = v.getDefiningOp <AffineMinOp>();
1131
- if (minOp && replaceAffineMin) {
1152
+ if (auto minOp = v.getDefiningOp <AffineMinOp>(); minOp && replaceAffineMin) {
1132
1153
AffineExpr dimOrSym = isDimReplacement ? getAffineDimExpr (pos, ctx)
1133
1154
: getAffineSymbolExpr (pos, ctx);
1134
- return replaceAffineMinBoundingBoxExpression (minOp, dimOrSym, map);
1155
+ return replaceAffineMinBoundingBoxExpression (minOp, dimOrSym, map, dims,
1156
+ syms);
1135
1157
}
1136
1158
1137
1159
auto affineApply = v.getDefiningOp <AffineApplyOp>();
0 commit comments