49 lines · plain
1// RUN: mlir-opt %s --transform-interpreter --verify-diagnostics --split-input-file2 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 %struct8 : (!transform.any_op) -> !transform.any_op9 transform.match.structured.yield %sp_kernel : !transform.any_op10 }11 transform.yield %0 : !transform.any_op12 }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_op16 transform.yield17 }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 %arg022 @match_sparse_structured -> @print_sparse_structured23 : (!transform.any_op) -> !transform.any_op24 transform.yield25 }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 : f6434 %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