test/Dialect/SparseTensor/sparse_fusion.mlir

// RUN: mlir-opt %s --linalg-fuse-elementwise-ops | FileCheck %s

#SV = #sparse_tensor.encoding<{ map = (d0) -> (d0 : compressed) }>

#trait = {
  indexing_maps = [
    affine_map<(i) -> (i)>, // A
    affine_map<(i) -> (i)>  // B (out)
  ],
  iterator_types = ["parallel"],
  doc = "B(i) = OP A(i)"
}

// CHECK-LABEL: func @sparse_fusion
// CHECK:     linalg.generic
// CHECK:       arith.addf
// CHECK:     linalg.generic
// CHECK:       math.exp
// CHECK:       arith.maximumf
// CHECK-NOT: linalg.generic
// CHECK:     return
func.func @sparse_fusion(%argA: tensor<100xf64, #SV>) -> tensor<100xf64> {
  %c1 = arith.constant 1.0 : f64
  %c100 = arith.constant 100.0 : f64

  //
  // Densifying op.
  // Should not be fused with subsequent dense ops.
  //
  %t0 = tensor.empty() : tensor<100xf64>
  %l0 = linalg.generic #trait
      ins(%argA: tensor<100xf64, #SV>) outs(%t0: tensor<100xf64>) {
    ^bb0(%in0: f64, %out0: f64):
      %b0 = arith.addf %in0, %c1 : f64
      linalg.yield %b0 : f64
  } -> tensor<100xf64>


  //
  // Two following dense ops.
  // Should be fused, but not with above.
  //
  %t1 = tensor.empty() : tensor<100xf64>
  %l1 = linalg.generic #trait
      ins(%l0: tensor<100xf64>) outs(%t1: tensor<100xf64>) {
    ^bb0(%in1: f64, %out1: f64):
      %b1 = math.exp %in1 : f64
      linalg.yield %b1 : f64
  } -> tensor<100xf64>
  %t2 = tensor.empty() : tensor<100xf64>
  %l2 = linalg.generic #trait
      ins(%l1: tensor<100xf64>) outs(%t2: tensor<100xf64>) {
    ^bb0(%in2: f64, %out2: f64):
      %b2 = arith.maximumf %in2, %c100 : f64
      linalg.yield %b2 : f64
  } -> tensor<100xf64>

  return %l2 : tensor<100xf64>
}