1// RUN: mlir-opt %s -pass-pipeline="func.func(convert-scf-to-cf,memref-expand,convert-arith-to-llvm),convert-memref-to-llvm,convert-func-to-llvm,reconcile-unrealized-casts" \
2// RUN: | mlir-cpu-runner -e main -entry-point-result=void \
3// RUN: -shared-libs=%mlir_runner_utils_dir/libmlir_runner_utils%shlibext,%mlir_runner_utils_dir/libmlir_c_runner_utils%shlibext \
4// RUN: | FileCheck %s
5
6
7func.func private @printMemrefF32(memref<*xf32>) attributes { llvm.emit_c_interface }
8
9func.func @main() -> () {
10  %c0 = arith.constant 0 : index
11  %c1 = arith.constant 1 : index
12
13  // Initialize input.
14  %input = memref.alloc() : memref<2x3xf32>
15  %dim_x = memref.dim %input, %c0 : memref<2x3xf32>
16  %dim_y = memref.dim %input, %c1 : memref<2x3xf32>
17  scf.parallel (%i, %j) = (%c0, %c0) to (%dim_x, %dim_y) step (%c1, %c1) {
18    %prod = arith.muli %i,  %dim_y : index
19    %val = arith.addi %prod, %j : index
20    %val_i64 = arith.index_cast %val : index to i64
21    %val_f32 = arith.sitofp %val_i64 : i64 to f32
22    memref.store %val_f32, %input[%i, %j] : memref<2x3xf32>
23  }
24  %unranked_input = memref.cast %input : memref<2x3xf32> to memref<*xf32>
25  call @printMemrefF32(%unranked_input) : (memref<*xf32>) -> ()
26  // CHECK: rank = 2 offset = 0 sizes = [2, 3] strides = [3, 1]
27  // CHECK-NEXT: [0,   1,   2]
28  // CHECK-NEXT: [3,   4,   5]
29
30  // Initialize shape.
31  %shape = memref.alloc() : memref<2xindex>
32  %c2 = arith.constant 2 : index
33  %c3 = arith.constant 3 : index
34  memref.store %c3, %shape[%c0] : memref<2xindex>
35  memref.store %c2, %shape[%c1] : memref<2xindex>
36
37  // Test cases.
38  call @reshape_ranked_memref_to_ranked(%input, %shape)
39    : (memref<2x3xf32>, memref<2xindex>) -> ()
40  call @reshape_unranked_memref_to_ranked(%input, %shape)
41    : (memref<2x3xf32>, memref<2xindex>) -> ()
42  call @reshape_ranked_memref_to_unranked(%input, %shape)
43    : (memref<2x3xf32>, memref<2xindex>) -> ()
44  call @reshape_unranked_memref_to_unranked(%input, %shape)
45    : (memref<2x3xf32>, memref<2xindex>) -> ()
46  memref.dealloc %input : memref<2x3xf32>
47  memref.dealloc %shape : memref<2xindex>
48  return
49}
50
51func.func @reshape_ranked_memref_to_ranked(%input : memref<2x3xf32>,
52                                      %shape : memref<2xindex>) {
53  %output = memref.reshape %input(%shape)
54                : (memref<2x3xf32>, memref<2xindex>) -> memref<?x?xf32>
55
56  %unranked_output = memref.cast %output : memref<?x?xf32> to memref<*xf32>
57  call @printMemrefF32(%unranked_output) : (memref<*xf32>) -> ()
58  // CHECK: rank = 2 offset = 0 sizes = [3, 2] strides = [2, 1] data =
59  // CHECK: [0,   1],
60  // CHECK: [2,   3],
61  // CHECK: [4,   5]
62  return
63}
64
65func.func @reshape_unranked_memref_to_ranked(%input : memref<2x3xf32>,
66                                        %shape : memref<2xindex>) {
67  %unranked_input = memref.cast %input : memref<2x3xf32> to memref<*xf32>
68  %output = memref.reshape %input(%shape)
69                : (memref<2x3xf32>, memref<2xindex>) -> memref<?x?xf32>
70
71  %unranked_output = memref.cast %output : memref<?x?xf32> to memref<*xf32>
72  call @printMemrefF32(%unranked_output) : (memref<*xf32>) -> ()
73  // CHECK: rank = 2 offset = 0 sizes = [3, 2] strides = [2, 1] data =
74  // CHECK: [0,   1],
75  // CHECK: [2,   3],
76  // CHECK: [4,   5]
77  return
78}
79
80func.func @reshape_ranked_memref_to_unranked(%input : memref<2x3xf32>,
81                                        %shape : memref<2xindex>) {
82  %dyn_size_shape = memref.cast %shape : memref<2xindex> to memref<?xindex>
83  %output = memref.reshape %input(%dyn_size_shape)
84                : (memref<2x3xf32>, memref<?xindex>) -> memref<*xf32>
85
86  call @printMemrefF32(%output) : (memref<*xf32>) -> ()
87  // CHECK: rank = 2 offset = 0 sizes = [3, 2] strides = [2, 1] data =
88  // CHECK: [0,   1],
89  // CHECK: [2,   3],
90  // CHECK: [4,   5]
91  return
92}
93
94func.func @reshape_unranked_memref_to_unranked(%input : memref<2x3xf32>,
95                                          %shape : memref<2xindex>) {
96  %unranked_input = memref.cast %input : memref<2x3xf32> to memref<*xf32>
97  %dyn_size_shape = memref.cast %shape : memref<2xindex> to memref<?xindex>
98  %output = memref.reshape %input(%dyn_size_shape)
99                : (memref<2x3xf32>, memref<?xindex>) -> memref<*xf32>
100
101  call @printMemrefF32(%output) : (memref<*xf32>) -> ()
102  // CHECK: rank = 2 offset = 0 sizes = [3, 2] strides = [2, 1] data =
103  // CHECK: [0,   1],
104  // CHECK: [2,   3],
105  // CHECK: [4,   5]
106  return
107}
108