mlir/test/Transforms/decompose-call-graph-types.mlir

   1 // RUN: mlir-opt %s -split-input-file -test-decompose-call-graph-types | FileCheck %s
   2
   3 // RUN: mlir-opt %s -split-input-file \
   4 // RUN:   -test-one-to-n-type-conversion="convert-func-ops" \
   5 // RUN: | FileCheck %s --check-prefix=CHECK-12N
   6
   7 // Test case: Most basic case of a 1:N decomposition, an identity function.
   8
   9 // CHECK-LABEL:   func @identity(
  10 // CHECK-SAME:                   %[[ARG0:.*]]: i1,
  11 // CHECK-SAME:                   %[[ARG1:.*]]: i32) -> (i1, i32) {
  12 // CHECK:           return %[[ARG0]], %[[ARG1]] : i1, i32
  13 // CHECK-12N-LABEL:   func @identity(
  14 // CHECK-12N-SAME:                   %[[ARG0:.*]]: i1,
  15 // CHECK-12N-SAME:                   %[[ARG1:.*]]: i32) -> (i1, i32) {
  16 // CHECK-12N:           return %[[ARG0]], %[[ARG1]] : i1, i32
  17 func.func @identity(%arg0: tuple<i1, i32>) -> tuple<i1, i32> {
  18   return %arg0 : tuple<i1, i32>
  19 }
  20
  21 // -----
  22
  23 // Test case: Ensure no materializations in the case of 1:1 decomposition.
  24
  25 // CHECK-LABEL:   func @identity_1_to_1_no_materializations(
  26 // CHECK-SAME:                                              %[[ARG0:.*]]: i1) -> i1 {
  27 // CHECK:           return %[[ARG0]] : i1
  28 // CHECK-12N-LABEL:   func @identity_1_to_1_no_materializations(
  29 // CHECK-12N-SAME:                                              %[[ARG0:.*]]: i1) -> i1 {
  30 // CHECK-12N:           return %[[ARG0]] : i1
  31 func.func @identity_1_to_1_no_materializations(%arg0: tuple<i1>) -> tuple<i1> {
  32   return %arg0 : tuple<i1>
  33 }
  34
  35 // -----
  36
  37 // Test case: Type that needs to be recursively decomposed.
  38
  39 // CHECK-LABEL:   func @recursive_decomposition(
  40 // CHECK-SAME:                                   %[[ARG0:.*]]: i1) -> i1 {
  41 // CHECK:           return %[[ARG0]] : i1
  42 // CHECK-12N-LABEL:   func @recursive_decomposition(
  43 // CHECK-12N-SAME:                                   %[[ARG0:.*]]: i1) -> i1 {
  44 // CHECK-12N:           return %[[ARG0]] : i1
  45 func.func @recursive_decomposition(%arg0: tuple<tuple<tuple<i1>>>) -> tuple<tuple<tuple<i1>>> {
  46   return %arg0 : tuple<tuple<tuple<i1>>>
  47 }
  48
  49 // -----
  50
  51 // Test case: Type that needs to be recursively decomposed at different recursion depths.
  52
  53 // CHECK-LABEL:   func @mixed_recursive_decomposition(
  54 // CHECK-SAME:                 %[[ARG0:.*]]: i1,
  55 // CHECK-SAME:                 %[[ARG1:.*]]: i2) -> (i1, i2) {
  56 // CHECK:           return %[[ARG0]], %[[ARG1]] : i1, i2
  57 // CHECK-12N-LABEL:   func @mixed_recursive_decomposition(
  58 // CHECK-12N-SAME:                 %[[ARG0:.*]]: i1,
  59 // CHECK-12N-SAME:                 %[[ARG1:.*]]: i2) -> (i1, i2) {
  60 // CHECK-12N:           return %[[ARG0]], %[[ARG1]] : i1, i2
  61 func.func @mixed_recursive_decomposition(%arg0: tuple<tuple<>, tuple<i1>, tuple<tuple<i2>>>) -> tuple<tuple<>, tuple<i1>, tuple<tuple<i2>>> {
  62   return %arg0 : tuple<tuple<>, tuple<i1>, tuple<tuple<i2>>>
  63 }
  64
  65 // -----
  66
  67 // Test case: Check decomposition of calls.
  68
  69 // CHECK-LABEL:   func private @callee(i1, i32) -> (i1, i32)
  70 // CHECK-12N-LABEL:   func private @callee(i1, i32) -> (i1, i32)
  71 func.func private @callee(tuple<i1, i32>) -> tuple<i1, i32>
  72
  73 // CHECK-LABEL:   func @caller(
  74 // CHECK-SAME:                 %[[ARG0:.*]]: i1,
  75 // CHECK-SAME:                 %[[ARG1:.*]]: i32) -> (i1, i32) {
  76 // CHECK:           %[[V0:.*]]:2 = call @callee(%[[ARG0]], %[[ARG1]]) : (i1, i32) -> (i1, i32)
  77 // CHECK:           return %[[V0]]#0, %[[V0]]#1 : i1, i32
  78 // CHECK-12N-LABEL:   func @caller(
  79 // CHECK-12N-SAME:                 %[[ARG0:.*]]: i1,
  80 // CHECK-12N-SAME:                 %[[ARG1:.*]]: i32) -> (i1, i32) {
  81 // CHECK-12N:           %[[V0:.*]]:2 = call @callee(%[[ARG0]], %[[ARG1]]) : (i1, i32) -> (i1, i32)
  82 // CHECK-12N:           return %[[V0]]#0, %[[V0]]#1 : i1, i32
  83 func.func @caller(%arg0: tuple<i1, i32>) -> tuple<i1, i32> {
  84   %0 = call @callee(%arg0) : (tuple<i1, i32>) -> tuple<i1, i32>
  85   return %0 : tuple<i1, i32>
  86 }
  87
  88 // -----
  89
  90 // Test case: Type that decomposes to nothing (that is, a 1:0 decomposition).
  91
  92 // CHECK-LABEL:   func private @callee()
  93 // CHECK-12N-LABEL:   func private @callee()
  94 func.func private @callee(tuple<>) -> tuple<>
  95
  96 // CHECK-LABEL:   func @caller() {
  97 // CHECK:           call @callee() : () -> ()
  98 // CHECK:           return
  99 // CHECK-12N-LABEL:   func @caller() {
 100 // CHECK-12N:           call @callee() : () -> ()
 101 // CHECK-12N:           return
 102 func.func @caller(%arg0: tuple<>) -> tuple<> {
 103   %0 = call @callee(%arg0) : (tuple<>) -> (tuple<>)
 104   return %0 : tuple<>
 105 }
 106
 107 // -----
 108
 109 // Test case: Ensure decompositions are inserted properly around results of
 110 // unconverted ops.
 111
 112 // CHECK-LABEL:   func @unconverted_op_result() -> (i1, i32) {
 113 // CHECK:           %[[UNCONVERTED_VALUE:.*]] = "test.source"() : () -> tuple<i1, i32>
 114 // CHECK:           %[[RET0:.*]] = "test.get_tuple_element"(%[[UNCONVERTED_VALUE]]) <{index = 0 : i32}> : (tuple<i1, i32>) -> i1
 115 // CHECK:           %[[RET1:.*]] = "test.get_tuple_element"(%[[UNCONVERTED_VALUE]]) <{index = 1 : i32}> : (tuple<i1, i32>) -> i32
 116 // CHECK:           return %[[RET0]], %[[RET1]] : i1, i32
 117 // CHECK-12N-LABEL:   func @unconverted_op_result() -> (i1, i32) {
 118 // CHECK-12N:           %[[UNCONVERTED_VALUE:.*]] = "test.source"() : () -> tuple<i1, i32>
 119 // CHECK-12N:           %[[RET0:.*]] = "test.get_tuple_element"(%[[UNCONVERTED_VALUE]]) <{index = 0 : i32}> : (tuple<i1, i32>) -> i1
 120 // CHECK-12N:           %[[RET1:.*]] = "test.get_tuple_element"(%[[UNCONVERTED_VALUE]]) <{index = 1 : i32}> : (tuple<i1, i32>) -> i32
 121 // CHECK-12N:           return %[[RET0]], %[[RET1]] : i1, i32
 122 func.func @unconverted_op_result() -> tuple<i1, i32> {
 123   %0 = "test.source"() : () -> (tuple<i1, i32>)
 124   return %0 : tuple<i1, i32>
 125 }
 126
 127 // -----
 128
 129 // Test case: Ensure decompositions are inserted properly around results of
 130 // unconverted ops in the case of different nesting levels.
 131
 132 // CHECK-LABEL:   func @nested_unconverted_op_result(
 133 // CHECK-SAME:                 %[[ARG0:.*]]: i1,
 134 // CHECK-SAME:                 %[[ARG1:.*]]: i32) -> (i1, i32) {
 135 // CHECK:           %[[V0:.*]] = "test.make_tuple"(%[[ARG1]]) : (i32) -> tuple<i32>
 136 // CHECK:           %[[V1:.*]] = "test.make_tuple"(%[[ARG0]], %[[V0]]) : (i1, tuple<i32>) -> tuple<i1, tuple<i32>>
 137 // CHECK:           %[[V2:.*]] = "test.op"(%[[V1]]) : (tuple<i1, tuple<i32>>) -> tuple<i1, tuple<i32>>
 138 // CHECK:           %[[V3:.*]] = "test.get_tuple_element"(%[[V2]]) <{index = 0 : i32}> : (tuple<i1, tuple<i32>>) -> i1
 139 // CHECK:           %[[V4:.*]] = "test.get_tuple_element"(%[[V2]]) <{index = 1 : i32}> : (tuple<i1, tuple<i32>>) -> tuple<i32>
 140 // CHECK:           %[[V5:.*]] = "test.get_tuple_element"(%[[V4]]) <{index = 0 : i32}> : (tuple<i32>) -> i32
 141 // CHECK:           return %[[V3]], %[[V5]] : i1, i32
 142 // CHECK-12N-LABEL:   func @nested_unconverted_op_result(
 143 // CHECK-12N-SAME:                 %[[ARG0:.*]]: i1,
 144 // CHECK-12N-SAME:                 %[[ARG1:.*]]: i32) -> (i1, i32) {
 145 // CHECK-12N:           %[[V0:.*]] = "test.make_tuple"(%[[ARG1]]) : (i32) -> tuple<i32>
 146 // CHECK-12N:           %[[V1:.*]] = "test.make_tuple"(%[[ARG0]], %[[V0]]) : (i1, tuple<i32>) -> tuple<i1, tuple<i32>>
 147 // CHECK-12N:           %[[V2:.*]] = "test.op"(%[[V1]]) : (tuple<i1, tuple<i32>>) -> tuple<i1, tuple<i32>>
 148 // CHECK-12N:           %[[V3:.*]] = "test.get_tuple_element"(%[[V2]]) <{index = 0 : i32}> : (tuple<i1, tuple<i32>>) -> i1
 149 // CHECK-12N:           %[[V4:.*]] = "test.get_tuple_element"(%[[V2]]) <{index = 1 : i32}> : (tuple<i1, tuple<i32>>) -> tuple<i32>
 150 // CHECK-12N:           %[[V5:.*]] = "test.get_tuple_element"(%[[V4]]) <{index = 0 : i32}> : (tuple<i32>) -> i32
 151 // CHECK-12N:           return %[[V3]], %[[V5]] : i1, i32
 152 func.func @nested_unconverted_op_result(%arg: tuple<i1, tuple<i32>>) -> tuple<i1, tuple<i32>> {
 153   %0 = "test.op"(%arg) : (tuple<i1, tuple<i32>>) -> (tuple<i1, tuple<i32>>)
 154   return %0 : tuple<i1, tuple<i32>>
 155 }
 156
 157 // -----
 158
 159 // Test case: Check mixed decomposed and non-decomposed args.
 160 // This makes sure to test the cases if 1:0, 1:1, and 1:N decompositions.
 161
 162 // CHECK-LABEL:   func private @callee(i1, i2, i3, i4, i5, i6) -> (i1, i2, i3, i4, i5, i6)
 163 // CHECK-12N-LABEL:   func private @callee(i1, i2, i3, i4, i5, i6) -> (i1, i2, i3, i4, i5, i6)
 164 func.func private @callee(tuple<>, i1, tuple<i2>, i3, tuple<i4, i5>, i6) -> (tuple<>, i1, tuple<i2>, i3, tuple<i4, i5>, i6)
 165
 166 // CHECK-LABEL:   func @caller(
 167 // CHECK-SAME:                 %[[I1:.*]]: i1,
 168 // CHECK-SAME:                 %[[I2:.*]]: i2,
 169 // CHECK-SAME:                 %[[I3:.*]]: i3,
 170 // CHECK-SAME:                 %[[I4:.*]]: i4,
 171 // CHECK-SAME:                 %[[I5:.*]]: i5,
 172 // CHECK-SAME:                 %[[I6:.*]]: i6) -> (i1, i2, i3, i4, i5, i6) {
 173 // CHECK:           %[[CALL:.*]]:6 = call @callee(%[[I1]], %[[I2]], %[[I3]], %[[I4]], %[[I5]], %[[I6]]) : (i1, i2, i3, i4, i5, i6) -> (i1, i2, i3, i4, i5, i6)
 174 // CHECK:           return %[[CALL]]#0, %[[CALL]]#1, %[[CALL]]#2, %[[CALL]]#3, %[[CALL]]#4, %[[CALL]]#5 : i1, i2, i3, i4, i5, i6
 175 // CHECK-12N-LABEL:   func @caller(
 176 // CHECK-12N-SAME:                 %[[I1:.*]]: i1,
 177 // CHECK-12N-SAME:                 %[[I2:.*]]: i2,
 178 // CHECK-12N-SAME:                 %[[I3:.*]]: i3,
 179 // CHECK-12N-SAME:                 %[[I4:.*]]: i4,
 180 // CHECK-12N-SAME:                 %[[I5:.*]]: i5,
 181 // CHECK-12N-SAME:                 %[[I6:.*]]: i6) -> (i1, i2, i3, i4, i5, i6) {
 182 // CHECK-12N:           %[[CALL:.*]]:6 = call @callee(%[[I1]], %[[I2]], %[[I3]], %[[I4]], %[[I5]], %[[I6]]) : (i1, i2, i3, i4, i5, i6) -> (i1, i2, i3, i4, i5, i6)
 183 // CHECK-12N:           return %[[CALL]]#0, %[[CALL]]#1, %[[CALL]]#2, %[[CALL]]#3, %[[CALL]]#4, %[[CALL]]#5 : i1, i2, i3, i4, i5, i6
 184 func.func @caller(%arg0: tuple<>, %arg1: i1, %arg2: tuple<i2>, %arg3: i3, %arg4: tuple<i4, i5>, %arg5: i6) -> (tuple<>, i1, tuple<i2>, i3, tuple<i4, i5>, i6) {
 185   %0, %1, %2, %3, %4, %5 = call @callee(%arg0, %arg1, %arg2, %arg3, %arg4, %arg5) : (tuple<>, i1, tuple<i2>, i3, tuple<i4, i5>, i6) -> (tuple<>, i1, tuple<i2>, i3, tuple<i4, i5>, i6)
 186   return %0, %1, %2, %3, %4, %5 : tuple<>, i1, tuple<i2>, i3, tuple<i4, i5>, i6
 187 }