mlir/test/IR/memory-ops.mlir

   1 // RUN: mlir-opt %s | FileCheck %s
   2
   3 // CHECK: #[[$MAP:.*]] = affine_map<(d0, d1)[s0] -> (d0 + s0, d1)>
   4
   5 // CHECK-LABEL: func @alloc() {
   6 func.func @alloc() {
   7 ^bb0:
   8   // Test simple alloc.
   9   // CHECK: %{{.*}} = memref.alloc() : memref<1024x64xf32, 1>
  10   %0 = memref.alloc() : memref<1024x64xf32, affine_map<(d0, d1) -> (d0, d1)>, 1>
  11
  12   %c0 = "arith.constant"() {value = 0: index} : () -> index
  13   %c1 = "arith.constant"() {value = 1: index} : () -> index
  14
  15   // Test alloc with dynamic dimensions.
  16   // CHECK: %{{.*}} = memref.alloc(%{{.*}}, %{{.*}}) : memref<?x?xf32, 1>
  17   %1 = memref.alloc(%c0, %c1) : memref<?x?xf32, affine_map<(d0, d1) -> (d0, d1)>, 1>
  18
  19   // Test alloc with no dynamic dimensions and one symbol.
  20   // CHECK: %{{.*}} = memref.alloc()[%{{.*}}] : memref<2x4xf32, #[[$MAP]], 1>
  21   %2 = memref.alloc()[%c0] : memref<2x4xf32, affine_map<(d0, d1)[s0] -> ((d0 + s0), d1)>, 1>
  22
  23   // Test alloc with dynamic dimensions and one symbol.
  24   // CHECK: %{{.*}} = memref.alloc(%{{.*}})[%{{.*}}] : memref<2x?xf32, #[[$MAP]], 1>
  25   %3 = memref.alloc(%c1)[%c0] : memref<2x?xf32, affine_map<(d0, d1)[s0] -> (d0 + s0, d1)>, 1>
  26
  27   // Alloc with no mappings.
  28   // b/116054838 Parser crash while parsing ill-formed AllocOp
  29   // CHECK: %{{.*}} = memref.alloc() : memref<2xi32>
  30   %4 = memref.alloc() : memref<2 x i32>
  31
  32   // CHECK:   return
  33   return
  34 }
  35
  36 // CHECK-LABEL: func @alloca() {
  37 func.func @alloca() {
  38 ^bb0:
  39   // Test simple alloc.
  40   // CHECK: %{{.*}} = memref.alloca() : memref<1024x64xf32, 1>
  41   %0 = memref.alloca() : memref<1024x64xf32, affine_map<(d0, d1) -> (d0, d1)>, 1>
  42
  43   %c0 = "arith.constant"() {value = 0: index} : () -> index
  44   %c1 = "arith.constant"() {value = 1: index} : () -> index
  45
  46   // Test alloca with dynamic dimensions.
  47   // CHECK: %{{.*}} = memref.alloca(%{{.*}}, %{{.*}}) : memref<?x?xf32, 1>
  48   %1 = memref.alloca(%c0, %c1) : memref<?x?xf32, affine_map<(d0, d1) -> (d0, d1)>, 1>
  49
  50   // Test alloca with no dynamic dimensions and one symbol.
  51   // CHECK: %{{.*}} = memref.alloca()[%{{.*}}] : memref<2x4xf32, #[[$MAP]], 1>
  52   %2 = memref.alloca()[%c0] : memref<2x4xf32, affine_map<(d0, d1)[s0] -> ((d0 + s0), d1)>, 1>
  53
  54   // Test alloca with dynamic dimensions and one symbol.
  55   // CHECK: %{{.*}} = memref.alloca(%{{.*}})[%{{.*}}] : memref<2x?xf32, #[[$MAP]], 1>
  56   %3 = memref.alloca(%c1)[%c0] : memref<2x?xf32, affine_map<(d0, d1)[s0] -> (d0 + s0, d1)>, 1>
  57
  58   // Alloca with no mappings, but with alignment.
  59   // CHECK: %{{.*}} = memref.alloca() {alignment = 64 : i64} : memref<2xi32>
  60   %4 = memref.alloca() {alignment = 64} : memref<2 x i32>
  61
  62   return
  63 }
  64
  65 // CHECK-LABEL: func @dealloc() {
  66 func.func @dealloc() {
  67 ^bb0:
  68   // CHECK: %{{.*}} = memref.alloc() : memref<1024x64xf32>
  69   %0 = memref.alloc() : memref<1024x64xf32, affine_map<(d0, d1) -> (d0, d1)>, 0>
  70
  71   // CHECK: memref.dealloc %{{.*}} : memref<1024x64xf32>
  72   memref.dealloc %0 : memref<1024x64xf32, affine_map<(d0, d1) -> (d0, d1)>, 0>
  73   return
  74 }
  75
  76 // CHECK-LABEL: func @load_store
  77 func.func @load_store() {
  78 ^bb0:
  79   // CHECK: %{{.*}} = memref.alloc() : memref<1024x64xf32, 1>
  80   %0 = memref.alloc() : memref<1024x64xf32, affine_map<(d0, d1) -> (d0, d1)>, 1>
  81
  82   %1 = arith.constant 0 : index
  83   %2 = arith.constant 1 : index
  84
  85   // CHECK: %{{.*}} = memref.load %{{.*}}[%{{.*}}, %{{.*}}] : memref<1024x64xf32, 1>
  86   %3 = memref.load %0[%1, %2] : memref<1024x64xf32, affine_map<(d0, d1) -> (d0, d1)>, 1>
  87
  88   // CHECK: memref.store %{{.*}}, %{{.*}}[%{{.*}}, %{{.*}}] : memref<1024x64xf32, 1>
  89   memref.store %3, %0[%1, %2] : memref<1024x64xf32, affine_map<(d0, d1) -> (d0, d1)>, 1>
  90
  91   return
  92 }
  93
  94 // CHECK-LABEL: func @dma_ops()
  95 func.func @dma_ops() {
  96   %c0 = arith.constant 0 : index
  97   %stride = arith.constant 32 : index
  98   %elt_per_stride = arith.constant 16 : index
  99
 100   %A = memref.alloc() : memref<256 x f32, affine_map<(d0) -> (d0)>, 0>
 101   %Ah = memref.alloc() : memref<256 x f32, affine_map<(d0) -> (d0)>, 1>
 102   %tag = memref.alloc() : memref<1 x f32>
 103
 104   %num_elements = arith.constant 256 : index
 105
 106   memref.dma_start %A[%c0], %Ah[%c0], %num_elements, %tag[%c0] : memref<256 x f32>, memref<256 x f32, 1>, memref<1 x f32>
 107   memref.dma_wait %tag[%c0], %num_elements : memref<1 x f32>
 108   // CHECK: dma_start %{{.*}}[%{{.*}}], %{{.*}}[%{{.*}}], %{{.*}}, %{{.*}}[%{{.*}}] : memref<256xf32>, memref<256xf32, 1>, memref<1xf32>
 109   // CHECK-NEXT:  dma_wait %{{.*}}[%{{.*}}], %{{.*}} : memref<1xf32>
 110
 111   // DMA with strides
 112   memref.dma_start %A[%c0], %Ah[%c0], %num_elements, %tag[%c0], %stride, %elt_per_stride : memref<256 x f32>, memref<256 x f32, 1>, memref<1 x f32>
 113   memref.dma_wait %tag[%c0], %num_elements : memref<1 x f32>
 114   // CHECK-NEXT:  dma_start %{{.*}}[%{{.*}}], %{{.*}}[%{{.*}}], %{{.*}}, %{{.*}}[%{{.*}}], %{{.*}}, %{{.*}} : memref<256xf32>, memref<256xf32, 1>, memref<1xf32>
 115   // CHECK-NEXT:  dma_wait %{{.*}}[%{{.*}}], %{{.*}} : memref<1xf32>
 116
 117   return
 118 }