xref: /llvm-project/mlir/test/Dialect/Transform/test-loop-transforms.mlir (revision 2798b72ae7e5caad793169b77cbac47fe2362d0f)
1b9fe461eSMatthias Springer// RUN: mlir-opt %s --transform-interpreter --split-input-file \
2b9fe461eSMatthias Springer// RUN:     --verify-diagnostics | FileCheck %s
3b9fe461eSMatthias Springer
4b9fe461eSMatthias Springer// UNSUPPORTED: target=aarch64-pc-windows-msvc
5b9fe461eSMatthias Springer
6b9fe461eSMatthias Springer// CHECK-LABEL: func @test_loop_invariant_subset_hoisting(
7b9fe461eSMatthias Springer//  CHECK-SAME:     %[[arg:.*]]: tensor<?xf32>
8b9fe461eSMatthias Springerfunc.func @test_loop_invariant_subset_hoisting(%arg: tensor<?xf32>) -> tensor<?xf32> {
9b9fe461eSMatthias Springer  %lb = "test.foo"() : () -> (index)
10b9fe461eSMatthias Springer  %ub = "test.foo"() : () -> (index)
11b9fe461eSMatthias Springer  %step = "test.foo"() : () -> (index)
12b9fe461eSMatthias Springer  // CHECK: %[[extract:.*]] = tensor.extract_slice %[[arg]]
13b9fe461eSMatthias Springer  // CHECK: %[[for:.*]]:2 = scf.for {{.*}} iter_args(%[[t:.*]] = %[[arg]], %[[hoisted:.*]] = %[[extract]])
14b9fe461eSMatthias Springer  // expected-remark @below{{new loop op}}
15b9fe461eSMatthias Springer  %0 = scf.for %iv = %lb to %ub step %step iter_args(%t = %arg) -> (tensor<?xf32>) {
16b9fe461eSMatthias Springer    %1 = tensor.extract_slice %t[0][5][1] : tensor<?xf32> to tensor<5xf32>
17b9fe461eSMatthias Springer    // CHECK: %[[foo:.*]] = "test.foo"(%[[hoisted]])
18b9fe461eSMatthias Springer    %2 = "test.foo"(%1) : (tensor<5xf32>) -> (tensor<5xf32>)
19b9fe461eSMatthias Springer    // Obfuscate the IR by inserting at offset %sub instead of 0; both of them
20b9fe461eSMatthias Springer    // have the same value.
21b9fe461eSMatthias Springer    %3 = tensor.insert_slice %2 into %t[0][5][1] : tensor<5xf32> into tensor<?xf32>
22b9fe461eSMatthias Springer    // CHECK: scf.yield %[[t]], %[[foo]]
23b9fe461eSMatthias Springer    scf.yield %3 : tensor<?xf32>
24b9fe461eSMatthias Springer  }
25b9fe461eSMatthias Springer  // CHECK: %[[insert:.*]] = tensor.insert_slice %[[for]]#1 into %[[for]]#0
26b9fe461eSMatthias Springer  // CHECK: return %[[insert]]
27b9fe461eSMatthias Springer  return %0 : tensor<?xf32>
28b9fe461eSMatthias Springer}
29b9fe461eSMatthias Springer
30b9fe461eSMatthias Springermodule attributes {transform.with_named_sequence} {
31b9fe461eSMatthias Springer  transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) {
32b9fe461eSMatthias Springer    %0 = transform.structured.match ops{["scf.for"]} in %arg0 : (!transform.any_op) -> !transform.any_op
33b9fe461eSMatthias Springer    %1 = transform.structured.match ops{["tensor.extract_slice"]} in %arg0 : (!transform.any_op) -> !transform.any_op
34b9fe461eSMatthias Springer    %2 = transform.structured.match ops{["tensor.insert_slice"]} in %arg0 : (!transform.any_op) -> !transform.any_op
35b9fe461eSMatthias Springer
36b9fe461eSMatthias Springer    transform.loop.hoist_loop_invariant_subsets %0 : !transform.any_op
37b9fe461eSMatthias Springer    // Make sure that the handles are still valid (and were updated in case of
38b9fe461eSMatthias Springer    // the loop).
39b9fe461eSMatthias Springer
40f90b6090SOleksandr "Alex" Zinenko    %p = transform.num_associations %0 : (!transform.any_op) -> !transform.param<i64>
41b9fe461eSMatthias Springer    // expected-remark @below{{1}}
42*2798b72aSOleksandr "Alex" Zinenko    transform.debug.emit_param_as_remark %p : !transform.param<i64>
43*2798b72aSOleksandr "Alex" Zinenko    transform.debug.emit_remark_at %0, "new loop op" : !transform.any_op
44f90b6090SOleksandr "Alex" Zinenko    %p2 = transform.num_associations %1 : (!transform.any_op) -> !transform.param<i64>
45b9fe461eSMatthias Springer    // expected-remark @below{{1}}
46*2798b72aSOleksandr "Alex" Zinenko    transform.debug.emit_param_as_remark %p2 : !transform.param<i64>
47f90b6090SOleksandr "Alex" Zinenko    %p3 = transform.num_associations %2 : (!transform.any_op) -> !transform.param<i64>
48b9fe461eSMatthias Springer    // expected-remark @below{{1}}
49*2798b72aSOleksandr "Alex" Zinenko    transform.debug.emit_param_as_remark %p3 : !transform.param<i64>
50b9fe461eSMatthias Springer
51b9fe461eSMatthias Springer    transform.yield
52b9fe461eSMatthias Springer  }
53b9fe461eSMatthias Springer}
54b9fe461eSMatthias Springer
55b9fe461eSMatthias Springer// -----
56b9fe461eSMatthias Springer
57b9fe461eSMatthias Springer// Checks that transform ops from LoopExtensionOps and SCFTransformOps can be
58b9fe461eSMatthias Springer// used together.
59b9fe461eSMatthias Springer
60b9fe461eSMatthias Springer// CHECK-LABEL: func @test_mixed_loop_extension_scf_transform(
61b9fe461eSMatthias Springerfunc.func @test_mixed_loop_extension_scf_transform(%arg: tensor<?xf32>) -> tensor<?xf32> {
62b9fe461eSMatthias Springer  %lb = "test.foo"() : () -> (index)
63b9fe461eSMatthias Springer  %ub = "test.foo"() : () -> (index)
64b9fe461eSMatthias Springer  %step = "test.foo"() : () -> (index)
65b9fe461eSMatthias Springer  // CHECK: scf.for
66b9fe461eSMatthias Springer  // CHECK: scf.for
67b9fe461eSMatthias Springer  %0 = scf.for %iv = %lb to %ub step %step iter_args(%t = %arg) -> (tensor<?xf32>) {
68b9fe461eSMatthias Springer    %1 = "test.foo"(%t) : (tensor<?xf32>) -> (tensor<?xf32>)
69b9fe461eSMatthias Springer    scf.yield %1 : tensor<?xf32>
70b9fe461eSMatthias Springer  }
71b9fe461eSMatthias Springer  return %0 : tensor<?xf32>
72b9fe461eSMatthias Springer}
73b9fe461eSMatthias Springer
74b9fe461eSMatthias Springermodule attributes {transform.with_named_sequence} {
75b9fe461eSMatthias Springer  transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) {
76b9fe461eSMatthias Springer    %0 = transform.structured.match ops{["scf.for"]} in %arg0 : (!transform.any_op) -> !transform.any_op
77b9fe461eSMatthias Springer    transform.loop.hoist_loop_invariant_subsets %0 : !transform.any_op
78b9fe461eSMatthias Springer    transform.loop.unroll %0 { factor = 4 } : !transform.any_op
79b9fe461eSMatthias Springer    transform.yield
80b9fe461eSMatthias Springer  }
81b9fe461eSMatthias Springer}
82