xref: /llvm-project/mlir/test/Analysis/test-match-reduction.mlir (revision 13bd41096286305ee603428f6adf161f52981827)
1// RUN: mlir-opt %s -pass-pipeline="builtin.module(func.func(test-match-reduction))" -verify-diagnostics -split-input-file
2
3// Verify that the generic reduction detection utility works on different
4// dialects.
5
6// expected-remark@below {{Testing function}}
7func.func @linalg_red_add(%in0t : tensor<?xf32>, %out0t : tensor<1xf32>) {
8  // expected-remark@below {{Reduction found in output #0!}}
9  // expected-remark@below {{Reduced Value: <block argument> of type 'f32' at index: 0}}
10  // expected-remark@below {{Combiner Op: %1 = arith.addf }}
11  %red = linalg.generic {indexing_maps = [affine_map<(d0) -> (d0)>,
12                                          affine_map<(d0) -> (0)>],
13                                          iterator_types = ["reduction"]}
14   ins(%in0t : tensor<?xf32>)
15   outs(%out0t : tensor<1xf32>) {
16    ^bb0(%in0: f32, %out0: f32):
17      %add = arith.addf %in0, %out0 : f32
18      linalg.yield %add : f32
19    } -> tensor<1xf32>
20  return
21}
22
23// -----
24
25// expected-remark@below {{Testing function}}
26func.func @affine_red_add(%in: memref<256x512xf32>, %out: memref<256xf32>) {
27 %cst = arith.constant 0.000000e+00 : f32
28 affine.for %i = 0 to 256 {
29   // expected-remark@below {{Reduction found in output #0!}}
30   // expected-remark@below {{Reduced Value: %1 = affine.load }}
31   // expected-remark@below {{Combiner Op: %2 = arith.addf }}
32   %final_red = affine.for %j = 0 to 512 iter_args(%red_iter = %cst) -> (f32) {
33     %ld = affine.load %in[%i, %j] : memref<256x512xf32>
34     %add = arith.addf %red_iter, %ld : f32
35     affine.yield %add : f32
36   }
37   affine.store %final_red, %out[%i] : memref<256xf32>
38 }
39 return
40}
41
42// -----
43
44// TODO: Iteration-carried values with multiple uses are not supported yet.
45// expected-remark@below {{Testing function}}
46func.func @linalg_red_max(%in0t: tensor<4x4xf32>, %out0t: tensor<4xf32>) {
47  // expected-remark@below {{Reduction NOT found in output #0!}}
48  %red = linalg.generic {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>,
49                                          affine_map<(d0, d1) -> (d0)>],
50   iterator_types = ["parallel", "reduction"]}
51   ins(%in0t : tensor<4x4xf32>)
52   outs(%out0t : tensor<4xf32>) {
53    ^bb0(%in0: f32, %out0: f32):
54      %cmp = arith.cmpf ogt, %in0, %out0 : f32
55      %sel = arith.select %cmp, %in0, %out0 : f32
56      linalg.yield %sel : f32
57    } -> tensor<4xf32>
58  return
59}
60
61// -----
62
63// expected-remark@below {{Testing function}}
64func.func @linalg_fused_red_add(%in0t: tensor<4x4xf32>, %out0t: tensor<4xf32>) {
65  // expected-remark@below {{Reduction found in output #0!}}
66  // expected-remark@below {{Reduced Value: %2 = arith.subf}}
67  // expected-remark@below {{Combiner Op: %3 = arith.addf}}
68  %red = linalg.generic {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>,
69                                          affine_map<(d0, d1) -> (d0)>],
70   iterator_types = ["parallel", "reduction"]}
71   ins(%in0t : tensor<4x4xf32>)
72   outs(%out0t : tensor<4xf32>) {
73    ^bb0(%in0: f32, %out0: f32):
74      %mul = arith.mulf %in0, %in0 : f32
75      %sub = arith.subf %mul, %in0 : f32
76      %add = arith.addf %sub, %out0 : f32
77      linalg.yield %add : f32
78    } -> tensor<4xf32>
79  return
80}
81
82// -----
83
84// expected-remark@below {{Testing function}}
85func.func @affine_no_red_rec(%in: memref<512xf32>) {
86 %cst = arith.constant 0.000000e+00 : f32
87 // %rec is the value loaded in the previous iteration.
88 // expected-remark@below {{Reduction NOT found in output #0!}}
89 %final_val = affine.for %j = 0 to 512 iter_args(%rec = %cst) -> (f32) {
90   %ld = affine.load %in[%j] : memref<512xf32>
91   %add = arith.addf %ld, %rec : f32
92   affine.yield %ld : f32
93 }
94 return
95}
96
97// -----
98
99// expected-remark@below {{Testing function}}
100func.func @affine_output_dep(%in: memref<512xf32>) {
101 %cst = arith.constant 0.000000e+00 : f32
102 // Reduction %red is not supported because it depends on another
103 // loop-carried dependence.
104 // expected-remark@below {{Reduction NOT found in output #0!}}
105 // expected-remark@below {{Reduction NOT found in output #1!}}
106 %final_red, %final_dep = affine.for %j = 0 to 512
107  iter_args(%red = %cst, %dep = %cst) -> (f32, f32) {
108   %ld = affine.load %in[%j] : memref<512xf32>
109   %add = arith.addf %dep, %red : f32
110   affine.yield %add, %ld : f32, f32
111 }
112 return
113}
114
115