mlir/test/Dialect/Affine/loop-fusion-scf-mixed.mlir

   1 // RUN: mlir-opt -pass-pipeline='builtin.module(func.func(affine-loop-fusion))' %s | FileCheck %s
   2
   3 // Test fusion of affine nests in the presence of other region-holding ops
   4 // (scf.for in the test case below) in the block.
   5
   6 // CHECK-LABEL: func @scf_and_affine
   7 func.func @scf_and_affine(%A : memref<10xf32>) {
   8   %c0 = arith.constant 0 : index
   9   %c1 = arith.constant 1 : index
  10   %c10 = arith.constant 10 : index
  11   %cst = arith.constant 0.0 : f32
  12
  13   %B = memref.alloc() : memref<10xf32>
  14   %C = memref.alloc() : memref<10xf32>
  15
  16   affine.for %j = 0 to 10 {
  17     %v = affine.load %A[%j] : memref<10xf32>
  18     affine.store %v, %B[%j] : memref<10xf32>
  19   }
  20
  21   affine.for %j = 0 to 10 {
  22     %v = affine.load %B[%j] : memref<10xf32>
  23     affine.store %v, %C[%j] : memref<10xf32>
  24   }
  25   // Nests are fused.
  26   // CHECK:     affine.for %{{.*}} = 0 to 10
  27   // CHECK-NOT: affine.for
  28   // CHECK:     scf.for
  29
  30   scf.for %i = %c0 to %c10 step %c1 {
  31     memref.store %cst, %B[%i] : memref<10xf32>
  32   }
  33
  34   // The nests below shouldn't be fused.
  35   affine.for %j = 0 to 10 {
  36     %v = affine.load %A[%j] : memref<10xf32>
  37     affine.store %v, %B[%j] : memref<10xf32>
  38   }
  39   scf.for %i = %c0 to %c10 step %c1 {
  40     memref.store %cst, %B[%i] : memref<10xf32>
  41   }
  42   affine.for %j = 0 to 10 {
  43     %v = affine.load %B[%j] : memref<10xf32>
  44     affine.store %v, %C[%j] : memref<10xf32>
  45   }
  46   // CHECK: affine.for
  47   // CHECK: scf.for
  48   // CHECK: affine.for
  49
  50   return
  51 }