1// RUN: mlir-opt %s -convert-scf-to-std -std-expand -convert-std-to-llvm \ 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 private @print_memref_f32(memref<*xf32>) attributes { llvm.emit_c_interface } 8 9func @main() -> () { 10 %c0 = constant 0 : index 11 %c1 = constant 1 : index 12 13 // Initialize input. 14 %input = alloc() : memref<2x3xf32> 15 %dim_x = dim %input, %c0 : memref<2x3xf32> 16 %dim_y = dim %input, %c1 : memref<2x3xf32> 17 scf.parallel (%i, %j) = (%c0, %c0) to (%dim_x, %dim_y) step (%c1, %c1) { 18 %prod = muli %i, %dim_y : index 19 %val = addi %prod, %j : index 20 %val_i64 = index_cast %val : index to i64 21 %val_f32 = sitofp %val_i64 : i64 to f32 22 store %val_f32, %input[%i, %j] : memref<2x3xf32> 23 } 24 %unranked_input = memref_cast %input : memref<2x3xf32> to memref<*xf32> 25 call @print_memref_f32(%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 = alloc() : memref<2xindex> 32 %c2 = constant 2 : index 33 %c3 = constant 3 : index 34 store %c3, %shape[%c0] : memref<2xindex> 35 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 return 47} 48 49func @reshape_ranked_memref_to_ranked(%input : memref<2x3xf32>, 50 %shape : memref<2xindex>) { 51 %output = memref_reshape %input(%shape) 52 : (memref<2x3xf32>, memref<2xindex>) -> memref<?x?xf32> 53 54 %unranked_output = memref_cast %output : memref<?x?xf32> to memref<*xf32> 55 call @print_memref_f32(%unranked_output) : (memref<*xf32>) -> () 56 // CHECK: rank = 2 offset = 0 sizes = [3, 2] strides = [2, 1] data = 57 // CHECK: [0, 1], 58 // CHECK: [2, 3], 59 // CHECK: [4, 5] 60 return 61} 62 63func @reshape_unranked_memref_to_ranked(%input : memref<2x3xf32>, 64 %shape : memref<2xindex>) { 65 %unranked_input = memref_cast %input : memref<2x3xf32> to memref<*xf32> 66 %output = memref_reshape %input(%shape) 67 : (memref<2x3xf32>, memref<2xindex>) -> memref<?x?xf32> 68 69 %unranked_output = memref_cast %output : memref<?x?xf32> to memref<*xf32> 70 call @print_memref_f32(%unranked_output) : (memref<*xf32>) -> () 71 // CHECK: rank = 2 offset = 0 sizes = [3, 2] strides = [2, 1] data = 72 // CHECK: [0, 1], 73 // CHECK: [2, 3], 74 // CHECK: [4, 5] 75 return 76} 77 78func @reshape_ranked_memref_to_unranked(%input : memref<2x3xf32>, 79 %shape : memref<2xindex>) { 80 %dyn_size_shape = memref_cast %shape : memref<2xindex> to memref<?xindex> 81 %output = memref_reshape %input(%dyn_size_shape) 82 : (memref<2x3xf32>, memref<?xindex>) -> memref<*xf32> 83 84 call @print_memref_f32(%output) : (memref<*xf32>) -> () 85 // CHECK: rank = 2 offset = 0 sizes = [3, 2] strides = [2, 1] data = 86 // CHECK: [0, 1], 87 // CHECK: [2, 3], 88 // CHECK: [4, 5] 89 return 90} 91 92func @reshape_unranked_memref_to_unranked(%input : memref<2x3xf32>, 93 %shape : memref<2xindex>) { 94 %unranked_input = memref_cast %input : memref<2x3xf32> to memref<*xf32> 95 %dyn_size_shape = memref_cast %shape : memref<2xindex> to memref<?xindex> 96 %output = memref_reshape %input(%dyn_size_shape) 97 : (memref<2x3xf32>, memref<?xindex>) -> memref<*xf32> 98 99 call @print_memref_f32(%output) : (memref<*xf32>) -> () 100 // CHECK: rank = 2 offset = 0 sizes = [3, 2] strides = [2, 1] data = 101 // CHECK: [0, 1], 102 // CHECK: [2, 3], 103 // CHECK: [4, 5] 104 return 105} 106