// RUN: mlir-opt %s --sparse-reinterpret-map --sparsification="sparse-emit-strategy=debug-interface" --canonicalize --cse --allow-unregistered-dialect | FileCheck %s

#map = affine_map<(d0, d1, d2, d3) -> (d0 + d2, d1 + d3)>
#map1 = affine_map<(d0, d1, d2, d3) -> (d2, d3)>
#map2 = affine_map<(d0, d1, d2, d3) -> (d0, d1)>

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


// CHECK-LABEL:   func.func @conv2d_all_sparse_CSR(
// CHECK:           "ne_sub<trivial<compressed[0,0]>>.begin"
// CHECK:           scf.while {{.*}} {
// CHECK:             "ne_sub<trivial<compressed[0,0]>>.not_end"
// CHECK:           } do {
// CHECK:             %[[D0:.*]] = "ne_sub<trivial<compressed[0,0]>>.deref"
// CHECK:             "ne_sub<trivial<compressed[0,1]>>.begin"
// CHECK:             scf.while {{.*}} {
// CHECK:               "ne_sub<trivial<compressed[0,1]>>.not_end"
// CHECK:             } do {
// CHECK:               %[[D1:.*]] = "ne_sub<trivial<compressed[0,1]>>.deref"
// CHECK:               "subsect<trivial<compressed[0,0]>>.begin"
// CHECK:               scf.while {{.*}} {
// CHECK:                 "subsect<trivial<compressed[0,0]>>.not_end
// CHECK:               } do {
// CHECK:                 %[[D2:.*]] = "subsect<trivial<compressed[0,0]>>.deref"
// CHECK:                 "trivial<batch[1,0]>.locate"(%{{.*}}, %[[D2]])
// CHECK:                 "subsect<trivial<compressed[0,1]>>.begin"
// CHECK:                 scf.while {{.*}} {
// CHECK:                   "subsect<trivial<compressed[0,1]>>.not_end"
// CHECK:                 } do {
// CHECK:                   %[[D3:.*]] = "subsect<trivial<compressed[0,1]>>.deref"
// CHECK:                   "trivial<batch[1,1]>.locate"(%{{.*}}, %[[D3]])
// CHECK:                   tensor.extract %{{.*}}{{\[}}%[[D2]], %[[D3]]]
// CHECK:                   arith.muli
// CHECK:                   arith.addi
// CHECK:                   "subsect<trivial<compressed[0,1]>>.next
// CHECK:                   scf.yield
// CHECK:                 }
// CHECK:                 "subsect<trivial<compressed[0,0]>>.next
// CHECK:                 scf.yield
// CHECK:               }
// CHECK:               scf.if {{.*}} {
// CHECK:                 tensor.insert %{{.*}} into %{{.*}}{{\[}}%[[D0]], %[[D1]]]
// CHECK:                 scf.yield
// CHECK:               } else {
// CHECK:                 scf.yield
// CHECK:               }
// CHECK:               "ne_sub<trivial<compressed[0,1]>>.next"
// CHECK:               scf.yield
// CHECK:             }
// CHECK:             "ne_sub<trivial<compressed[0,0]>>.next"
// CHECK:             scf.yield
// CHECK:           }
// CHECK:           sparse_tensor.load
// CHECK:           return
// CHECK:         }
func.func @conv2d_all_sparse_CSR(%arg0: tensor<8x8xi32, #DCSR>,
                                 %arg1: tensor<3x3xi32>) -> tensor<6x6xi32, #DCSR> {
  %0 = tensor.empty() : tensor<6x6xi32, #DCSR>
  %1 = linalg.generic {
         indexing_maps = [#map, #map1, #map2],
         iterator_types = ["parallel", "parallel", "reduction", "reduction"]}
         ins(%arg0, %arg1 : tensor<8x8xi32, #DCSR>, tensor<3x3xi32>)
         outs(%0 : tensor<6x6xi32, #DCSR>) {
    ^bb0(%in: i32, %in_0: i32, %out: i32):
      %2 = arith.muli %in, %in_0 : i32
      %3 = arith.addi %out, %2 : i32
      linalg.yield %3 : i32
    } -> tensor<6x6xi32, #DCSR>
  return %1 : tensor<6x6xi32, #DCSR>
}