detensorize_if.mlir | Explore in Territory

// RUN: mlir-opt %s -split-input-file -allow-unregistered-dialect -pass-pipeline="builtin.module(func.func(linalg-detensorize))" | FileCheck %s

#map0 = affine_map<() -> ()>

#attrs = {
  indexing_maps = [#map0, #map0, #map0],
  iterator_types = []
}

func.func @main() -> (tensor<i32>) attributes {} {
  %c0 = arith.constant 0 : i32
  %0 = tensor.from_elements %c0 : tensor<i32>
  %c10 = arith.constant 10 : i32
  %1 = tensor.from_elements %c10 : tensor<i32>
  cf.br ^bb1(%0 : tensor<i32>)

^bb1(%2: tensor<i32>):  // 2 preds: ^bb0, ^bb2
  %3 = tensor.empty() : tensor<i1>
  %4 = linalg.generic #attrs
    ins(%2, %1 : tensor<i32>, tensor<i32>)
    outs(%3 : tensor<i1>) {
    ^bb0(%arg0: i32, %arg1: i32, %arg2: i1):
      %8 = arith.cmpi slt, %arg0, %arg1 : i32
      linalg.yield %8 : i1
  } -> tensor<i1>
  %5 = tensor.extract %4[] : tensor<i1>
  cf.cond_br %5, ^bb2(%2 : tensor<i32>), ^bb3(%2 : tensor<i32>)

^bb2(%6: tensor<i32>):  // pred: ^bb1
  %7 = tensor.empty() : tensor<i32>
  %8 = linalg.generic #attrs
    ins(%6, %6 : tensor<i32>, tensor<i32>)
    outs(%7 : tensor<i32>) {
    ^bb0(%arg0: i32, %arg1: i32, %arg2: i32):
      %9 = arith.addi %arg0, %arg1 : i32
      linalg.yield %9 : i32
  } -> tensor<i32>
  cf.br ^bb3(%8 : tensor<i32>)

^bb3(%10: tensor<i32>):  // pred: ^bb1
  return %10 : tensor<i32>
}

// CHECK-LABEL:  func @main()
// CHECK-DAG:     %[[cst:.*]] = arith.constant dense<0>
// CHECK-DAG:     arith.constant true
// CHECK:         cf.br
// CHECK-NEXT:   ^[[bb1:.*]]:
// CHECK-NEXT:     cf.cond_br %{{.*}}, ^[[bb2:.*]], ^bb3
// CHECK-NEXT:   ^[[bb2]]
// CHECK-NEXT:     cf.br ^[[bb3:.*]]
// CHECK-NEXT:   ^[[bb3]]
// CHECK-NEXT:     return %[[cst]]
// CHECK-NEXT:   }

// -----

// Similar to the above test with one change: one of the block after the
// if-condition passes/forwards its tensor argument to another block.

#map0 = affine_map<() -> ()>

#attrs = {
  indexing_maps = [#map0, #map0, #map0],
  iterator_types = []
}

func.func @main() -> (tensor<i32>) attributes {} {
  %c0 = arith.constant 0 : i32
  %0 = tensor.from_elements %c0 : tensor<i32>
  %c10 = arith.constant 10 : i32
  %1 = tensor.from_elements %c10 : tensor<i32>
  cf.br ^bb1(%0 : tensor<i32>)

^bb1(%2: tensor<i32>):  // 2 preds: ^bb0, ^bb2
  %3 = tensor.empty() : tensor<i1>
  %4 = linalg.generic #attrs
    ins(%2, %1 : tensor<i32>, tensor<i32>)
    outs(%3 : tensor<i1>) {
    ^bb0(%arg0: i32, %arg1: i32, %arg2: i1):
      %8 = arith.cmpi slt, %arg0, %arg1 : i32
      linalg.yield %8 : i1
  } -> tensor<i1>
  %5 = tensor.extract %4[] : tensor<i1>
  cf.cond_br %5, ^bb2(%2 : tensor<i32>), ^bb3(%2 : tensor<i32>)

^bb2(%6: tensor<i32>):  // pred: ^bb1
  %7 = tensor.empty() : tensor<i32>
  %8 = linalg.generic #attrs
    ins(%6, %6 : tensor<i32>, tensor<i32>)
    outs(%7 : tensor<i32>) {
    ^bb0(%arg0: i32, %arg1: i32, %arg2: i32):
      %9 = arith.addi %arg0, %arg1 : i32
      linalg.yield %9 : i32
  } -> tensor<i32>
  cf.br ^bb3(%8 : tensor<i32>)

^bb3(%10: tensor<i32>):  // pred: ^bb1
  cf.br ^bb4(%10 : tensor<i32>)

^bb4(%11: tensor<i32>):  // pred: ^bb1
  return %11 : tensor<i32>
}

