115 lines · plain
1// RUN: mlir-opt %s -pass-pipeline="builtin.module(func.func(test-match-reduction))" -verify-diagnostics -split-input-file2 3// Verify that the generic reduction detection utility works on different4// 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 : f3218 linalg.yield %add : f3219 } -> tensor<1xf32>20 return21}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 : f3228 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 : f3235 affine.yield %add : f3236 }37 affine.store %final_red, %out[%i] : memref<256xf32>38 }39 return40}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 : f3255 %sel = arith.select %cmp, %in0, %out0 : f3256 linalg.yield %sel : f3257 } -> tensor<4xf32>58 return59}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 : f3275 %sub = arith.subf %mul, %in0 : f3276 %add = arith.addf %sub, %out0 : f3277 linalg.yield %add : f3278 } -> tensor<4xf32>79 return80}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 : f3287 // %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 : f3292 affine.yield %ld : f3293 }94 return95}96 97// -----98 99// expected-remark@below {{Testing function}}100func.func @affine_output_dep(%in: memref<512xf32>) {101 %cst = arith.constant 0.000000e+00 : f32102 // Reduction %red is not supported because it depends on another103 // 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 512107 iter_args(%red = %cst, %dep = %cst) -> (f32, f32) {108 %ld = affine.load %in[%j] : memref<512xf32>109 %add = arith.addf %dep, %red : f32110 affine.yield %add, %ld : f32, f32111 }112 return113}114 115