[mlir][LLVMIR] Fix fusion for rank-0 tensors

[asaadaldien]

Summary

This diff fixes fusion craching for ops with rank-0 tensors

[mravishankar]

mlir/lib/IR/AffineMap.cpp

260–261

There are going to be so many of such cases. The get method that gets context from the result expression should just be removed IMO

[asaadaldien]

mlir/lib/IR/AffineMap.cpp

260–261

IMO we should have a single ::get method that handle any results exprs including empty e.g

AffineMap AffineMap::get(unsigned dimCount, unsigned symbolCount,
                         ArrayRef<AffineExpr> results) {
  
  if (!results.empty()) {
  getImpl(dimCount, symbolCount, results, results[0].getContext());
 } else {
   assert(symbolCount == 0);
   return get(numResultDims, 0, this->getContext());
 }
}

Not sure if Nicolas prefer to have an explicit get method for rank-0 tensors for a specific reason.

[mravishankar]

mlir/lib/IR/AffineMap.cpp

260–261

There is no this->getContext() . This is a static class method. So this would not compile.

I think the reason for the separate method is because of the current assert(!result.empty()). Should just have a single method

AffineMap AffineMap::get(unsigned dimCount, unsigned symbolCount, ArrayRef<AffineExpr> results, MLIRContext *context) {
  return getImpl(dimCount, symbolCount, results, context);
}

[asaadaldien]

mlir/lib/IR/AffineMap.cpp

260–261

ah Didn't notice its static method. having both results & context and a single get SGTM. We can do it in a follow up diff it will touch many places

[rriddle]

mlir/lib/IR/AffineMap.cpp

260–261

nit: Drop the this->

[asaadaldien]

resolve nit

mlir/lib/IR/AffineMap.cpp

AffineMap AffineMap::replaceDimsAndSymbols(ArrayRef<AffineExpr> dimReplacements,
                                            unsigned numResultDims,
                                            unsigned numResultSyms) {
   SmallVector<AffineExpr, 8> results;
   results.reserve(getNumResults());
   for (auto expr : getResults())
     results.push_back(
         expr.replaceDimsAndSymbols(dimReplacements, symReplacements));
 
-  return get(numResultDims, numResultSyms, results);
+  return results.empty() ? get(numResultDims, 0, getContext())
+                         : get(numResultDims, numResultSyms, results);
 }
 
 AffineMap AffineMap::compose(AffineMap map) {
   assert(getNumDims() == map.getNumResults() && "Number of results mismatch");
   // Prepare `map` by concatenating the symbols and rewriting its exprs.
   unsigned numDims = map.getNumDims();
   unsigned numSymbolsThisMap = getNumSymbols();
   unsigned numSymbols = numSymbolsThisMap + map.getNumSymbols();

mlir/test/Dialect/Linalg/fusion-tensor.mlir

func @add_broadcast_mul_fusion(%arg0: tensor<?xf32>, %arg1 : tensor<?xf32>, %arg2 : tensor<?x?xf32>) -> tensor<?x?xf32>
   // CHECK-SAME: indexing_maps = {{\[}}[[MAP1]], [[MAP1]], [[MAP0]], [[MAP0]]
   %2 = linalg.generic {args_in = 2 : i64, args_out = 1 : i64, indexing_maps = [#map1, #map0, #map0], iterator_types = ["parallel", "parallel"]} %0, %arg2 {
     ^bb0(%arg5: f32, %arg6: f32):       // no predecessors
       %3 = mulf %arg5, %arg6 : f32
       linalg.yield %3 : f32
     }: tensor<?xf32>, tensor<?x?xf32> -> tensor<?x?xf32>
   return %2 : tensor<?x?xf32>
 }
+
+// -----
+
+// CHECK: #[[MAP0:.*]] = affine_map<() -> ()>
+#map0 = affine_map<() -> ()>
+
+// CHECK-LABEL: @add_mul_scalar_fusion
+func @add_mul_scalar_fusion(%arg0: tensor<f32>, %arg1: tensor<f32>, %arg2: tensor<f32>) -> tensor<f32>
+{
+  %0 = linalg.generic {args_in = 2 : i64, args_out = 1 : i64, indexing_maps = [#map0, #map0, #map0], iterator_types = []} %arg0, %arg1 {
+    ^bb0(%arg3: f32, %arg4: f32):       // no predecessors
+      %1 = addf %arg3, %arg4 : f32
+      linalg.yield %1 : f32
+  }: tensor<f32>, tensor<f32> -> tensor<f32>
+  // CHECK: linalg.generic {args_in = 2 : i64, args_out = 1 : i64
+  // CHECK: addf
+  // CHECK: mulf
+  %1 = linalg.generic {args_in = 2 : i64, args_out = 1 : i64, indexing_maps = [#map0, #map0, #map0], iterator_types = []} %0, %arg2 {
+    ^bb0(%arg3: f32, %arg4: f32):       // no predecessors
+      %1 = mulf %arg3, %arg4 : f32
+      linalg.yield %1 : f32
+  }: tensor<f32>, tensor<f32> -> tensor<f32>
+
+  return %1 : tensor<f32>
+}