107 lines · plain
1// RUN: mlir-opt %s -pass-pipeline="builtin.module(func.func(convert-scf-to-cf),finalize-memref-to-llvm,func.func(convert-arith-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 6func.func private @printMemrefF32(memref<*xf32>) attributes { llvm.emit_c_interface }7 8func.func @main() -> () {9 %c0 = arith.constant 0 : index10 %c1 = arith.constant 1 : index11 12 // Initialize input.13 %input = memref.alloc() : memref<2x3xf32>14 %dim_x = memref.dim %input, %c0 : memref<2x3xf32>15 %dim_y = memref.dim %input, %c1 : memref<2x3xf32>16 scf.parallel (%i, %j) = (%c0, %c0) to (%dim_x, %dim_y) step (%c1, %c1) {17 %prod = arith.muli %i, %dim_y : index18 %val = arith.addi %prod, %j : index19 %val_i64 = arith.index_cast %val : index to i6420 %val_f32 = arith.sitofp %val_i64 : i64 to f3221 memref.store %val_f32, %input[%i, %j] : memref<2x3xf32>22 }23 %unranked_input = memref.cast %input : memref<2x3xf32> to memref<*xf32>24 call @printMemrefF32(%unranked_input) : (memref<*xf32>) -> ()25 // CHECK: rank = 2 offset = 0 sizes = [2, 3] strides = [3, 1]26 // CHECK-NEXT: [0, 1, 2]27 // CHECK-NEXT: [3, 4, 5]28 29 // Test cases.30 call @cast_ranked_memref_to_static_shape(%input) : (memref<2x3xf32>) -> ()31 call @cast_ranked_memref_to_dynamic_shape(%input) : (memref<2x3xf32>) -> ()32 call @cast_unranked_memref_to_static_shape(%input) : (memref<2x3xf32>) -> ()33 call @cast_unranked_memref_to_dynamic_shape(%input) : (memref<2x3xf32>) -> ()34 memref.dealloc %input : memref<2x3xf32>35 return36}37 38func.func @cast_ranked_memref_to_static_shape(%input : memref<2x3xf32>) {39 %output = memref.reinterpret_cast %input to40 offset: [0], sizes: [6, 1], strides: [1, 1]41 : memref<2x3xf32> to memref<6x1xf32>42 43 %unranked_output = memref.cast %output44 : memref<6x1xf32> to memref<*xf32>45 call @printMemrefF32(%unranked_output) : (memref<*xf32>) -> ()46 // CHECK: rank = 2 offset = 0 sizes = [6, 1] strides = [1, 1] data =47 // CHECK-NEXT: [0],48 // CHECK-NEXT: [1],49 // CHECK-NEXT: [2],50 // CHECK-NEXT: [3],51 // CHECK-NEXT: [4],52 // CHECK-NEXT: [5]53 return54}55 56func.func @cast_ranked_memref_to_dynamic_shape(%input : memref<2x3xf32>) {57 %c0 = arith.constant 0 : index58 %c1 = arith.constant 1 : index59 %c6 = arith.constant 6 : index60 %output = memref.reinterpret_cast %input to61 offset: [%c0], sizes: [%c1, %c6], strides: [%c6, %c1]62 : memref<2x3xf32> to memref<?x?xf32, strided<[?, ?], offset: ?>>63 64 %unranked_output = memref.cast %output65 : memref<?x?xf32, strided<[?, ?], offset: ?>> to memref<*xf32>66 call @printMemrefF32(%unranked_output) : (memref<*xf32>) -> ()67 // CHECK: rank = 2 offset = 0 sizes = [1, 6] strides = [6, 1] data =68 // CHECK-NEXT: [0, 1, 2, 3, 4, 5]69 return70}71 72func.func @cast_unranked_memref_to_static_shape(%input : memref<2x3xf32>) {73 %unranked_input = memref.cast %input : memref<2x3xf32> to memref<*xf32>74 %output = memref.reinterpret_cast %unranked_input to75 offset: [0], sizes: [6, 1], strides: [1, 1]76 : memref<*xf32> to memref<6x1xf32>77 78 %unranked_output = memref.cast %output79 : memref<6x1xf32> to memref<*xf32>80 call @printMemrefF32(%unranked_output) : (memref<*xf32>) -> ()81 // CHECK: rank = 2 offset = 0 sizes = [6, 1] strides = [1, 1] data =82 // CHECK-NEXT: [0],83 // CHECK-NEXT: [1],84 // CHECK-NEXT: [2],85 // CHECK-NEXT: [3],86 // CHECK-NEXT: [4],87 // CHECK-NEXT: [5]88 return89}90 91func.func @cast_unranked_memref_to_dynamic_shape(%input : memref<2x3xf32>) {92 %unranked_input = memref.cast %input : memref<2x3xf32> to memref<*xf32>93 %c0 = arith.constant 0 : index94 %c1 = arith.constant 1 : index95 %c6 = arith.constant 6 : index96 %output = memref.reinterpret_cast %unranked_input to97 offset: [%c0], sizes: [%c1, %c6], strides: [%c6, %c1]98 : memref<*xf32> to memref<?x?xf32, strided<[?, ?], offset: ?>>99 100 %unranked_output = memref.cast %output101 : memref<?x?xf32, strided<[?, ?], offset: ?>> to memref<*xf32>102 call @printMemrefF32(%unranked_output) : (memref<*xf32>) -> ()103 // CHECK: rank = 2 offset = 0 sizes = [1, 6] strides = [6, 1] data =104 // CHECK-NEXT: [0, 1, 2, 3, 4, 5]105 return106}107