From f984a805f3f92ea3066ea13ba2684a7947d95225 Mon Sep 17 00:00:00 2001
From: KareemErgawy-TomTom <kareem.ergawy@gmail.com>
Date: Tue, 3 Aug 2021 18:08:00 +0200
Subject: [PATCH] [MLIR][Linalg] Extend detensoring control flow model.

This patch extends the PureControlFlowDetectionModel to consider
detensoring br and cond_br operands.

See: https://github.com/google/iree/issues/1159#issuecomment-884322687,
for a disccusion on the need for such extension.

Reviewed By: silvas

Differential Revision: https://reviews.llvm.org/D107358
---
 .../Dialect/Linalg/Transforms/Detensorize.cpp | 13 ++++-
 .../Linalg/detensorize_br_operands.mlir       | 49 +++++++++++++++++++
 2 files changed, 60 insertions(+), 2 deletions(-)
 create mode 100644 mlir/test/Dialect/Linalg/detensorize_br_operands.mlir

diff --git a/mlir/lib/Dialect/Linalg/Transforms/Detensorize.cpp b/mlir/lib/Dialect/Linalg/Transforms/Detensorize.cpp
index 920e34899755..8fe6ac5be980 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/Detensorize.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/Detensorize.cpp
@@ -297,8 +297,17 @@ struct LinalgDetensorize : public LinalgDetensorizeBase<LinalgDetensorize> {
                  DenseSet<BlockArgument> &blockArgsToDetensor) override {
       SmallVector<Value> workList;
 
-      func.walk(
-          [&](CondBranchOp condBr) { workList.push_back(condBr.condition()); });
+      func.walk([&](CondBranchOp condBr) {
+        for (auto operand : condBr.getOperands()) {
+          workList.push_back(operand);
+        }
+      });
+
+      func.walk([&](BranchOp br) {
+        for (auto operand : br.getOperands()) {
+          workList.push_back(operand);
+        }
+      });
 
       DenseSet<Value> visitedValues;
       DenseSet<Operation *> visitedOps;
diff --git a/mlir/test/Dialect/Linalg/detensorize_br_operands.mlir b/mlir/test/Dialect/Linalg/detensorize_br_operands.mlir
new file mode 100644
index 000000000000..0c9d132eb086
--- /dev/null
+++ b/mlir/test/Dialect/Linalg/detensorize_br_operands.mlir
@@ -0,0 +1,49 @@
+// RUN: mlir-opt %s -split-input-file -allow-unregistered-dialect -linalg-detensorize | FileCheck %s
+
+// TODO: Detensoring breaks if %arg0 or %arg1 are passed directly as tensors. Fix that.
+func @if_true_test(%arg0: i1, %arg1: i32) -> tensor<i32> attributes {} {
+  %arg0_t = tensor.from_elements %arg0 : tensor<1xi1>
+  %arg0_t2 = linalg.tensor_collapse_shape %arg0_t [] : tensor<1xi1> into tensor<i1>
+
+  %arg1_t = tensor.from_elements %arg1 : tensor<1xi32>
+  %arg1_t2 = linalg.tensor_collapse_shape %arg1_t [] : tensor<1xi32> into tensor<i32>
+
+  %cst = constant dense<10> : tensor<i32>
+  %2 = linalg.init_tensor [] : tensor<i8>
+  %3 = linalg.generic
+    {indexing_maps = [affine_map<() -> ()>, affine_map<() -> ()>], iterator_types = []}
+    ins(%arg0_t2 : tensor<i1>)
+    outs(%2 : tensor<i8>) {
+  ^bb0(%arg2: i1, %arg3: i8):  // no predecessors
+    %10 = zexti %arg2 : i1 to i8
+    linalg.yield %10 : i8
+  } -> tensor<i8>
+  %4 = tensor.extract %3[] : tensor<i8>
+  %5 = trunci %4 : i8 to i1
+  cond_br %5, ^bb1, ^bb2(%arg1_t2 : tensor<i32>)
+^bb1:
+  %6 = linalg.init_tensor [] : tensor<i32>
+  %7 = linalg.generic
+    {indexing_maps = [affine_map<() -> ()>, affine_map<() -> ()>, affine_map<() -> ()>], iterator_types = []}
+    ins(%arg1_t2, %cst : tensor<i32>, tensor<i32>)
+    outs(%6 : tensor<i32>) {
+  ^bb0(%arg2: i32, %arg3: i32, %arg4: i32):  // no predecessors
+    %10 = addi %arg2, %arg3 : i32
+    linalg.yield %10 : i32
+  } -> tensor<i32>
+  br ^bb2(%7 : tensor<i32>)
+^bb2(%8: tensor<i32>):
+  return %8 : tensor<i32>
+}
+
+// CHECK-LABEL:  func @if_true_test
+// CHECK-SAME:     (%[[arg0:.*]]: i1, %[[arg1:.*]]: i32)
+// CHECK-NEXT:     constant 10 : i32
+// CHECK-NEXT:     cond_br %[[arg0]], ^[[bb1:.*]], ^[[bb2:.*]](%[[arg1]] : i32)
+// CHECK-NEXT:   ^[[bb1]]:
+// CHECK-NEXT:     %[[add_res:.*]] = addi
+// CHECK-NEXT:     br ^[[bb2]](%[[add_res]] : i32)
+// CHECK-NEXT:   ^[[bb2]]
+// CHECK-NEXT:     tensor.from_elements
+// CHECK-NEXT:     %[[func_res:.*]] = linalg.tensor_collapse_shape
+// CHECK-NEXT:     return %[[func_res]]