| // RUN: mlir-opt %s --transform-interpreter --verify-diagnostics --split-input-file |
| |
| module attributes { transform.with_named_sequence } { |
| transform.named_sequence @match_sparse_structured(%arg0: !transform.any_op {transform.readonly}) -> !transform.any_op { |
| %0 = transform.match.structured %arg0 : (!transform.any_op) -> !transform.any_op { |
| ^bb0(%struct: !transform.any_op): |
| %sp_kernel = transform.sparse_tensor.match.sparse_inout %struct |
| : (!transform.any_op) -> !transform.any_op |
| transform.match.structured.yield %sp_kernel : !transform.any_op |
| } |
| transform.yield %0 : !transform.any_op |
| } |
| |
| transform.named_sequence @print_sparse_structured(%arg0: !transform.any_op {transform.readonly}) { |
| transform.debug.emit_remark_at %arg0, "sparse_kernel" : !transform.any_op |
| transform.yield |
| } |
| |
| // Entry point. Match any structured sparse operation and emit at remark. |
| transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.consumed}) { |
| transform.foreach_match in %arg0 |
| @match_sparse_structured -> @print_sparse_structured |
| : (!transform.any_op) -> !transform.any_op |
| transform.yield |
| } |
| } |
| |
| #CSR = #sparse_tensor.encoding<{map = (d0, d1) -> (d0 : dense, d1 : compressed)}> |
| |
| func.func @payload(%lhs: tensor<10x20xf16>, |
| %sp_lhs: tensor<10x20xf16, #CSR>, |
| %rhs: tensor<20x15xf32>) -> tensor<10x15xf64>{ |
| %cst = arith.constant 0.0 : f64 |
| %empty = tensor.empty() : tensor<10x15xf64> |
| %fill = linalg.fill ins(%cst : f64) outs(%empty : tensor<10x15xf64>) -> tensor<10x15xf64> |
| |
| %result = linalg.matmul ins(%lhs, %rhs: tensor<10x20xf16>, tensor<20x15xf32>) |
| outs(%fill: tensor<10x15xf64>) -> tensor<10x15xf64> |
| // expected-remark @below {{sparse_kernel}} |
| %sp_in = linalg.matmul ins(%sp_lhs, %rhs: tensor<10x20xf16, #CSR>, tensor<20x15xf32>) |
| outs(%fill: tensor<10x15xf64>) -> tensor<10x15xf64> |
| |
| %sp_empty = tensor.empty() : tensor<10x15xf64, #CSR> |
| // expected-remark @below {{sparse_kernel}} |
| %sp_out = linalg.matmul ins(%lhs, %rhs: tensor<10x20xf16>, tensor<20x15xf32>) |
| outs(%sp_empty: tensor<10x15xf64, #CSR>) -> tensor<10x15xf64, #CSR> |
| return %result : tensor<10x15xf64> |
| } |