1e4384149SOleksandr "Alex" Zinenko// RUN: mlir-opt %s --transform-interpreter --verify-diagnostics --split-input-file 2849187c9SPeiming Liu 3849187c9SPeiming Liumodule attributes { transform.with_named_sequence } { 4849187c9SPeiming Liu transform.named_sequence @match_sparse_structured(%arg0: !transform.any_op {transform.readonly}) -> !transform.any_op { 5849187c9SPeiming Liu %0 = transform.match.structured %arg0 : (!transform.any_op) -> !transform.any_op { 6849187c9SPeiming Liu ^bb0(%struct: !transform.any_op): 7849187c9SPeiming Liu %sp_kernel = transform.sparse_tensor.match.sparse_inout %struct 8849187c9SPeiming Liu : (!transform.any_op) -> !transform.any_op 9849187c9SPeiming Liu transform.match.structured.yield %sp_kernel : !transform.any_op 10849187c9SPeiming Liu } 11849187c9SPeiming Liu transform.yield %0 : !transform.any_op 12849187c9SPeiming Liu } 13849187c9SPeiming Liu 14849187c9SPeiming Liu transform.named_sequence @print_sparse_structured(%arg0: !transform.any_op {transform.readonly}) { 15*2798b72aSOleksandr "Alex" Zinenko transform.debug.emit_remark_at %arg0, "sparse_kernel" : !transform.any_op 16849187c9SPeiming Liu transform.yield 17849187c9SPeiming Liu } 18849187c9SPeiming Liu 19849187c9SPeiming Liu // Entry point. Match any structured sparse operation and emit at remark. 20e4384149SOleksandr "Alex" Zinenko transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.consumed}) { 21849187c9SPeiming Liu transform.foreach_match in %arg0 22849187c9SPeiming Liu @match_sparse_structured -> @print_sparse_structured 23849187c9SPeiming Liu : (!transform.any_op) -> !transform.any_op 24e4384149SOleksandr "Alex" Zinenko transform.yield 25849187c9SPeiming Liu } 26849187c9SPeiming Liu} 27849187c9SPeiming Liu 28e2e429d9SYinying Li#CSR = #sparse_tensor.encoding<{map = (d0, d1) -> (d0 : dense, d1 : compressed)}> 29849187c9SPeiming Liu 30849187c9SPeiming Liufunc.func @payload(%lhs: tensor<10x20xf16>, 31849187c9SPeiming Liu %sp_lhs: tensor<10x20xf16, #CSR>, 32849187c9SPeiming Liu %rhs: tensor<20x15xf32>) -> tensor<10x15xf64>{ 33849187c9SPeiming Liu %cst = arith.constant 0.0 : f64 34849187c9SPeiming Liu %empty = tensor.empty() : tensor<10x15xf64> 35849187c9SPeiming Liu %fill = linalg.fill ins(%cst : f64) outs(%empty : tensor<10x15xf64>) -> tensor<10x15xf64> 36849187c9SPeiming Liu 37849187c9SPeiming Liu %result = linalg.matmul ins(%lhs, %rhs: tensor<10x20xf16>, tensor<20x15xf32>) 38849187c9SPeiming Liu outs(%fill: tensor<10x15xf64>) -> tensor<10x15xf64> 39849187c9SPeiming Liu // expected-remark @below {{sparse_kernel}} 40849187c9SPeiming Liu %sp_in = linalg.matmul ins(%sp_lhs, %rhs: tensor<10x20xf16, #CSR>, tensor<20x15xf32>) 41849187c9SPeiming Liu outs(%fill: tensor<10x15xf64>) -> tensor<10x15xf64> 42849187c9SPeiming Liu 43849187c9SPeiming Liu %sp_empty = tensor.empty() : tensor<10x15xf64, #CSR> 44849187c9SPeiming Liu // expected-remark @below {{sparse_kernel}} 45849187c9SPeiming Liu %sp_out = linalg.matmul ins(%lhs, %rhs: tensor<10x20xf16>, tensor<20x15xf32>) 46849187c9SPeiming Liu outs(%sp_empty: tensor<10x15xf64, #CSR>) -> tensor<10x15xf64, #CSR> 47849187c9SPeiming Liu return %result : tensor<10x15xf64> 48849187c9SPeiming Liu} 49