xref: /llvm-project/mlir/test/Dialect/SCF/forall-to-parallel.mlir (revision 0b665c3dd206abb3deab4c3af9dfeb43000ab0c7)
1// RUN: mlir-opt %s -pass-pipeline='builtin.module(func.func(scf-forall-to-parallel))' -split-input-file | FileCheck %s
2
3func.func private @callee(%i: index, %j: index)
4
5// CHECK-LABEL: @two_iters
6// CHECK-SAME: %[[UB1:.+]]: index, %[[UB2:.+]]: index
7func.func @two_iters(%ub1: index, %ub2: index) {
8  scf.forall (%i, %j) in (%ub1, %ub2) {
9    func.call @callee(%i, %j) : (index, index) -> ()
10  }
11
12  // CHECK: scf.parallel (%[[IV1:.+]], %[[IV2:.+]]) = (%{{.*}}, %{{.*}}) to (%[[UB1]], %[[UB2]])
13  // CHECK:   func.call @callee(%[[IV1]], %[[IV2]]) : (index, index) -> ()
14  // CHECK:   scf.reduce
15  return
16}
17
18// -----
19
20func.func private @callee(%i: index, %j: index)
21
22// CHECK-LABEL: @repeated
23// CHECK-SAME: %[[UB1:.+]]: index, %[[UB2:.+]]: index
24func.func @repeated(%ub1: index, %ub2: index) {
25  scf.forall (%i, %j) in (%ub1, %ub2) {
26    func.call @callee(%i, %j) : (index, index) -> ()
27  }
28
29  // CHECK: scf.parallel (%[[IV1:.+]], %[[IV2:.+]]) = (%{{.*}}, %{{.*}}) to (%[[UB1]], %[[UB2]])
30  // CHECK:   func.call @callee(%[[IV1]], %[[IV2]]) : (index, index) -> ()
31  // CHECK:   scf.reduce
32  scf.forall (%i, %j) in (%ub1, %ub2) {
33    func.call @callee(%i, %j) : (index, index) -> ()
34  }
35
36  // CHECK: scf.parallel (%[[IV3:.+]], %[[IV4:.+]]) = (%{{.*}}, %{{.*}}) to (%[[UB1]], %[[UB2]])
37  // CHECK:   func.call @callee(%[[IV3]], %[[IV4]])
38  // CHECK:   scf.reduce
39  return
40}
41
42// -----
43
44func.func private @callee(%i: index, %j: index, %k: index, %l: index)
45
46// CHECK-LABEL: @nested
47// CHECK-SAME: %[[UB1:.+]]: index, %[[UB2:.+]]: index, %[[UB3:.+]]: index, %[[UB4:.+]]: index
48func.func @nested(%ub1: index, %ub2: index, %ub3: index, %ub4: index) {
49  // CHECK: scf.parallel (%[[IV1:.+]], %[[IV2:.+]]) = (%{{.*}}, %{{.*}}) to (%[[UB1]], %[[UB2]]) step (%{{.*}}, %{{.*}}) {
50  // CHECK:   scf.parallel (%[[IV3:.+]], %[[IV4:.+]]) = (%{{.*}}, %{{.*}}) to (%[[UB3]], %[[UB4]]) step (%{{.*}}, %{{.*}}) {
51  // CHECK:     func.call @callee(%[[IV1]], %[[IV2]], %[[IV3]], %[[IV4]])
52  // CHECK:     scf.reduce
53  // CHECK:   }
54  // CHECK:   scf.reduce
55  // CHECK: }
56  scf.forall (%i, %j) in (%ub1, %ub2) {
57    scf.forall (%k, %l) in (%ub3, %ub4) {
58      func.call @callee(%i, %j, %k, %l) : (index, index, index, index) -> ()
59    }
60  }
61  return
62}
63
64// -----
65
66// CHECK-LABEL: @mapping_attr
67func.func @mapping_attr() -> () {
68  // CHECK: scf.parallel
69  // CHECK:   scf.reduce
70  // CHECK: {mapping = [#gpu.thread<x>]}
71
72  %num_threads = arith.constant 100 : index
73
74  scf.forall (%thread_idx) in (%num_threads) {
75    scf.forall.in_parallel {
76    }
77  } {mapping = [#gpu.thread<x>]}
78  return
79
80}
81