xref: /llvm-project/mlir/test/Dialect/Func/func-transform.mlir (revision 42b160356fe5d3b41bf07c428d0142d3721b1d44)
1// RUN: mlir-opt %s --transform-interpreter -allow-unregistered-dialect --split-input-file | FileCheck %s
2
3// CHECK-LABEL: func.func @basic_cast_and_call
4func.func @basic_cast_and_call() {
5  // CHECK-NEXT: call @second()
6  "test.foo"() : () -> ()
7  // CHECK-NEXT: test.foo
8  // CHECK-NEXT: call @third()
9  func.return
10}
11
12func.func @second() {
13  "test.bar"() : () -> ()
14  func.return
15}
16
17func.func private @third()
18
19module attributes {transform.with_named_sequence} {
20  transform.named_sequence @__transform_main(%arg0: !transform.any_op) {
21    %funcs = transform.structured.match ops{["func.func"]} in %arg0 : (!transform.any_op) -> !transform.any_op
22    %f:3 = transform.split_handle %funcs : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op)
23    %foo = transform.structured.match ops{["test.foo"]} in %f#0 : (!transform.any_op) -> !transform.any_op
24    transform.func.cast_and_call @second before %foo : (!transform.any_op) -> !transform.any_op
25    transform.func.cast_and_call %f#2 after %foo : (!transform.any_op, !transform.any_op) -> !transform.any_op
26    transform.yield
27  }
28}
29
30// -----
31
32// CHECK-LABEL: func.func @non_empty_arg_and_out
33func.func @non_empty_arg_and_out(%arg0 : index) -> i32 {
34  // CHECK-NEXT: %[[FOO:.+]] = "test.foo"
35  %0 = "test.foo"(%arg0) : (index) -> (index)
36  // CHECK-NEXT: %[[CALL:.+]] = call @second(%[[FOO]]) : (index) -> i32
37  %1 = "test.bar"(%0) : (index) -> (i32)
38  // CHECK: return %[[CALL]] : i32
39  func.return %1 : i32
40}
41
42func.func private @second(%arg1 : index) -> i32
43
44module attributes {transform.with_named_sequence} {
45  transform.named_sequence @__transform_main(%arg0: !transform.any_op) {
46    %funcs = transform.structured.match ops{["func.func"]} in %arg0 : (!transform.any_op) -> !transform.any_op
47    %f:2 = transform.split_handle %funcs : (!transform.any_op) -> (!transform.any_op, !transform.any_op)
48    %foo = transform.structured.match ops{["test.foo"]} in %f#0 : (!transform.any_op) -> !transform.any_op
49    %bar = transform.structured.match ops{["test.bar"]} in %f#0 : (!transform.any_op) -> !transform.any_op
50    %in = transform.get_result %foo[0] : (!transform.any_op) -> !transform.any_value
51    %out = transform.get_result %bar[0] : (!transform.any_op) -> !transform.any_value
52    transform.func.cast_and_call %f#1(%in) -> %out before %bar
53        : (!transform.any_op, !transform.any_value,
54           !transform.any_value, !transform.any_op) -> !transform.any_op
55    transform.yield
56  }
57}
58
59// -----
60
61// CHECK-LABEL: func.func @multi_arg_and_result
62func.func @multi_arg_and_result(%arg0 : index) -> (index, index) {
63  // CHECK-NEXT: %[[FOO:.+]] = "test.foo"
64  %0 = "test.foo"(%arg0) : (index) -> (index)
65  %1 = "test.bar"(%0) : (index) -> (index)
66  %2 = "test.bar"(%0) : (index) -> (index)
67  // CHECK: %[[CALL:.+]]:2 = call @second(%[[FOO]], %[[FOO]]) : (index, index) -> (index, index)
68  // CHECK: return %[[CALL]]#0, %[[CALL]]#1 : index, index
69  func.return %1, %2 : index, index
70}
71
72func.func private @second(%arg1: index, %arg2: index) -> (index, index)
73
74module attributes {transform.with_named_sequence} {
75  transform.named_sequence @__transform_main(%arg0: !transform.any_op) {
76    %funcs = transform.structured.match ops{["func.func"]} in %arg0 : (!transform.any_op) -> !transform.any_op
77    %f:2 = transform.split_handle %funcs : (!transform.any_op) -> (!transform.any_op, !transform.any_op)
78    %foo = transform.structured.match ops{["test.foo"]} in %f#0 : (!transform.any_op) -> !transform.any_op
79    %bars = transform.structured.match ops{["test.bar"]} in %f#0 : (!transform.any_op) -> !transform.any_op
80    %in0 = transform.get_result %foo[0] : (!transform.any_op) -> !transform.any_value
81    %in1 = transform.get_result %foo[0] : (!transform.any_op) -> !transform.any_value
82    %ins = transform.merge_handles %in0, %in1 : !transform.any_value
83
84    %outs = transform.get_result %bars[0] : (!transform.any_op) -> !transform.any_value
85
86    transform.func.cast_and_call %f#1(%ins) -> %outs after %foo
87        : (!transform.any_op, !transform.any_value,
88           !transform.any_value, !transform.any_op) -> !transform.any_op
89    transform.yield
90  }
91}
92
93// -----
94
95// CHECK-LABEL: func.func @nested_call
96func.func @nested_call() {
97  // CHECK-NEXT: %[[ARG:.+]] = "test.arg"
98  // CHECK-NEXT: test.foo
99  %0 = "test.arg"() : () -> (index)
100  "test.foo"() ({
101    // CHECK-NEXT: call @second(%[[ARG]]) : (index) -> ()
102    "test.bar"(%0) : (index) -> ()
103  }) : () -> ()
104}
105
106func.func private @second(%arg1: index) -> ()
107
108module attributes {transform.with_named_sequence} {
109  transform.named_sequence @__transform_main(%arg0: !transform.any_op) {
110    %funcs = transform.structured.match ops{["func.func"]} in %arg0 : (!transform.any_op) -> !transform.any_op
111    %f:2 = transform.split_handle %funcs : (!transform.any_op) -> (!transform.any_op, !transform.any_op)
112    %arg = transform.structured.match ops{["test.arg"]} in %f#0 : (!transform.any_op) -> !transform.any_op
113    %bar = transform.structured.match ops{["test.bar"]} in %f#0 : (!transform.any_op) -> !transform.any_op
114    %in = transform.get_result %arg[0] : (!transform.any_op) -> !transform.any_value
115
116    transform.func.cast_and_call %f#1(%in) before %bar
117        : (!transform.any_op, !transform.any_value, !transform.any_op) -> !transform.any_op
118    transform.yield
119  }
120}
121