mlir/test/mlir-cpu-runner/memref-reshape.mlir

   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,reconcile-unrealized-casts)" \
   2 // RUN: | mlir-cpu-runner -e main -entry-point-result=void \
   3 // RUN: -shared-libs=%mlir_runner_utils,%mlir_c_runner_utils \
   4 // RUN: | FileCheck %s
   5
   6
   7 func.func private @printMemrefF32(memref<*xf32>) attributes { llvm.emit_c_interface }
   8
   9 func.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
  51 func.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
  65 func.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
  80 func.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
  94 func.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 }