[PATCH] D76479: [mlir][LLVMIR] Fix fusion for rank-0 tensors

[Ahmed S. Taei via Phabricator via llvm-commits]

asaadaldien created this revision.
asaadaldien added reviewers: mravishankar, nicolasvasilache.
Herald added subscribers: llvm-commits, Joonsoo, liufengdb, lucyrfox, mgester, arpith-jacob, antiagainst, shauheen, burmako, jpienaar, rriddle, mehdi_amini.
Herald added a reviewer: rriddle.
Herald added a project: LLVM.

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


Repository:
  rG LLVM Github Monorepo

https://reviews.llvm.org/D76479

Files:
  mlir/lib/IR/AffineMap.cpp
  mlir/test/Dialect/Linalg/fusion-tensor.mlir


Index: mlir/test/Dialect/Linalg/fusion-tensor.mlir
===================================================================
--- mlir/test/Dialect/Linalg/fusion-tensor.mlir
+++ mlir/test/Dialect/Linalg/fusion-tensor.mlir
@@ -105,3 +105,28 @@
     }: 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>
+}
Index: mlir/lib/IR/AffineMap.cpp
===================================================================
--- mlir/lib/IR/AffineMap.cpp
+++ mlir/lib/IR/AffineMap.cpp
@@ -257,7 +257,8 @@
     results.push_back(
         expr.replaceDimsAndSymbols(dimReplacements, symReplacements));
 
-  return get(numResultDims, numResultSyms, results);
+  return results.empty() ? get(numResultDims, 0, this->getContext())
+                         : get(numResultDims, numResultSyms, results);
 }
 
 AffineMap AffineMap::compose(AffineMap map) {


-------------- next part --------------
A non-text attachment was scrubbed...
Name: D76479.251558.patch
Type: text/x-patch
Size: 1845 bytes
Desc: not available
URL: <http://lists.llvm.org/pipermail/llvm-commits/attachments/20200320/2b93024e/attachment.bin>