From a5bfd32c07b9a90e8fac18a3ebb722b12038222f Mon Sep 17 00:00:00 2001
From: Lei Zhang <antiagainst@google.com>
Date: Mon, 27 Apr 2020 11:47:39 -0400
Subject: [PATCH] [mlir][linalg] Fix crash in
 tileAndFuseLinalgOpToParallelLoopsAndSetMarker

Instead of using llvm_unreachable to guard against fusing linalg.conv,
reject fusing linalg.conv in isFusableInto.

tileLinalgOpImpl is a templated function now and it can operate on
loop.parellel. So we should avoid calling into getForInductionVarOwner
which always assumes loop.for.

Differential Revision: https://reviews.llvm.org/D78936
---
 mlir/lib/Dialect/Linalg/Transforms/Fusion.cpp | 21 +++++++++----------
 mlir/lib/Dialect/Linalg/Transforms/Tiling.cpp |  6 ++++--
 2 files changed, 14 insertions(+), 13 deletions(-)

diff --git a/mlir/lib/Dialect/Linalg/Transforms/Fusion.cpp b/mlir/lib/Dialect/Linalg/Transforms/Fusion.cpp
index cd6301ae249c..cea91f6dc532 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/Fusion.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/Fusion.cpp
@@ -161,17 +161,6 @@ static LinalgOp fuse(Value producedView, LinalgOp producer, LinalgOp consumer,
   assert(consumer.hasBufferSemantics() &&
          "expected linalg op with buffer semantics");
 
-  if (auto convOp = dyn_cast<linalg::ConvOp>(producer.getOperation())) {
-    // TODO(ntv): add a level of indirection to linalg.generic.
-    if (convOp.padding())
-      llvm_unreachable("Unexpected conv with padding");
-  }
-  if (auto convOp = dyn_cast<linalg::ConvOp>(consumer.getOperation())) {
-    // TODO(ntv): add a level of indirection to linalg.generic.
-    if (convOp.padding())
-      llvm_unreachable("Unexpected conv with padding");
-  }
-
   auto subView = dyn_cast_or_null<SubViewOp>(
       consumer.getBuffer(consumerIdx).getDefiningOp());
   auto slice = dyn_cast_or_null<SliceOp>(
@@ -287,6 +276,16 @@ bool mlir::linalg::isFusableInto(const LinalgDependenceGraph &graph,
                       << *producer.getOperation());
     return false;
   }
+  if (auto convOp = dyn_cast<linalg::ConvOp>(producer.getOperation())) {
+    // TODO(ntv): add a level of indirection to linalg.generic.
+    if (convOp.padding())
+      return false;
+  }
+  if (auto convOp = dyn_cast<linalg::ConvOp>(consumer.getOperation())) {
+    // TODO(ntv): add a level of indirection to linalg.generic.
+    if (convOp.padding())
+      return false;
+  }
   return true;
 }
 
diff --git a/mlir/lib/Dialect/Linalg/Transforms/Tiling.cpp b/mlir/lib/Dialect/Linalg/Transforms/Tiling.cpp
index 4dd8c5ad0e50..dd872a69d28d 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/Tiling.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/Tiling.cpp
@@ -409,8 +409,10 @@ Optional<TiledLinalgOp> static tileLinalgOpImpl(OpBuilder &b, LinalgOp op,
   // 5. Gather the newly created loops and return them with the new op.
   SmallVector<Operation *, 8> loops;
   loops.reserve(ivs.size());
-  for (auto iv : ivs)
-    loops.push_back(loop::getForInductionVarOwner(iv));
+  for (auto iv : ivs) {
+    loops.push_back(iv.cast<BlockArgument>().getOwner()->getParentOp());
+    assert(loops.back() && "no owner found for induction variable!");
+  }
 
   return TiledLinalgOp{res, loops};
 }