LinalgElementwiseFusion: fix linalg.index computation while fusing (#410)

 Fix num of loops calculation

Author: josel-amd

mlir/lib/Dialect/Linalg/Transforms/ElementwiseOpFusion.cpp

@@ -223,8 +223,7 @@ static void generateFusedElementwiseOpRegion(
   // `consumerToProducerLoopsMap` to map the producer indices.
   if (producer.hasIndexSemantics()) {
     // Add an index operation for every fused loop dimension.
-    unsigned numFusedOpLoops =
-        std::max(producer.getNumLoops(), consumer.getNumLoops());
+    unsigned numFusedOpLoops = consumerToProducerLoopsMap.getNumDims();
     SmallVector<Value> fusedIndices;
     fusedIndices.reserve(numFusedOpLoops);
     llvm::transform(llvm::seq<uint64_t>(0, numFusedOpLoops),

mlir/test/Dialect/Linalg/fusion-elementwise-ops.mlir

@@ -977,3 +977,32 @@ module {
 //   CHECK-DAG:     %[[T3:.+]] = arith.addf %[[T2]], %[[B1]]
 //       CHECK:     linalg.yield %[[T3]] : f32
 //       CHECK:   return %[[GENERIC]]
+
+// -----
+
+#map = affine_map<(d0, d1, d2) -> (d1, d2)>
+#map1 = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
+#map2 = affine_map<(d0, d1) -> (d0, d1 floordiv 4, d1 mod 4)>
+#map3 = affine_map<(d0, d1) -> (d0, d1)>
+
+// CHECK-DAG: [[$MAP0:#[a-zA-Z0-9_]*]] = affine_map<(d0, d1) -> (d0, d1)>
+// CHECK-DAG: [[$MAP1:#[a-zA-Z0-9_]*]] = affine_map<(d0) -> (d0 floordiv 4)>
+
+func.func @fuse_and_collapse(%arg0: tensor<3x4xindex>) -> tensor<2x12xindex> {
+  %1 = tensor.empty() : tensor<2x3x4xindex>
+  // CHECK: linalg.generic {
+  // CHECK: %[[INDEX1:[a-zA-Z0-9_]+]] = linalg.index 1 : index
+  // CHECK-NEXT: %[[MAP:[a-zA-Z0-9_]+]] = affine.apply #map1(%[[INDEX1]])
+  // CHECK-NEXT: linalg.yield %[[MAP]] : index
+  %2 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel"]} ins(%arg0: tensor<3x4xindex>) outs(%1 : tensor<2x3x4xindex>) {
+  ^bb0(%in: index, %out: index):
+    %3 = linalg.index 1 : index
+    linalg.yield %3: index
+  } -> tensor<2x3x4xindex>
+  %7 = tensor.empty() : tensor<2x12xindex>
+  %8 = linalg.generic {indexing_maps = [#map2, #map3], iterator_types = ["parallel", "parallel"]} ins(%2 : tensor<2x3x4xindex>) outs(%7 : tensor<2x12xindex>) {
+  ^bb0(%in: index, %out: index):
+    linalg.yield %in : index
+  } -> tensor<2x12xindex>
+  return %8 : tensor<2x12xindex>
+}