[MLIR] Fix affine_map compose with multi-symbols

Fix bug: https://bugs.llvm.org/show_bug.cgi?id=46845

Differential Revision: https://reviews.llvm.org/D93831
This commit is contained in:
Chengji Yao 2021-01-02 06:55:40 +05:30 committed by Uday Bondhugula
parent e4337159e3
commit 3bcca6b12d
2 changed files with 13 additions and 1 deletions

View File

@ -319,7 +319,7 @@ AffineMap AffineMap::compose(AffineMap map) {
for (unsigned idx = 0; idx < numDims; ++idx) {
newDims[idx] = getAffineDimExpr(idx, getContext());
}
SmallVector<AffineExpr, 8> newSymbols(numSymbols);
SmallVector<AffineExpr, 8> newSymbols(numSymbols - numSymbolsThisMap);
for (unsigned idx = numSymbolsThisMap; idx < numSymbols; ++idx) {
newSymbols[idx - numSymbolsThisMap] =
getAffineSymbolExpr(idx, getContext());

View File

@ -24,6 +24,9 @@
// CHECK-DAG: [[$MAP13A:#map[0-9]+]] = affine_map<(d0) -> ((d0 + 6) ceildiv 8)>
// CHECK-DAG: [[$MAP13B:#map[0-9]+]] = affine_map<(d0) -> ((d0 * 4 - 4) floordiv 3)>
// Affine maps for test case: compose_affine_maps_multiple_symbols
// CHECK-DAG: [[$MAP14:#map[0-9]+]] = affine_map<()[s0, s1] -> (((s1 + s0) * 4) floordiv s0)>
// Affine maps for test case: partial_fold_map
// CHECK-DAG: [[$MAP15:#map[0-9]+]] = affine_map<()[s0] -> (s0 - 42)>
@ -218,6 +221,15 @@ func @compose_affine_maps_diamond_dependency(%arg0: f32, %arg1: memref<4x4xf32>)
return
}
// CHECK-LABEL: func @compose_affine_maps_multiple_symbols
func @compose_affine_maps_multiple_symbols(%arg0: index, %arg1: index) -> index {
%a = affine.apply affine_map<(d0)[s0] -> (s0 + d0)> (%arg0)[%arg1]
%c = affine.apply affine_map<(d0) -> (d0 * 4)> (%a)
%e = affine.apply affine_map<(d0)[s0] -> (d0 floordiv s0)> (%c)[%arg1]
// CHECK: [[I0:%[0-9]+]] = affine.apply [[$MAP14]]()[%{{.*}}, %{{.*}}]
return %e : index
}
// CHECK-LABEL: func @arg_used_as_dim_and_symbol
func @arg_used_as_dim_and_symbol(%arg0: memref<100x100xf32>, %arg1: index, %arg2: f32) {
%c9 = constant 9 : index