| // RUN: mlir-opt %s --transform-interpreter --verify-diagnostics |
| |
| module attributes { transform.with_named_sequence } { |
| transform.named_sequence @_reduce_leading_trailing(%entry: !transform.any_op {transform.readonly}) |
| -> (!transform.any_op) { |
| %c1 = transform.param.constant 1 : i64 -> !transform.param<i64> |
| |
| transform.match.structured %entry : !transform.any_op { |
| ^bb0(%struct: !transform.any_op): |
| transform.match.structured.dim %struct[all] {parallel} : !transform.any_op |
| transform.match.structured.input %struct[all] {projected_permutation} : !transform.any_op |
| transform.match.structured.init %struct[all] {permutation} : !transform.any_op |
| %ni = transform.match.structured.num_inits %struct : (!transform.any_op) -> !transform.param<i64> |
| transform.match.param.cmpi eq %ni, %c1 : !transform.param<i64> |
| } |
| transform.yield %entry : !transform.any_op |
| } |
| |
| transform.named_sequence @fill_reduce_leading_trailing(%entry: !transform.any_op {transform.readonly}) |
| -> (!transform.any_op, !transform.any_op, !transform.any_op, !transform.any_op, |
| !transform.param<i64>, !transform.param<i64>, !transform.param<i64>) { |
| %c1 = transform.param.constant 1 : i64 -> !transform.param<i64> |
| %c2 = transform.param.constant 2 : i64 -> !transform.param<i64> |
| %c4 = transform.param.constant 4 : i64 -> !transform.param<i64> |
| |
| %rk, %dms, %bw, %operand_o, %init_v, %trailing_o = transform.match.structured failures(propagate) %entry |
| : (!transform.any_op) -> (!transform.param<i64>, !transform.param<i64>, !transform.param<i64>, |
| !transform.any_op, !transform.any_value, !transform.any_op) { |
| ^bb0(%struct: !transform.any_op): |
| %rank = transform.match.structured.rank %struct : (!transform.any_op) -> !transform.param<i64> |
| transform.match.param.cmpi ge %rank, %c2 : !transform.param<i64> |
| transform.match.param.cmpi le %rank, %c4 : !transform.param<i64> |
| |
| transform.match.structured.dim %struct[-1] {reduction} : !transform.any_op |
| transform.match.structured.dim %struct[except(-1)] {parallel} : !transform.any_op |
| %dims = transform.match.structured.dim %struct[all] : (!transform.any_op) -> !transform.param<i64> |
| |
| %n_inputs = transform.match.structured.num_inputs %struct : (!transform.any_op) -> !transform.param<i64> |
| %n_outputs = transform.match.structured.num_inits %struct : (!transform.any_op) -> !transform.param<i64> |
| transform.match.param.cmpi eq %n_inputs, %c1 : !transform.param<i64> |
| transform.match.param.cmpi eq %n_outputs, %c1 : !transform.param<i64> |
| |
| transform.match.structured.input %struct[0] {projected_permutation} : !transform.any_op |
| transform.match.structured.init %struct[0] {projected_permutation} : !transform.any_op |
| %init = transform.match.structured.init %struct[0] : (!transform.any_op) -> !transform.any_value |
| |
| // This danse is necessary to create an empty handle if there is no single |
| // user without failing the entire match |
| %trailing_optional = transform.sequence %struct : (!transform.any_op) -> !transform.any_op failures(suppress) { |
| ^bb0(%struct_inner: !transform.any_op): |
| %result = transform.match.structured failures(propagate) %struct_inner : (!transform.any_op) -> !transform.any_op { |
| ^bb0(%struct_inner_inner: !transform.any_op): |
| %result_inner = transform.match.structured.result %struct_inner_inner[0] {single} : (!transform.any_op) -> !transform.any_op |
| %trailing = transform.include @_reduce_leading_trailing failures(propagate) (%result_inner) : (!transform.any_op) -> !transform.any_op |
| transform.match.structured.yield %trailing : !transform.any_op |
| } |
| transform.yield %result: !transform.any_op |
| } |
| |
| // Suppress errors as a way to implement optionality. We cannot suppress them in |
| // the include because it keeps matching after "get_defining_op" fails, which |
| // breaks the single-op precondition of the following ops. We don't want to |
| // propagate that failure though. |
| // |
| // Additionally, we cannot put the sequence inside the call because its first |
| // operand must be an operation handle (the verifier asserts!) and there is |
| // no such handle available there. |
| // |
| // TODO: extend the structured matching to gracefully handle empty handles |
| // or provide the suppress-errors-but-stop failure mode for includes to |
| // implement optionality. |
| %operand_optional = transform.sequence %struct : (!transform.any_op) -> !transform.any_op failures(suppress) { |
| ^bb0(%struct_inner: !transform.any_op): |
| %operand3 = transform.match.structured failures(propagate) %struct_inner : (!transform.any_op) -> !transform.any_op { |
| ^bb1(%struct_inner_inner: !transform.any_op): |
| %operand = transform.match.structured.input %struct_inner_inner[0] : (!transform.any_op) -> !transform.any_op |
| %operand2 = transform.include @_reduce_leading_trailing failures(propagate) (%operand) : (!transform.any_op) -> !transform.any_op |
| transform.match.structured.yield %operand2 : !transform.any_op |
| } |
| transform.yield %operand3 : !transform.any_op |
| } |
| |
| %bitwidth = transform.match.structured.elemental_bitwidth %init : (!transform.any_value) -> !transform.param<i64> |
| |
| transform.match.structured.body %struct { reduction_position = 0 } : !transform.any_op |
| transform.match.structured.yield %rank, %dims, %bitwidth, %operand_optional, %init, %trailing_optional |
| : !transform.param<i64>, !transform.param<i64>, !transform.param<i64>, |
| !transform.any_op, !transform.any_value, !transform.any_op |
| } |
| |
| %init_o = transform.get_defining_op %init_v : (!transform.any_value) -> !transform.any_op |
| transform.match.operation_name %init_o ["linalg.fill"] : !transform.any_op |
| |
| transform.yield %operand_o, %init_o, %entry, %trailing_o, %rk, %dms, %bw |
| : !transform.any_op, !transform.any_op, !transform.any_op, !transform.any_op, |
| !transform.param<i64>, !transform.param<i64>, !transform.param<i64> |
| } |
| |
| transform.named_sequence @print_reduce_leading_trailing( |
| %leading: !transform.any_op {transform.readonly}, |
| %fill: !transform.any_op {transform.readonly}, |
| %reduction: !transform.any_op {transform.readonly}, |
| %trailing: !transform.any_op {transform.readonly}, |
| %rank: !transform.param<i64> {transform.readonly}, |
| %dims: !transform.param<i64> {transform.readonly}, |
| %bitwidth: !transform.param<i64> {transform.readonly}) { |
| transform.debug.emit_remark_at %leading, "leading" : !transform.any_op |
| transform.debug.emit_remark_at %fill, "fill" : !transform.any_op |
| transform.debug.emit_remark_at %reduction, "reduction" : !transform.any_op |
| transform.debug.emit_remark_at %trailing, "trailing" : !transform.any_op |
| transform.debug.emit_param_as_remark %rank, "rank" at %reduction : !transform.param<i64>, !transform.any_op |
| transform.debug.emit_param_as_remark %dims, "dimensions" at %reduction : !transform.param<i64>, !transform.any_op |
| transform.debug.emit_param_as_remark %bitwidth, "bitwidth" at %reduction : !transform.param<i64>, !transform.any_op |
| transform.yield |
| } |
| |
| transform.named_sequence @__transform_main(%root: !transform.any_op {transform.consumed}) { |
| transform.foreach_match in %root |
| @fill_reduce_leading_trailing -> @print_reduce_leading_trailing |
| : (!transform.any_op) -> !transform.any_op |
| transform.yield |
| } |
| } |
| |
| !in_tensor_t = tensor<8x64xf32> |
| !out_tensor_t = tensor<8xf32> |
| |
| func.func @eltwise_reduce(%arg : !in_tensor_t) -> (!out_tensor_t) { |
| %cst = arith.constant -0.000000e+00 : f32 |
| |
| %0 = tensor.empty() : !out_tensor_t |
| // expected-remark @below {{fill}} |
| %1 = linalg.fill ins(%cst : f32) outs(%0 : !out_tensor_t) -> !out_tensor_t |
| %2 = tensor.empty() : !in_tensor_t |
| // expected-remark @below {{leading}} |
| %3 = linalg.generic { |
| indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, |
| affine_map<(d0, d1) -> (d0, d1)>], |
| iterator_types = ["parallel", "parallel"]} |
| ins(%arg : !in_tensor_t) outs(%2 : !in_tensor_t) { |
| ^bb0(%arg3: f32, %arg4: f32): |
| %4 = arith.addf %arg3, %arg3 : f32 |
| %5 = arith.addf %4, %4 : f32 |
| linalg.yield %5 : f32 |
| } -> !in_tensor_t |
| |
| // expected-remark @below {{reduction}} |
| // expected-remark @below {{rank 2}} |
| // expected-remark @below {{dimensions 8 : i64, 64 : i64}} |
| // expected-remark @below {{bitwidth 32 : i64}} |
| %6 = linalg.generic { |
| indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, |
| affine_map<(d0, d1) -> (d0)>], |
| iterator_types = ["parallel", "reduction"]} |
| ins(%3 : !in_tensor_t) outs(%1 : !out_tensor_t) { |
| ^bb0(%arg3: f32, %arg4: f32): |
| %4 = arith.addf %arg3, %arg4 : f32 |
| linalg.yield %4 : f32 |
| } -> !out_tensor_t |
| |
| return %6 : !out_tensor_t |
| } |
| |
| func.func @reduce_eltwise(%arg : !in_tensor_t) -> (!out_tensor_t) { |
| %cst = arith.constant -0.000000e+00 : f32 |
| |
| %0 = tensor.empty() : !out_tensor_t |
| // expected-remark @below {{fill}} |
| %1 = linalg.fill ins(%cst : f32) outs(%0 : !out_tensor_t) -> !out_tensor_t |
| // expected-remark @below {{reduction}} |
| // expected-remark @below {{rank 2}} |
| // expected-remark @below {{dimensions 8 : i64, 64 : i64}} |
| // expected-remark @below {{bitwidth 32 : i64}} |
| %5 = linalg.generic { |
| indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, |
| affine_map<(d0, d1) -> (d0)>], |
| iterator_types = ["parallel", "reduction"]} |
| ins(%arg : !in_tensor_t) outs(%1 : !out_tensor_t) { |
| ^bb0(%arg3: f32, %arg4: f32): |
| %4 = arith.addf %arg3, %arg4 : f32 |
| linalg.yield %4 : f32 |
| } -> !out_tensor_t |
| |
| %6 = tensor.empty() : !out_tensor_t |
| // expected-remark @below {{trailing}} |
| %7 = linalg.generic { |
| indexing_maps = [affine_map<(d0) -> (d0)>, |
| affine_map<(d0) -> (d0)>], |
| iterator_types = ["parallel"]} |
| ins(%5 : !out_tensor_t) outs(%6 : !out_tensor_t) { |
| ^bb0(%arg3: f32, %arg4: f32): |
| %4 = math.sqrt %arg3 : f32 |
| linalg.yield %4 : f32 |
| } -> !out_tensor_t |
| return %7 : !out_tensor_t |
| } |
| |
| func.func @eltwise_reduce_eltwise(%arg : !in_tensor_t) -> (!out_tensor_t) { |
| %cst = arith.constant -0.000000e+00 : f32 |
| |
| %0 = tensor.empty() : !out_tensor_t |
| // expected-remark @below {{fill}} |
| %1 = linalg.fill ins(%cst : f32) outs(%0 : !out_tensor_t) -> !out_tensor_t |
| %2 = tensor.empty() : !in_tensor_t |
| // expected-remark @below {{leading}} |
| %3 = linalg.generic { |
| indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, |
| affine_map<(d0, d1) -> (d0, d1)>], |
| iterator_types = ["parallel", "parallel"]} |
| ins(%arg : !in_tensor_t) outs(%2 : !in_tensor_t) { |
| ^bb0(%arg3: f32, %arg4: f32): |
| %4 = arith.addf %arg3, %arg3 : f32 |
| %5 = arith.addf %4, %4 : f32 |
| linalg.yield %5 : f32 |
| } -> !in_tensor_t |
| |
| // expected-remark @below {{reduction}} |
| // expected-remark @below {{rank 2}} |
| // expected-remark @below {{dimensions 8 : i64, 64 : i64}} |
| // expected-remark @below {{bitwidth 32 : i64}} |
| %6 = linalg.generic { |
| indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, |
| affine_map<(d0, d1) -> (d0)>], |
| iterator_types = ["parallel", "reduction"]} |
| ins(%3 : !in_tensor_t) outs(%1 : !out_tensor_t) { |
| ^bb0(%arg3: f32, %arg4: f32): |
| %4 = arith.addf %arg3, %arg4 : f32 |
| linalg.yield %4 : f32 |
| } -> !out_tensor_t |
| |
| %7 = tensor.empty() : !out_tensor_t |
| // expected-remark @below {{trailing}} |
| %8 = linalg.generic { |
| indexing_maps = [affine_map<(d0) -> (d0)>, |
| affine_map<(d0) -> (d0)>], |
| iterator_types = ["parallel"]} |
| ins(%6 : !out_tensor_t) outs(%7 : !out_tensor_t) { |
| ^bb0(%arg3: f32, %arg4: f32): |
| %4 = math.sqrt %arg3 : f32 |
| linalg.yield %4 : f32 |
| } -> !out_tensor_t |
| |
| |
| return %8 : !out_tensor_t |
| } |
| |
| func.func @eltwise_reduce_eltwise_swapped(%arg : !in_tensor_t) -> (!out_tensor_t) { |
| %cst = arith.constant -0.000000e+00 : f32 |
| |
| %2 = tensor.empty() : !in_tensor_t |
| // expected-remark @below {{leading}} |
| %3 = linalg.generic { |
| indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, |
| affine_map<(d0, d1) -> (d0, d1)>], |
| iterator_types = ["parallel", "parallel"]} |
| ins(%arg : !in_tensor_t) outs(%2 : !in_tensor_t) { |
| ^bb0(%arg3: f32, %arg4: f32): |
| %4 = arith.addf %arg3, %arg3 : f32 |
| %5 = arith.addf %4, %4 : f32 |
| linalg.yield %5 : f32 |
| } -> !in_tensor_t |
| |
| %0 = tensor.empty() : !out_tensor_t |
| // expected-remark @below {{fill}} |
| %1 = linalg.fill ins(%cst : f32) outs(%0 : !out_tensor_t) -> !out_tensor_t |
| // expected-remark @below {{reduction}} |
| // expected-remark @below {{rank 2}} |
| // expected-remark @below {{dimensions 8 : i64, 64 : i64}} |
| // expected-remark @below {{bitwidth 32 : i64}} |
| %6 = linalg.generic { |
| indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, |
| affine_map<(d0, d1) -> (d0)>], |
| iterator_types = ["parallel", "reduction"]} |
| ins(%3 : !in_tensor_t) outs(%1 : !out_tensor_t) { |
| ^bb0(%arg3: f32, %arg4: f32): |
| %4 = arith.addf %arg3, %arg4 : f32 |
| linalg.yield %4 : f32 |
| } -> !out_tensor_t |
| |
| %7 = tensor.empty() : !out_tensor_t |
| // expected-remark @below {{trailing}} |
| %8 = linalg.generic { |
| indexing_maps = [affine_map<(d0) -> (d0)>, |
| affine_map<(d0) -> (d0)>], |
| iterator_types = ["parallel"]} |
| ins(%6 : !out_tensor_t) outs(%7 : !out_tensor_t) { |
| ^bb0(%arg3: f32, %arg4: f32): |
| %4 = math.sqrt %arg3 : f32 |
| linalg.yield %4 : f32 |
| } -> !out_tensor_t |
| |
| |
| return %8 : !out_tensor_t |
| } |
| |
| func.func @reduction_with_extra_op_in_func(%arg0: tensor<8x479xf32>, %arg1: tensor<32x32xf32>) -> (tensor<8xf32>, tensor<32xf32>) { |
| %cst = arith.constant 0.0 : f32 |
| %empty = tensor.empty() : tensor<8xf32> |
| // expected-remark @below {{fill}} |
| %fill = linalg.fill ins(%cst : f32) outs(%empty : tensor<8xf32>) -> tensor<8xf32> |
| // expected-remark @below {{reduction}} |
| // expected-remark @below {{rank 2}} |
| // expected-remark @below {{dimensions 8 : i64, 479 : i64}} |
| // expected-remark @below {{bitwidth 32 : i64}} |
| %result = linalg.generic { |
| indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, |
| affine_map<(d0, d1) -> (d0)>], |
| iterator_types = ["parallel", "reduction"]} |
| ins(%arg0 : tensor<8x479xf32>) |
| outs(%fill : tensor<8xf32>) { |
| ^bb0(%in: f32, %out: f32): |
| %6 = arith.addf %in, %out : f32 |
| linalg.yield %6 : f32 |
| } -> tensor<8xf32> |
| |
| %empty2 = tensor.empty() : tensor<32xf32> |
| %fill2 = linalg.fill ins(%cst : f32) outs(%empty2 : tensor<32xf32>) -> tensor<32xf32> |
| return %result, %fill2 : tensor<8xf32>, tensor<32xf32> |
| } |
