xref: /llvm-project/mlir/test/Dialect/MemRef/make-loop-independent.mlir (revision e4384149b58f7c3d19c5d38bc46038c660b77ca9)
176100870SMatthias Springer// RUN: mlir-opt %s -allow-unregistered-dialect \
2*e4384149SOleksandr "Alex" Zinenko// RUN:     -transform-interpreter -canonicalize \
376100870SMatthias Springer// RUN:     -split-input-file -verify-diagnostics | FileCheck %s
476100870SMatthias Springer
576100870SMatthias Springer// CHECK: #[[$map:.*]] = affine_map<()[s0] -> (s0 - 1)>
676100870SMatthias Springer// CHECK-LABEL: func @make_alloca_loop_independent(
776100870SMatthias Springer//  CHECK-SAME:     %[[lb:.*]]: index, %[[ub:.*]]: index, %[[step:.*]]: index)
876100870SMatthias Springerfunc.func @make_alloca_loop_independent(%lb: index, %ub: index, %step: index) {
976100870SMatthias Springer  %cst = arith.constant 5.5 : f32
1076100870SMatthias Springer  %c0 = arith.constant 0 : index
1176100870SMatthias Springer  // CHECK: scf.for %[[iv:.*]] = %[[lb]] to %[[ub]]
1276100870SMatthias Springer  scf.for %i = %lb to %ub step %step {
1376100870SMatthias Springer    // CHECK: %[[sz:.*]] = affine.apply #[[$map]]()[%[[ub]]]
1476100870SMatthias Springer    // CHECK: %[[alloca:.*]] = memref.alloca(%[[sz]])
1576100870SMatthias Springer    // CHECK: %[[subview:.*]] = memref.subview %[[alloca]][0] [%[[iv]]] [1] : memref<?xf32> to memref<?xf32, strided<[1]>>
1676100870SMatthias Springer    // CHECK: %[[cast:.*]] = builtin.unrealized_conversion_cast %[[subview]] : memref<?xf32, strided<[1]>> to memref<?xf32>
1776100870SMatthias Springer    %alloc = memref.alloca(%i) : memref<?xf32>
1876100870SMatthias Springer
1976100870SMatthias Springer    // memref.subview has special handling.
2076100870SMatthias Springer    // CHECK: %[[subview2:.*]] = memref.subview %[[subview]][1] [5] [1] : memref<?xf32, strided<[1]>> to memref<5xf32, strided<[1], offset: 1>>
2176100870SMatthias Springer    %view = memref.subview %alloc[1][5][1] : memref<?xf32> to memref<5xf32, strided<[1], offset: 1>>
2276100870SMatthias Springer
2376100870SMatthias Springer    // This op takes a memref but does not produce one. The new alloc is used
2476100870SMatthias Springer    // directly.
2576100870SMatthias Springer    // CHECK: "test.some_use"(%[[subview2]])
2676100870SMatthias Springer    "test.some_use"(%view) : (memref<5xf32, strided<[1], offset: 1>>) -> ()
2776100870SMatthias Springer
2876100870SMatthias Springer    // This op produces a memref, so the new alloc cannot be used directly.
2976100870SMatthias Springer    // It is wrapped in a unrealized_conversion_cast.
3076100870SMatthias Springer    // CHECK: "test.another_use"(%[[cast]]) : (memref<?xf32>) -> memref<?xf32>
3176100870SMatthias Springer    "test.another_use"(%alloc) : (memref<?xf32>) -> (memref<?xf32>)
3276100870SMatthias Springer
3376100870SMatthias Springer    // CHECK: memref.store %{{.*}}, %[[subview]]
3476100870SMatthias Springer    memref.store %cst, %alloc[%c0] : memref<?xf32>
3576100870SMatthias Springer  }
3676100870SMatthias Springer  return
3776100870SMatthias Springer}
38*e4384149SOleksandr "Alex" Zinenkomodule attributes {transform.with_named_sequence} {
39*e4384149SOleksandr "Alex" Zinenko  transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
402f3ac28cSAlex Zinenko    %0 = transform.structured.match ops{["memref.alloca"]} in %arg1 : (!transform.any_op) -> !transform.any_op
412f3ac28cSAlex Zinenko    %1 = transform.memref.make_loop_independent %0 {num_loops = 1} : (!transform.any_op) -> !transform.any_op
42*e4384149SOleksandr "Alex" Zinenko    transform.yield
43*e4384149SOleksandr "Alex" Zinenko  }
4476100870SMatthias Springer}
4576100870SMatthias Springer
4676100870SMatthias Springer// -----
4776100870SMatthias Springer
4876100870SMatthias Springer// CHECK: #[[$map:.*]] = affine_map<(d0) -> (-d0 + 128)>
4976100870SMatthias Springer// CHECK-LABEL: func @make_alloca_loop_independent_static(
5076100870SMatthias Springerfunc.func @make_alloca_loop_independent_static(%step: index) {
5176100870SMatthias Springer  %cst = arith.constant 5.5 : f32
5276100870SMatthias Springer  %c0 = arith.constant 0 : index
5376100870SMatthias Springer  %ub = arith.constant 128 : index
5476100870SMatthias Springer  // CHECK: scf.for %[[iv:.*]] =
5576100870SMatthias Springer  scf.for %i = %c0 to %ub step %step {
5676100870SMatthias Springer    // CHECK: %[[sz:.*]] = affine.apply #[[$map]](%[[iv]])
5776100870SMatthias Springer    %sz = affine.apply affine_map<(d0)[s0] -> (-d0 + s0)>(%i)[%ub]
5876100870SMatthias Springer
5976100870SMatthias Springer    // CHECK: %[[alloca:.*]] = memref.alloca() : memref<128xf32>
6076100870SMatthias Springer    // CHECK: %[[subview:.*]] = memref.subview %[[alloca]][0] [%[[sz]]] [1] : memref<128xf32> to memref<?xf32, strided<[1]>>
6176100870SMatthias Springer    %alloc = memref.alloca(%sz) : memref<?xf32>
6276100870SMatthias Springer
6376100870SMatthias Springer    // CHECK: memref.store %{{.*}}, %[[subview]]
6476100870SMatthias Springer    memref.store %cst, %alloc[%c0] : memref<?xf32>
6576100870SMatthias Springer
6676100870SMatthias Springer    // CHECK: vector.print %[[sz]]
6776100870SMatthias Springer    %dim = memref.dim %alloc, %c0 : memref<?xf32>
6876100870SMatthias Springer    vector.print %dim : index
6976100870SMatthias Springer  }
7076100870SMatthias Springer  return
7176100870SMatthias Springer}
72*e4384149SOleksandr "Alex" Zinenkomodule attributes {transform.with_named_sequence} {
73*e4384149SOleksandr "Alex" Zinenko  transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
742f3ac28cSAlex Zinenko    %0 = transform.structured.match ops{["memref.alloca"]} in %arg1 : (!transform.any_op) -> !transform.any_op
752f3ac28cSAlex Zinenko    %1 = transform.memref.make_loop_independent %0 {num_loops = 1} : (!transform.any_op) -> !transform.any_op
76*e4384149SOleksandr "Alex" Zinenko    transform.yield
77*e4384149SOleksandr "Alex" Zinenko  }
7876100870SMatthias Springer}
79