diff --git a/mlir/include/mlir/Dialect/AffineOps/AffineOps.h b/mlir/include/mlir/Dialect/AffineOps/AffineOps.h --- a/mlir/include/mlir/Dialect/AffineOps/AffineOps.h +++ b/mlir/include/mlir/Dialect/AffineOps/AffineOps.h @@ -654,6 +654,10 @@ SmallVector reorderedDims; SmallVector concatenatedSymbols; + /// The number of symbols in concatenated symbols that belong to the original + /// map as opposed to those concatendated during map composition. + unsigned numProperSymbols; + AffineMap affineMap; /// Used with RAII to control the depth at which AffineApply are composed @@ -667,7 +671,7 @@ } static constexpr unsigned kMaxAffineApplyDepth = 1; - AffineApplyNormalizer() { affineApplyDepth()++; } + AffineApplyNormalizer() : numProperSymbols(0) { affineApplyDepth()++; } public: ~AffineApplyNormalizer() { affineApplyDepth()--; } diff --git a/mlir/lib/Dialect/AffineOps/AffineOps.cpp b/mlir/lib/Dialect/AffineOps/AffineOps.cpp --- a/mlir/lib/Dialect/AffineOps/AffineOps.cpp +++ b/mlir/lib/Dialect/AffineOps/AffineOps.cpp @@ -536,8 +536,12 @@ auxiliaryExprs.push_back(renumberOneDim(t)); } else { // c. The mathematical composition of AffineMap concatenates symbols. - // We do the same for symbol operands. - concatenatedSymbols.push_back(t); + // Note that the map composition will put symbols already present + // in the map before any symbols coming from the auxiliary map, so + // we insert them before any symbols that are due to renumbering, + // and after the proper symbols we have seen already. + concatenatedSymbols.insert( + std::next(concatenatedSymbols.begin(), numProperSymbols++), t); } } } diff --git a/mlir/test/Dialect/AffineOps/canonicalize.mlir b/mlir/test/Dialect/AffineOps/canonicalize.mlir --- a/mlir/test/Dialect/AffineOps/canonicalize.mlir +++ b/mlir/test/Dialect/AffineOps/canonicalize.mlir @@ -17,7 +17,7 @@ // Affine maps for test case: compose_affine_maps_dependent_loads // CHECK-DAG: [[MAP9:#map[0-9]+]] = affine_map<(d0) -> (d0 + 3)> // CHECK-DAG: [[MAP10:#map[0-9]+]] = affine_map<(d0) -> (d0 * 3)> -// CHECK-DAG: [[MAP11:#map[0-9]+]] = affine_map<(d0) -> ((d0 + 7) ceildiv 3)> +// CHECK-DAG: [[MAP11:#map[0-9]+]] = affine_map<(d0) -> ((d0 + 3) ceildiv 3)> // CHECK-DAG: [[MAP12:#map[0-9]+]] = affine_map<(d0) -> (d0 * 7 - 49)> // Affine maps for test case: compose_affine_maps_diamond_dependency @@ -187,9 +187,9 @@ // Swizzle %x00, %x01 and %c3, %c7 %x10 = affine.apply affine_map<(d0, d1)[s0, s1] -> (d0 * s1)> - (%x01, %x00)[%c7, %c3] + (%x01, %x00)[%c3, %c7] %x11 = affine.apply affine_map<(d0, d1)[s0, s1] -> (d1 ceildiv s0)> - (%x01, %x00)[%c7, %c3] + (%x01, %x00)[%c3, %c7] // CHECK-NEXT: [[I2A:%[0-9]+]] = affine.apply [[MAP12]](%{{.*}}) // CHECK-NEXT: [[I2B:%[0-9]+]] = affine.apply [[MAP11]](%{{.*}}) @@ -568,3 +568,29 @@ } return } + +// ----- + +// Reproducer for PR45031. This used to fold into an incorrect map because +// symbols were concatenated in the wrong order during map folding. Map +// composition places the symbols of the original map before those of the map +// it is composed with, e.g. A.compose(B) will first have all symbols of A, +// then all symbols of B. + +#map1 = affine_map<(d0)[s0, s1] -> (d0 * s0 + s1)> +#map2 = affine_map<(d0)[s0] -> (1024, -d0 + s0)> + +// CHECK: #[[MAP:.*]] = affine_map<()[s0, s1] -> (1024, s1 * -1024 + s0)> + +// CHECK: func @rep(%[[ARG0:.*]]: index, %[[ARG1:.*]]: index) +func @rep(%arg0 : index, %arg1 : index) -> index { + // CHECK-NOT: constant + %c0 = constant 0 : index + %c1024 = constant 1024 : index + // CHECK-NOT: affine.apply + %0 = affine.apply #map1(%arg0)[%c1024, %c0] + + // CHECK: affine.min #[[MAP]]()[%[[ARG1]], %[[ARG0]]] + %1 = affine.min #map2(%0)[%arg1] + return %1 : index +}