1// RUN: mlir-opt %s --transform-preload-library='transform-library-paths=%p/match_matmul_common.mlir' --transform-interpreter --verify-diagnostics 2 3module attributes { transform.with_named_sequence } { 4 transform.named_sequence @_match_matmul_like( 5 %entry: !transform.any_op {transform.readonly}, 6 %rank: !transform.param<i64> {transform.readonly}) 7 -> (!transform.any_op, !transform.any_op, !transform.param<i64>, 8 !transform.type, !transform.type, !transform.type, 9 !transform.param<i64>, !transform.param<i64>, !transform.param<i64>, !transform.param<i64>) 10 11 transform.named_sequence @match_matmul(%entry: !transform.any_op {transform.readonly}) 12 -> (!transform.any_op, !transform.any_op, !transform.param<i64>, 13 !transform.type, !transform.type, !transform.type) { 14 transform.match.operation_name %entry ["linalg.matmul", "linalg.generic"] : !transform.any_op 15 %c3 = transform.param.constant 3 : i64 -> !transform.param<i64> 16 %fill, %matmul, %dims, %lhs_type, %rhs_type, %res_type, %kinds:4 = 17 transform.include @_match_matmul_like failures(propagate) (%entry, %c3) 18 : (!transform.any_op, !transform.param<i64>) 19 -> (!transform.any_op, !transform.any_op, !transform.param<i64>, 20 !transform.type, !transform.type, !transform.type, 21 !transform.param<i64>, !transform.param<i64>, !transform.param<i64>, !transform.param<i64>) 22 23 transform.yield %fill, %matmul, %dims, %lhs_type, %rhs_type, %res_type 24 : !transform.any_op, !transform.any_op, !transform.param<i64>, !transform.type, !transform.type, !transform.type 25 } 26 27 transform.named_sequence @print_matmul( 28 %fill: !transform.any_op {transform.readonly}, 29 %matmul: !transform.any_op {transform.readonly}, 30 %dims: !transform.param<i64> {transform.readonly}, 31 %lhs_type: !transform.type {transform.readonly}, 32 %rhs_type: !transform.type {transform.readonly}, 33 %res_type: !transform.type {transform.readonly}) { 34 transform.debug.emit_remark_at %fill, "fill" : !transform.any_op 35 transform.debug.emit_remark_at %matmul, "matmul" : !transform.any_op 36 transform.debug.emit_param_as_remark %dims, "dimensions" at %matmul : !transform.param<i64>, !transform.any_op 37 transform.debug.emit_param_as_remark %lhs_type, "LHS type" at %matmul : !transform.type, !transform.any_op 38 transform.debug.emit_param_as_remark %rhs_type, "RHS type" at %matmul : !transform.type, !transform.any_op 39 transform.debug.emit_param_as_remark %res_type, "result type" at %matmul : !transform.type, !transform.any_op 40 transform.yield 41 } 42 43 transform.named_sequence @__transform_main(%root: !transform.any_op {transform.consumed}) { 44 transform.foreach_match in %root 45 @match_matmul -> @print_matmul 46 : (!transform.any_op) -> !transform.any_op 47 transform.yield 48 } 49} 50 51func.func @matmul_simple(%lhs: tensor<10x20xf16>, %rhs: tensor<20x15xf32>) -> tensor<10x15xf64>{ 52 %cst = arith.constant 0.0 : f64 53 %empty = tensor.empty() : tensor<10x15xf64> 54 // expected-remark @below {{fill}} 55 %fill = linalg.fill ins(%cst : f64) outs(%empty : tensor<10x15xf64>) -> tensor<10x15xf64> 56 // expected-remark @below {{matmul}} 57 // expected-remark @below {{dimensions 10 : i64, 15 : i64, 20 : i64}} 58 // expected-remark @below {{LHS type f16}} 59 // expected-remark @below {{RHS type f32}} 60 // expected-remark @below {{result type f64}} 61 %result = linalg.matmul ins(%lhs, %rhs: tensor<10x20xf16>, tensor<20x15xf32>) outs(%fill: tensor<10x15xf64>) -> tensor<10x15xf64> 62 return %result : tensor<10x15xf64> 63} 64 65func.func @matmul_with_extra_ops_in_func(%lhs: tensor<10x20xf32>, %rhs: tensor<20x15xf32>) -> tensor<10x15xf32> { 66 %cst = arith.constant 0.0 : f64 67 %empty = tensor.empty() : tensor<10x15xf32> 68 69 // expected-remark @below {{fill}} 70 %fill = linalg.fill ins(%cst : f64) outs(%empty : tensor<10x15xf32>) -> tensor<10x15xf32> 71 72 %real_lhs = linalg.elemwise_binary { fun = #linalg.binary_fn<mul> } 73 ins(%lhs, %lhs : tensor<10x20xf32>, tensor<10x20xf32>) outs(%lhs : tensor<10x20xf32>) -> tensor<10x20xf32> 74 75 // expected-remark @below {{matmul}} 76 // expected-remark @below {{dimensions 10 : i64, 15 : i64, 20 : i64}} 77 // expected-remark @below {{LHS type f32}} 78 // expected-remark @below {{RHS type f32}} 79 // expected-remark @below {{result type f32}} 80 %result = linalg.matmul ins(%real_lhs, %rhs: tensor<10x20xf32>, tensor<20x15xf32>) outs(%fill: tensor<10x15xf32>) -> tensor<10x15xf32> 81 return %result : tensor<10x15xf32> 82} 83 84func.func @matmul_generic(%lhs: tensor<10x20xf16>, %rhs: tensor<20x15xf32>) -> tensor<10x15xf64>{ 85 %cst = arith.constant 0.0 : f64 86 %empty = tensor.empty() : tensor<10x15xf64> 87 // expected-remark @below {{fill}} 88 %fill = linalg.fill ins(%cst : f64) outs(%empty : tensor<10x15xf64>) -> tensor<10x15xf64> 89 // expected-remark @below {{matmul}} 90 // expected-remark @below {{dimensions 10 : i64, 15 : i64, 20 : i64}} 91 // expected-remark @below {{LHS type f16}} 92 // expected-remark @below {{RHS type f32}} 93 // expected-remark @below {{result type f64}} 94 %result = linalg.generic { 95 indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d2)>, 96 affine_map<(d0, d1, d2) -> (d2, d1)>, 97 affine_map<(d0, d1, d2) -> (d0, d1)>], 98 iterator_types = ["parallel", "parallel", "reduction"] 99 } ins(%lhs, %rhs: tensor<10x20xf16>, tensor<20x15xf32>) outs(%fill: tensor<10x15xf64>) { 100 ^bb(%arg0: f16, %arg1: f32, %arg2: f64): 101 %0 = arith.extf %arg0 : f16 to f32 102 %1 = arith.mulf %0, %arg1 : f32 103 %2 = arith.extf %1 : f32 to f64 104 %3 = arith.addf %2, %arg2 : f64 105 linalg.yield %3 : f64 106 }-> tensor<10x15xf64> 107 return %result : tensor<10x15xf64> 108} 109