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