mlir/test/Conversion/OneToNTypeConversion/one-to-n-type-conversion.mlir

   1 // RUN: mlir-opt %s -split-input-file \
   2 // RUN:   -test-one-to-n-type-conversion="convert-tuple-ops" \
   3 // RUN: | FileCheck --check-prefix=CHECK-TUP %s
   4
   5 // RUN: mlir-opt %s -split-input-file \
   6 // RUN:   -test-one-to-n-type-conversion="convert-func-ops" \
   7 // RUN: | FileCheck --check-prefix=CHECK-FUNC %s
   8
   9 // RUN: mlir-opt %s -split-input-file \
  10 // RUN:   -test-one-to-n-type-conversion="convert-func-ops convert-tuple-ops" \
  11 // RUN: | FileCheck --check-prefix=CHECK-BOTH %s
  12
  13 // Test case: Matching nested packs and unpacks just disappear.
  14
  15 // CHECK-TUP-LABEL: func.func @pack_unpack(
  16 // CHECK-TUP-SAME:                          %[[ARG0:.*]]: i1,
  17 // CHECK-TUP-SAME:                          %[[ARG1:.*]]: i2) -> (i1, i2) {
  18 // CHECK-TUP-DAG:     return %[[ARG0]], %[[ARG1]] : i1, i2
  19 func.func @pack_unpack(%arg0: i1, %arg1: i2) -> (i1, i2) {
  20   %0 = "test.make_tuple"() : () -> tuple<>
  21   %1 = "test.make_tuple"(%arg1) : (i2) -> tuple<i2>
  22   %2 = "test.make_tuple"(%1) : (tuple<i2>) -> tuple<tuple<i2>>
  23   %3 = "test.make_tuple"(%0, %arg0, %2) : (tuple<>, i1, tuple<tuple<i2>>) -> tuple<tuple<>, i1, tuple<tuple<i2>>>
  24   %4 = "test.get_tuple_element"(%3) {index = 0 : i32} : (tuple<tuple<>, i1, tuple<tuple<i2>>>) -> tuple<>
  25   %5 = "test.get_tuple_element"(%3) {index = 1 : i32} : (tuple<tuple<>, i1, tuple<tuple<i2>>>) -> i1
  26   %6 = "test.get_tuple_element"(%3) {index = 2 : i32} : (tuple<tuple<>, i1, tuple<tuple<i2>>>) -> tuple<tuple<i2>>
  27   %7 = "test.get_tuple_element"(%6) {index = 0 : i32} : (tuple<tuple<i2>>) -> tuple<i2>
  28   %8 = "test.get_tuple_element"(%7) {index = 0 : i32} : (tuple<i2>) -> i2
  29   return %5, %8 : i1, i2
  30 }
  31
  32 // -----
  33
  34 // Test case: Appropriate materializations are created depending on which ops
  35 // are converted.
  36
  37 // If we only convert the tuple ops, the original `get_tuple_element` ops will
  38 // disappear but one target materialization will be inserted from the
  39 // unconverted function arguments to each of the return values (which have
  40 // redundancy among themselves).
  41 //
  42 // CHECK-TUP-LABEL: func.func @materializations_tuple_args(
  43 // CHECK-TUP-SAME:                                         %[[ARG0:.*]]: tuple<tuple<>, i1, tuple<tuple<i2>>>) -> (i1, i2) {
  44 // CHECK-TUP-DAG:     %[[V0:.*]] = "test.get_tuple_element"(%[[ARG0]]) <{index = 0 : i32}> : (tuple<tuple<>, i1, tuple<tuple<i2>>>) -> tuple<>
  45 // CHECK-TUP-DAG:     %[[V1:.*]] = "test.get_tuple_element"(%[[ARG0]]) <{index = 1 : i32}> : (tuple<tuple<>, i1, tuple<tuple<i2>>>) -> i1
  46 // CHECK-TUP-DAG:     %[[V2:.*]] = "test.get_tuple_element"(%[[ARG0]]) <{index = 2 : i32}> : (tuple<tuple<>, i1, tuple<tuple<i2>>>) -> tuple<tuple<i2>>
  47 // CHECK-TUP-DAG:     %[[V3:.*]] = "test.get_tuple_element"(%[[V2]]) <{index = 0 : i32}> : (tuple<tuple<i2>>) -> tuple<i2>
  48 // CHECK-TUP-DAG:     %[[V4:.*]] = "test.get_tuple_element"(%[[V3]]) <{index = 0 : i32}> : (tuple<i2>) -> i2
  49 // CHECK-TUP-DAG:     %[[V5:.*]] = "test.get_tuple_element"(%[[ARG0]]) <{index = 0 : i32}> : (tuple<tuple<>, i1, tuple<tuple<i2>>>) -> tuple<>
  50 // CHECK-TUP-DAG:     %[[V6:.*]] = "test.get_tuple_element"(%[[ARG0]]) <{index = 1 : i32}> : (tuple<tuple<>, i1, tuple<tuple<i2>>>) -> i1
  51 // CHECK-TUP-DAG:     %[[V7:.*]] = "test.get_tuple_element"(%[[ARG0]]) <{index = 2 : i32}> : (tuple<tuple<>, i1, tuple<tuple<i2>>>) -> tuple<tuple<i2>>
  52 // CHECK-TUP-DAG:     %[[V8:.*]] = "test.get_tuple_element"(%[[V7]]) <{index = 0 : i32}> : (tuple<tuple<i2>>) -> tuple<i2>
  53 // CHECK-TUP-DAG:     %[[V9:.*]] = "test.get_tuple_element"(%[[V8]]) <{index = 0 : i32}> : (tuple<i2>) -> i2
  54 // CHECK-TUP-DAG:     return %[[V1]], %[[V9]] : i1, i2
  55
  56 // If we only convert the func ops, argument materializations are created from
  57 // the converted tuple elements back to the tuples that the `get_tuple_element`
  58 // ops expect.
  59 //
  60 // CHECK-FUNC-LABEL: func.func @materializations_tuple_args(
  61 // CHECK-FUNC-SAME:                                         %[[ARG0:.*]]: i1,
  62 // CHECK-FUNC-SAME:                                         %[[ARG1:.*]]: i2) -> (i1, i2) {
  63 // CHECK-FUNC-DAG:     %[[V0:.*]] = "test.make_tuple"() : () -> tuple<>
  64 // CHECK-FUNC-DAG:     %[[V1:.*]] = "test.make_tuple"(%[[ARG1]]) : (i2) -> tuple<i2>
  65 // CHECK-FUNC-DAG:     %[[V2:.*]] = "test.make_tuple"(%[[V1]]) : (tuple<i2>) -> tuple<tuple<i2>>
  66 // CHECK-FUNC-DAG:     %[[V3:.*]] = "test.make_tuple"(%[[V0]], %[[ARG0]], %[[V2]]) : (tuple<>, i1, tuple<tuple<i2>>) -> tuple<tuple<>, i1, tuple<tuple<i2>>>
  67 // CHECK-FUNC-DAG:     %[[V4:.*]] = "test.get_tuple_element"(%[[V3]]) <{index = 0 : i32}> : (tuple<tuple<>, i1, tuple<tuple<i2>>>) -> tuple<>
  68 // CHECK-FUNC-DAG:     %[[V5:.*]] = "test.get_tuple_element"(%[[V3]]) <{index = 1 : i32}> : (tuple<tuple<>, i1, tuple<tuple<i2>>>) -> i1
  69 // CHECK-FUNC-DAG:     %[[V6:.*]] = "test.get_tuple_element"(%[[V3]]) <{index = 2 : i32}> : (tuple<tuple<>, i1, tuple<tuple<i2>>>) -> tuple<tuple<i2>>
  70 // CHECK-FUNC-DAG:     %[[V7:.*]] = "test.get_tuple_element"(%[[V6]]) <{index = 0 : i32}> : (tuple<tuple<i2>>) -> tuple<i2>
  71 // CHECK-FUNC-DAG:     %[[V8:.*]] = "test.get_tuple_element"(%[[V7]]) <{index = 0 : i32}> : (tuple<i2>) -> i2
  72 // CHECK-FUNC-DAG:     return %[[V5]], %[[V8]] : i1, i2
  73
  74 // If we convert both tuple and func ops, basically everything disappears.
  75 //
  76 // CHECK-BOTH-LABEL: func.func @materializations_tuple_args(
  77 // CHECK-BOTH-SAME:                                         %[[ARG0:.*]]: i1,
  78 // CHECK-BOTH-SAME:                                         %[[ARG1:.*]]: i2) -> (i1, i2) {
  79 // CHECK-BOTH-DAG:     return %[[ARG0]], %[[ARG1]] : i1, i2
  80
  81 func.func @materializations_tuple_args(%arg0: tuple<tuple<>, i1, tuple<tuple<i2>>>) -> (i1, i2) {
  82   %0 = "test.get_tuple_element"(%arg0) {index = 0 : i32} : (tuple<tuple<>, i1, tuple<tuple<i2>>>) -> tuple<>
  83   %1 = "test.get_tuple_element"(%arg0) {index = 1 : i32} : (tuple<tuple<>, i1, tuple<tuple<i2>>>) -> i1
  84   %2 = "test.get_tuple_element"(%arg0) {index = 2 : i32} : (tuple<tuple<>, i1, tuple<tuple<i2>>>) -> tuple<tuple<i2>>
  85   %3 = "test.get_tuple_element"(%2) {index = 0 : i32} : (tuple<tuple<i2>>) -> tuple<i2>
  86   %4 = "test.get_tuple_element"(%3) {index = 0 : i32} : (tuple<i2>) -> i2
  87   return %1, %4 : i1, i2
  88 }
  89 // -----
  90
  91 // Test case: Appropriate materializations are created depending on which ops
  92 // are converted.
  93
  94 // If we only convert the tuple ops, the original `make_tuple` ops will
  95 // disappear but a source materialization will be inserted from the result of
  96 // conversion (which, for `make_tuple`, are the original ops that get forwarded)
  97 // to the operands of the unconverted op with the original type (i.e.,
  98 // `return`).
  99
 100 // CHECK-TUP-LABEL: func.func @materializations_tuple_return(
 101 // CHECK-TUP-SAME:                                           %[[ARG0:.*]]: i1,
 102 // CHECK-TUP-SAME:                                           %[[ARG1:.*]]: i2) -> tuple<tuple<>, i1, tuple<tuple<i2>>> {
 103 // CHECK-TUP-DAG:     %[[V0:.*]] = "test.make_tuple"() : () -> tuple<>
 104 // CHECK-TUP-DAG:     %[[V1:.*]] = "test.make_tuple"(%[[ARG1]]) : (i2) -> tuple<i2>
 105 // CHECK-TUP-DAG:     %[[V2:.*]] = "test.make_tuple"(%[[V1]]) : (tuple<i2>) -> tuple<tuple<i2>>
 106 // CHECK-TUP-DAG:     %[[V3:.*]] = "test.make_tuple"(%[[V0]], %[[ARG0]], %[[V2]]) : (tuple<>, i1, tuple<tuple<i2>>) -> tuple<tuple<>, i1, tuple<tuple<i2>>>
 107 // CHECK-TUP-DAG:     return %[[V3]] : tuple<tuple<>, i1, tuple<tuple<i2>>>
 108
 109 // If we only convert the func ops, target materializations are created from
 110 // original tuples produced by `make_tuple` to its constituent elements that the
 111 // converted op (i.e., `return`) expect.
 112 //
 113 // CHECK-FUNC-LABEL: func.func @materializations_tuple_return(
 114 // CHECK-FUNC-SAME:                                           %[[ARG0:.*]]: i1,
 115 // CHECK-FUNC-SAME:                                           %[[ARG1:.*]]: i2) -> (i1, i2) {
 116 // CHECK-FUNC-DAG:     %[[V0:.*]] = "test.make_tuple"() : () -> tuple<>
 117 // CHECK-FUNC-DAG:     %[[V1:.*]] = "test.make_tuple"(%[[ARG1]]) : (i2) -> tuple<i2>
 118 // CHECK-FUNC-DAG:     %[[V2:.*]] = "test.make_tuple"(%[[V1]]) : (tuple<i2>) -> tuple<tuple<i2>>
 119 // CHECK-FUNC-DAG:     %[[V3:.*]] = "test.make_tuple"(%[[V0]], %[[ARG0]], %[[V2]]) : (tuple<>, i1, tuple<tuple<i2>>) -> tuple<tuple<>, i1, tuple<tuple<i2>>>
 120 // CHECK-FUNC-DAG:     %[[V4:.*]] = "test.get_tuple_element"(%[[V3]]) <{index = 0 : i32}> : (tuple<tuple<>, i1, tuple<tuple<i2>>>) -> tuple<>
 121 // CHECK-FUNC-DAG:     %[[V5:.*]] = "test.get_tuple_element"(%[[V3]]) <{index = 1 : i32}> : (tuple<tuple<>, i1, tuple<tuple<i2>>>) -> i1
 122 // CHECK-FUNC-DAG:     %[[V6:.*]] = "test.get_tuple_element"(%[[V3]]) <{index = 2 : i32}> : (tuple<tuple<>, i1, tuple<tuple<i2>>>) -> tuple<tuple<i2>>
 123 // CHECK-FUNC-DAG:     %[[V7:.*]] = "test.get_tuple_element"(%[[V6]]) <{index = 0 : i32}> : (tuple<tuple<i2>>) -> tuple<i2>
 124 // CHECK-FUNC-DAG:     %[[V8:.*]] = "test.get_tuple_element"(%[[V7]]) <{index = 0 : i32}> : (tuple<i2>) -> i2
 125 // CHECK-FUNC-DAG:     return %[[V5]], %[[V8]] : i1, i2
 126
 127 // If we convert both tuple and func ops, basically everything disappears.
 128 //
 129 // CHECK-BOTH-LABEL: func.func @materializations_tuple_return(
 130 // CHECK-BOTH-SAME:                                           %[[ARG0:.*]]: i1,
 131 // CHECK-BOTH-SAME:                                           %[[ARG1:.*]]: i2) -> (i1, i2) {
 132 // CHECK-BOTH-DAG:     return %[[ARG0]], %[[ARG1]] : i1, i2
 133
 134 func.func @materializations_tuple_return(%arg0: i1, %arg1: i2) -> tuple<tuple<>, i1, tuple<tuple<i2>>> {
 135   %0 = "test.make_tuple"() : () -> tuple<>
 136   %1 = "test.make_tuple"(%arg1) : (i2) -> tuple<i2>
 137   %2 = "test.make_tuple"(%1) : (tuple<i2>) -> tuple<tuple<i2>>
 138   %3 = "test.make_tuple"(%0, %arg0, %2) : (tuple<>, i1, tuple<tuple<i2>>) -> tuple<tuple<>, i1, tuple<tuple<i2>>>
 139   return %3 : tuple<tuple<>, i1, tuple<tuple<i2>>>
 140 }