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 // Test case: Most basic case of a 1:N decomposition, an identity function.
   4
   5 // CHECK-LABEL:   func @identity(
   6 // CHECK-SAME:                   %[[ARG0:.*]]: i1,
   7 // CHECK-SAME:                   %[[ARG1:.*]]: i32) -> (i1, i32) {
   8 // CHECK:           %[[ARG_MATERIALIZED:.*]] = "test.make_tuple"(%[[ARG0]], %[[ARG1]]) : (i1, i32) -> tuple<i1, i32>
   9 // CHECK:           %[[RET0:.*]] = "test.get_tuple_element"(%[[ARG_MATERIALIZED]]) {index = 0 : i32} : (tuple<i1, i32>) -> i1
  10 // CHECK:           %[[RET1:.*]] = "test.get_tuple_element"(%[[ARG_MATERIALIZED]]) {index = 1 : i32} : (tuple<i1, i32>) -> i32
  11 // CHECK:           return %[[RET0]], %[[RET1]] : i1, i32
  12 func @identity(%arg0: tuple<i1, i32>) -> tuple<i1, i32> {
  13   return %arg0 : tuple<i1, i32>
  14 }
  15
  16 // -----
  17
  18 // Test case: Ensure no materializations in the case of 1:1 decomposition.
  19
  20 // CHECK-LABEL:   func @identity_1_to_1_no_materializations(
  21 // CHECK-SAME:                                              %[[ARG0:.*]]: i1) -> i1 {
  22 // CHECK:           return %[[ARG0]] : i1
  23 func @identity_1_to_1_no_materializations(%arg0: tuple<i1>) -> tuple<i1> {
  24   return %arg0 : tuple<i1>
  25 }
  26
  27 // -----
  28
  29 // Test case: Type that needs to be recursively decomposed.
  30
  31 // CHECK-LABEL:   func @recursive_decomposition(
  32 // CHECK-SAME:                                   %[[ARG0:.*]]: i1) -> i1 {
  33 // CHECK:           return %[[ARG0]] : i1
  34 func @recursive_decomposition(%arg0: tuple<tuple<tuple<i1>>>) -> tuple<tuple<tuple<i1>>> {
  35   return %arg0 : tuple<tuple<tuple<i1>>>
  36 }
  37
  38 // -----
  39
  40 // Test case: Check decomposition of calls.
  41
  42 // CHECK-LABEL:   func private @callee(i1, i32) -> (i1, i32)
  43 func private @callee(tuple<i1, i32>) -> tuple<i1, i32>
  44
  45 // CHECK-LABEL:   func @caller(
  46 // CHECK-SAME:                 %[[ARG0:.*]]: i1,
  47 // CHECK-SAME:                 %[[ARG1:.*]]: i32) -> (i1, i32) {
  48 // CHECK:           %[[ARG_MATERIALIZED:.*]] = "test.make_tuple"(%[[ARG0]], %[[ARG1]]) : (i1, i32) -> tuple<i1, i32>
  49 // CHECK:           %[[CALL_ARG0:.*]] = "test.get_tuple_element"(%[[ARG_MATERIALIZED]]) {index = 0 : i32} : (tuple<i1, i32>) -> i1
  50 // CHECK:           %[[CALL_ARG1:.*]] = "test.get_tuple_element"(%[[ARG_MATERIALIZED]]) {index = 1 : i32} : (tuple<i1, i32>) -> i32
  51 // CHECK:           %[[DECOMPOSED:.*]]:2 = call @callee(%[[CALL_ARG0]], %[[CALL_ARG1]]) : (i1, i32) -> (i1, i32)
  52 // CHECK:           %[[CALL_RESULT_RECOMPOSED:.*]] = "test.make_tuple"(%[[DECOMPOSED]]#0, %[[DECOMPOSED]]#1) : (i1, i32) -> tuple<i1, i32>
  53 // CHECK:           %[[RET0:.*]] = "test.get_tuple_element"(%[[CALL_RESULT_RECOMPOSED]]) {index = 0 : i32} : (tuple<i1, i32>) -> i1
  54 // CHECK:           %[[RET1:.*]] = "test.get_tuple_element"(%[[CALL_RESULT_RECOMPOSED]]) {index = 1 : i32} : (tuple<i1, i32>) -> i32
  55 // CHECK:           return %[[RET0]], %[[RET1]] : i1, i32
  56 func @caller(%arg0: tuple<i1, i32>) -> tuple<i1, i32> {
  57   %0 = call @callee(%arg0) : (tuple<i1, i32>) -> tuple<i1, i32>
  58   return %0 : tuple<i1, i32>
  59 }
  60
  61 // -----
  62
  63 // Test case: Type that decomposes to nothing (that is, a 1:0 decomposition).
  64
  65 // CHECK-LABEL:   func private @callee()
  66 func private @callee(tuple<>) -> tuple<>
  67 // CHECK-LABEL:   func @caller() {
  68 // CHECK:           call @callee() : () -> ()
  69 // CHECK:           return
  70 func @caller(%arg0: tuple<>) -> tuple<> {
  71   %0 = call @callee(%arg0) : (tuple<>) -> (tuple<>)
  72   return %0 : tuple<>
  73 }
  74
  75 // -----
  76
  77 // Test case: Ensure decompositions are inserted properly around results of
  78 // unconverted ops.
  79
  80 // CHECK-LABEL:   func @unconverted_op_result() -> (i1, i32) {
  81 // CHECK:           %[[UNCONVERTED_VALUE:.*]] = "test.source"() : () -> tuple<i1, i32>
  82 // CHECK:           %[[RET0:.*]] = "test.get_tuple_element"(%[[UNCONVERTED_VALUE]]) {index = 0 : i32} : (tuple<i1, i32>) -> i1
  83 // CHECK:           %[[RET1:.*]] = "test.get_tuple_element"(%[[UNCONVERTED_VALUE]]) {index = 1 : i32} : (tuple<i1, i32>) -> i32
  84 // CHECK:           return %[[RET0]], %[[RET1]] : i1, i32
  85 func @unconverted_op_result() -> tuple<i1, i32> {
  86   %0 = "test.source"() : () -> (tuple<i1, i32>)
  87   return %0 : tuple<i1, i32>
  88 }
  89
  90 // -----
  91
  92 // Test case: Check mixed decomposed and non-decomposed args.
  93 // This makes sure to test the cases if 1:0, 1:1, and 1:N decompositions.
  94
  95 // CHECK-LABEL:   func private @callee(i1, i2, i3, i4, i5, i6) -> (i1, i2, i3, i4, i5, i6)
  96 func private @callee(tuple<>, i1, tuple<i2>, i3, tuple<i4, i5>, i6) -> (tuple<>, i1, tuple<i2>, i3, tuple<i4, i5>, i6)
  97
  98 // CHECK-LABEL:   func @caller(
  99 // CHECK-SAME:                 %[[I1:.*]]: i1,
 100 // CHECK-SAME:                 %[[I2:.*]]: i2,
 101 // CHECK-SAME:                 %[[I3:.*]]: i3,
 102 // CHECK-SAME:                 %[[I4:.*]]: i4,
 103 // CHECK-SAME:                 %[[I5:.*]]: i5,
 104 // CHECK-SAME:                 %[[I6:.*]]: i6) -> (i1, i2, i3, i4, i5, i6) {
 105 // CHECK:           %[[ARG_TUPLE:.*]] = "test.make_tuple"(%[[I4]], %[[I5]]) : (i4, i5) -> tuple<i4, i5>
 106 // CHECK:           %[[ARG_TUPLE_0:.*]] = "test.get_tuple_element"(%[[ARG_TUPLE]]) {index = 0 : i32} : (tuple<i4, i5>) -> i4
 107 // CHECK:           %[[ARG_TUPLE_1:.*]] = "test.get_tuple_element"(%[[ARG_TUPLE]]) {index = 1 : i32} : (tuple<i4, i5>) -> i5
 108 // CHECK:           %[[CALL:.*]]:6 = call @callee(%[[I1]], %[[I2]], %[[I3]], %[[ARG_TUPLE_0]], %[[ARG_TUPLE_1]], %[[I6]]) : (i1, i2, i3, i4, i5, i6) -> (i1, i2, i3, i4, i5, i6)
 109 // CHECK:           %[[RET_TUPLE:.*]] = "test.make_tuple"(%[[CALL]]#3, %[[CALL]]#4) : (i4, i5) -> tuple<i4, i5>
 110 // CHECK:           %[[RET_TUPLE_0:.*]] = "test.get_tuple_element"(%[[RET_TUPLE]]) {index = 0 : i32} : (tuple<i4, i5>) -> i4
 111 // CHECK:           %[[RET_TUPLE_1:.*]] = "test.get_tuple_element"(%[[RET_TUPLE]]) {index = 1 : i32} : (tuple<i4, i5>) -> i5
 112 // CHECK:           return %[[CALL]]#0, %[[CALL]]#1, %[[CALL]]#2, %[[RET_TUPLE_0]], %[[RET_TUPLE_1]], %[[CALL]]#5 : i1, i2, i3, i4, i5, i6
 113 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) {
 114   %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)
 115   return %0, %1, %2, %3, %4, %5 : tuple<>, i1, tuple<i2>, i3, tuple<i4, i5>, i6
 116 }