fix affine_map compose with multi-symbols

[yaochengji]

Summary

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

[bondhugula]

Thanks very much for fixing this!

mlir/lib/IR/AffineMap.cpp

AffineMap AffineMap::compose(AffineMap map) {
   // Prepare `map` by concatenating the symbols and rewriting its exprs.
   unsigned numDims = map.getNumDims();
   unsigned numSymbolsThisMap = getNumSymbols();
   unsigned numSymbols = numSymbolsThisMap + map.getNumSymbols();
   SmallVector<AffineExpr, 8> newDims(numDims);
   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());
   }
   auto newMap =
       map.replaceDimsAndSymbols(newDims, newSymbols, numDims, numSymbols);
   SmallVector<AffineExpr, 8> exprs;
   exprs.reserve(getResults().size());

mlir/test/Dialect/Affine/canonicalize.mlir

 // CHECK-DAG: [[$MAP10:#map[0-9]+]] = affine_map<(d0) -> (d0 * 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
 // 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)>
 
 // Affine maps for test cases: symbolic_composition_*
 // CHECK-DAG: [[$MAP_symbolic_composition_a:#map[0-9]+]] = affine_map<()[s0] -> (s0 * 512)>
 // CHECK-DAG: [[$MAP_symbolic_composition_b:#map[0-9]+]] = affine_map<()[s0] -> (s0 * 4)>
 // CHECK-DAG: [[$MAP_symbolic_composition_c:#map[0-9]+]] = affine_map<()[s0, s1] -> (s0 * 3 + s1)>
 // CHECK-DAG: [[$MAP_symbolic_composition_d:#map[0-9]+]] = affine_map<()[s0, s1] -> (s1 * 3 + s0)>
affine.for %i0 = 0 to 15 {
     // CHECK: [[I1:%[0-9]+]] = affine.apply [[$MAP13B]](%{{.*}})
     // CHECK-NEXT: store %arg0, %arg1{{\[}}[[I0]], [[I1]]{{\]}}
     store %arg0, %arg1[%d0, %d1] : 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
   %1 = alloc() : memref<100x100xf32, 1>
   %2 = alloc() : memref<1xi32>
   affine.for %i0 = 0 to 100 {
     affine.for %i1 = 0 to 100 {
       %3 = affine.apply affine_map<(d0, d1)[s0, s1] -> (d1 + s0 + s1)>