108 lines · plain
1// RUN: mlir-opt %s -pass-pipeline="builtin.module(func.func(convert-scf-to-cf,memref-expand,convert-arith-to-llvm),finalize-memref-to-llvm,convert-func-to-llvm,convert-cf-to-llvm,reconcile-unrealized-casts)" \2// RUN: | mlir-runner -e main -entry-point-result=void \3// RUN: -shared-libs=%mlir_runner_utils,%mlir_c_runner_utils \4// RUN: | FileCheck %s5 6 7func.func private @printMemrefF32(memref<*xf32>) attributes { llvm.emit_c_interface }8 9func.func @main() -> () {10 %c0 = arith.constant 0 : index11 %c1 = arith.constant 1 : index12 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 : index19 %val = arith.addi %prod, %j : index20 %val_i64 = arith.index_cast %val : index to i6421 %val_f32 = arith.sitofp %val_i64 : i64 to f3222 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 : index33 %c3 = arith.constant 3 : index34 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 return49}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 return63}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 %unranked_input(%shape)69 : (memref<*xf32>, 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 return78}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 return92}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 %unranked_input(%dyn_size_shape)99 : (memref<*xf32>, 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 return107}108