1// RUN: mlir-opt %s --transform-interpreter --verify-diagnostics --split-input-file 2 3module attributes { transform.with_named_sequence } { 4 transform.named_sequence @match_sparse_structured(%arg0: !transform.any_op {transform.readonly}) -> !transform.any_op { 5 %0 = transform.match.structured %arg0 : (!transform.any_op) -> !transform.any_op { 6 ^bb0(%struct: !transform.any_op): 7 %sp_kernel = transform.sparse_tensor.match.sparse_inout %struct 8 : (!transform.any_op) -> !transform.any_op 9 transform.match.structured.yield %sp_kernel : !transform.any_op 10 } 11 transform.yield %0 : !transform.any_op 12 } 13 14 transform.named_sequence @print_sparse_structured(%arg0: !transform.any_op {transform.readonly}) { 15 transform.debug.emit_remark_at %arg0, "sparse_kernel" : !transform.any_op 16 transform.yield 17 } 18 19 // Entry point. Match any structured sparse operation and emit at remark. 20 transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.consumed}) { 21 transform.foreach_match in %arg0 22 @match_sparse_structured -> @print_sparse_structured 23 : (!transform.any_op) -> !transform.any_op 24 transform.yield 25 } 26} 27 28#CSR = #sparse_tensor.encoding<{map = (d0, d1) -> (d0 : dense, d1 : compressed)}> 29 30func.func @payload(%lhs: tensor<10x20xf16>, 31 %sp_lhs: tensor<10x20xf16, #CSR>, 32 %rhs: tensor<20x15xf32>) -> tensor<10x15xf64>{ 33 %cst = arith.constant 0.0 : f64 34 %empty = tensor.empty() : tensor<10x15xf64> 35 %fill = linalg.fill ins(%cst : f64) outs(%empty : tensor<10x15xf64>) -> tensor<10x15xf64> 36 37 %result = linalg.matmul ins(%lhs, %rhs: tensor<10x20xf16>, tensor<20x15xf32>) 38 outs(%fill: tensor<10x15xf64>) -> tensor<10x15xf64> 39 // expected-remark @below {{sparse_kernel}} 40 %sp_in = linalg.matmul ins(%sp_lhs, %rhs: tensor<10x20xf16, #CSR>, tensor<20x15xf32>) 41 outs(%fill: tensor<10x15xf64>) -> tensor<10x15xf64> 42 43 %sp_empty = tensor.empty() : tensor<10x15xf64, #CSR> 44 // expected-remark @below {{sparse_kernel}} 45 %sp_out = linalg.matmul ins(%lhs, %rhs: tensor<10x20xf16>, tensor<20x15xf32>) 46 outs(%sp_empty: tensor<10x15xf64, #CSR>) -> tensor<10x15xf64, #CSR> 47 return %result : tensor<10x15xf64> 48} 49