// CHECK-LABEL:  func @main()
// CHECK-DAG:     %[[cst:.*]] = arith.constant dense<0>
// CHECK-DAG:     arith.constant true
// CHECK:         cf.br ^[[bb1:.*]]
// CHECK-NEXT:   ^[[bb1:.*]]:
// CHECK-NEXT:     cf.cond_br %{{.*}}, ^[[bb2:.*]], ^bb3
// CHECK-NEXT:   ^[[bb2]]:
// CHECK-NEXT:     cf.br ^[[bb3:.*]]
// CHECK-NEXT:   ^[[bb3]]:
// CHECK-NEXT:     cf.br ^[[bb4:.*]]
// CHECK-NEXT:   ^[[bb4]]:
// CHECK-NEXT:     return %[[cst]]
// CHECK-NEXT:   }

// -----

#map0 = affine_map<() -> ()>

#attrs = {
  indexing_maps = [#map0, #map0, #map0],
  iterator_types = []
}

func.func @main() -> (tensor<i32>) attributes {} {
  %c0 = arith.constant 0 : i32
  %0 = tensor.from_elements %c0 : tensor<i32>
  %c10 = arith.constant 10 : i32
  %1 = tensor.from_elements %c10 : tensor<i32>
  cf.br ^bb1(%0 : tensor<i32>)

^bb1(%2: tensor<i32>):  // 2 preds: ^bb0, ^bb2
  %3 = tensor.empty() : tensor<i1>
  %4 = linalg.generic #attrs
    ins(%2, %1 : tensor<i32>, tensor<i32>)
    outs(%3 : tensor<i1>) {
    ^bb0(%arg0: i32, %arg1: i32, %arg2: i1):
      %8 = arith.cmpi slt, %arg0, %arg1 : i32
      linalg.yield %8 : i1
  } -> tensor<i1>
  %5 = tensor.extract %4[] : tensor<i1>
  // This cf.cond_br intentionally has bb2 as it's target for both branches. This
  // is to make sure that the "forward phase" of the cost-model correctly adds
  // the users of a block argument (in this case bb2's argument) to the work
  // list.
  cf.cond_br %5, ^bb2(%2 : tensor<i32>), ^bb2(%2 : tensor<i32>)

^bb2(%6: tensor<i32>):  // pred: ^bb1
  %12 = tensor.from_elements %c10 : tensor<i32>
  %7 = tensor.empty() : tensor<i32>
  %8 = linalg.generic #attrs
    ins(%6, %12 : tensor<i32>, tensor<i32>)
    outs(%7 : tensor<i32>) {
    ^bb0(%arg0: i32, %arg1: i32, %arg2: i32):
      %9 = arith.addi %arg0, %arg1 : i32
      linalg.yield %9 : i32
  } -> tensor<i32>
  cf.br ^bb3(%8 : tensor<i32>)

^bb3(%10: tensor<i32>):  // pred: ^bb1
  return %10 : tensor<i32>
}

// CHECK-LABEL:  func @main()
// CHECK-DAG:     %[[cst:.*]] = arith.constant dense<10>
// CHECK-DAG:     arith.constant true
// CHECK:         cf.br ^[[bb1:.*]]
// CHECK-NEXT:   ^[[bb1]]:
// CHECK-NEXT:     cf.cond_br %{{.*}}, ^[[bb2:.*]], ^bb2
// CHECK-NEXT:   ^[[bb2]]
// CHECK-NEXT:     cf.br ^[[bb3:.*]]
// CHECK-NEXT:   ^[[bb3]]
// CHECK-NEXT:     return %[[cst]]
// CHECK-NEXT:   }
llvm/mlir/test/Dialect/Linalg/detensorize_if.mlir
