diff --git a/mlir/lib/Dialect/Linalg/Transforms/Detensorize.cpp b/mlir/lib/Dialect/Linalg/Transforms/Detensorize.cpp --- a/mlir/lib/Dialect/Linalg/Transforms/Detensorize.cpp +++ b/mlir/lib/Dialect/Linalg/Transforms/Detensorize.cpp @@ -297,8 +297,17 @@ DenseSet &blockArgsToDetensor) override { SmallVector 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 visitedValues; DenseSet 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 --- /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 attributes {} { + %arg0_t = tensor.from_elements %arg0 : tensor<1xi1> + %arg0_t2 = linalg.tensor_collapse_shape %arg0_t [] : tensor<1xi1> into tensor + + %arg1_t = tensor.from_elements %arg1 : tensor<1xi32> + %arg1_t2 = linalg.tensor_collapse_shape %arg1_t [] : tensor<1xi32> into tensor + + %cst = constant dense<10> : tensor + %2 = linalg.init_tensor [] : tensor + %3 = linalg.generic + {indexing_maps = [affine_map<() -> ()>, affine_map<() -> ()>], iterator_types = []} + ins(%arg0_t2 : tensor) + outs(%2 : tensor) { + ^bb0(%arg2: i1, %arg3: i8): // no predecessors + %10 = zexti %arg2 : i1 to i8 + linalg.yield %10 : i8 + } -> tensor + %4 = tensor.extract %3[] : tensor + %5 = trunci %4 : i8 to i1 + cond_br %5, ^bb1, ^bb2(%arg1_t2 : tensor) +^bb1: + %6 = linalg.init_tensor [] : tensor + %7 = linalg.generic + {indexing_maps = [affine_map<() -> ()>, affine_map<() -> ()>, affine_map<() -> ()>], iterator_types = []} + ins(%arg1_t2, %cst : tensor, tensor) + outs(%6 : tensor) { + ^bb0(%arg2: i32, %arg3: i32, %arg4: i32): // no predecessors + %10 = addi %arg2, %arg3 : i32 + linalg.yield %10 : i32 + } -> tensor + br ^bb2(%7 : tensor) +^bb2(%8: tensor): + return %8 : tensor +} + +// 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]]