mlir/test/Dialect/Affine/dma.mlir

   1 // RUN: mlir-opt %s -split-input-file | FileCheck %s
   2
   3 // -----
   4
   5 // Test with loop IVs.
   6 func.func @test0(%arg0 : index, %arg1 : index) {
   7   %0 = memref.alloc() : memref<100x100xf32>
   8   %1 = memref.alloc() : memref<100x100xf32, affine_map<(d0, d1) -> (d0, d1)>, 2>
   9   %2 = memref.alloc() : memref<1xi32>
  10   %c0 = arith.constant 0 : index
  11   %c64 = arith.constant 64 : index
  12   affine.for %i0 = 0 to 10 {
  13     affine.for %i1 = 0 to 10 {
  14       affine.dma_start %0[%i0, %i1], %1[%i0, %i1], %2[%c0], %c64
  15         : memref<100x100xf32>, memref<100x100xf32, 2>, memref<1xi32>
  16       affine.dma_wait %2[%c0], %c64 : memref<1xi32>
  17 // CHECK: affine.dma_start %{{.*}}[%{{.*}}, %{{.*}}], %{{.*}}[%{{.*}}, %{{.*}}], %{{.*}}[%{{.*}}], %{{.*}} : memref<100x100xf32>, memref<100x100xf32, 2>, memref<1xi32>
  18 // CHECK: affine.dma_wait %{{.*}}[%{{.*}}], %{{.*}} : memref<1xi32>
  19     }
  20   }
  21   return
  22 }
  23
  24 // -----
  25
  26 // Test with loop IVs and optional stride arguments.
  27 func.func @test1(%arg0 : index, %arg1 : index) {
  28   %0 = memref.alloc() : memref<100x100xf32>
  29   %1 = memref.alloc() : memref<100x100xf32, affine_map<(d0, d1) -> (d0, d1)>, 2>
  30   %2 = memref.alloc() : memref<1xi32>
  31   %c0 = arith.constant 0 : index
  32   %c64 = arith.constant 64 : index
  33   %c128 = arith.constant 128 : index
  34   %c256 = arith.constant 256 : index
  35   affine.for %i0 = 0 to 10 {
  36     affine.for %i1 = 0 to 10 {
  37       affine.dma_start %0[%i0, %i1], %1[%i0, %i1], %2[%c0], %c64, %c128, %c256
  38         : memref<100x100xf32>, memref<100x100xf32, 2>, memref<1xi32>
  39       affine.dma_wait %2[%c0], %c64 : memref<1xi32>
  40 // CHECK: affine.dma_start %{{.*}}[%{{.*}}, %{{.*}}], %{{.*}}[%{{.*}}, %{{.*}}], %{{.*}}[%{{.*}}], %{{.*}}, %{{.*}}, %{{.*}} : memref<100x100xf32>, memref<100x100xf32, 2>, memref<1xi32>
  41 // CHECK: affine.dma_wait %{{.*}}[%{{.*}}], %{{.*}} : memref<1xi32>
  42     }
  43   }
  44   return
  45 }
  46
  47 // -----
  48
  49 // Test with loop IVs and symbols (without symbol keyword).
  50 func.func @test2(%arg0 : index, %arg1 : index) {
  51   %0 = memref.alloc() : memref<100x100xf32>
  52   %1 = memref.alloc() : memref<100x100xf32, affine_map<(d0, d1) -> (d0, d1)>, 2>
  53   %2 = memref.alloc() : memref<1xi32>
  54   %c0 = arith.constant 0 : index
  55   %c64 = arith.constant 64 : index
  56   affine.for %i0 = 0 to 10 {
  57     affine.for %i1 = 0 to 10 {
  58       affine.dma_start %0[%i0 + %arg0, %i1], %1[%i0, %i1 + %arg1 + 5],
  59                        %2[%c0], %c64
  60         : memref<100x100xf32>, memref<100x100xf32, 2>, memref<1xi32>
  61       affine.dma_wait %2[%c0], %c64 : memref<1xi32>
  62 // CHECK: affine.dma_start %{{.*}}[%{{.*}} + %{{.*}}, %{{.*}}], %{{.*}}[%{{.*}}, %{{.*}} + %{{.*}} + 5], %{{.*}}[%{{.*}}], %{{.*}} : memref<100x100xf32>, memref<100x100xf32, 2>, memref<1xi32>
  63 // CHECK: affine.dma_wait %{{.*}}[%{{.*}}], %{{.*}} : memref<1xi32>
  64     }
  65   }
  66   return
  67 }
  68
  69 // -----
  70
  71 // Test with loop IVs and symbols (with symbol keyword).
  72 func.func @test3(%arg0 : index, %arg1 : index) {
  73   %0 = memref.alloc() : memref<100x100xf32>
  74   %1 = memref.alloc() : memref<100x100xf32, affine_map<(d0, d1) -> (d0, d1)>, 2>
  75   %2 = memref.alloc() : memref<1xi32>
  76   %c0 = arith.constant 0 : index
  77   %c64 = arith.constant 64 : index
  78   affine.for %i0 = 0 to 10 {
  79     affine.for %i1 = 0 to 10 {
  80       affine.dma_start %0[%i0 + symbol(%arg0), %i1],
  81                        %1[%i0, %i1 + symbol(%arg1) + 7],
  82                        %2[%i0 + %i1 + 11], %c64
  83         : memref<100x100xf32>, memref<100x100xf32, 2>, memref<1xi32>
  84       affine.dma_wait %2[%c0], %c64 : memref<1xi32>
  85 // CHECK: affine.dma_start %{{.*}}[%{{.*}} + symbol(%{{.*}}), %{{.*}}], %{{.*}}[%{{.*}}, %{{.*}} + symbol(%{{.*}}) + 7], %{{.*}}[%{{.*}} + %{{.*}} + 11], %{{.*}} : memref<100x100xf32>, memref<100x100xf32, 2>, memref<1xi32>
  86 // CHECK: affine.dma_wait %{{.*}}[%{{.*}}], %{{.*}} : memref<1xi32>
  87     }
  88   }
  89   return
  90 }
  91
  92 // -----
  93
  94 // Test with loop IVs, symbols and constants in nested affine expressions.
  95 func.func @test4(%arg0 : index, %arg1 : index) {
  96   %0 = memref.alloc() : memref<100x100xf32>
  97   %1 = memref.alloc() : memref<100x100xf32, 2>
  98   %2 = memref.alloc() : memref<1xi32>
  99   %c64 = arith.constant 64 : index
 100   affine.for %i0 = 0 to 10 {
 101     affine.for %i1 = 0 to 10 {
 102       affine.dma_start %0[(%i0 + symbol(%arg0)) floordiv 3, %i1],
 103                        %1[%i0, (%i1 + symbol(%arg1)) mod 9 + 7],
 104                        %2[%i0 + %i1 + 11], %c64
 105         : memref<100x100xf32>, memref<100x100xf32, 2>, memref<1xi32>
 106       affine.dma_wait %2[%i0 + %i1 + 11], %c64 : memref<1xi32>
 107 // CHECK: affine.dma_start %{{.*}}[(%{{.*}} + symbol(%{{.*}})) floordiv 3, %{{.*}}], %{{.*}}[%{{.*}}, (%{{.*}} + symbol(%{{.*}})) mod 9 + 7], %{{.*}}[%{{.*}} + %{{.*}} + 11], %{{.*}} : memref<100x100xf32>, memref<100x100xf32, 2>, memref<1xi32>
 108 // CHECK: affine.dma_wait %{{.*}}[%{{.*}} + %{{.*}} + 11], %{{.*}} : memref<1xi32>
 109     }
 110   }
 111   return
 112 }