1// RUN: mlir-opt %s --sparse-space-collapse | FileCheck %s 2 3#COO = #sparse_tensor.encoding<{ 4 map = (i, j) -> ( 5 i : compressed(nonunique), 6 j : singleton(soa) 7 ) 8}> 9 10// CHECK-LABEL: func.func @sparse_sparse_collapse( 11// CHECK-SAME: %[[VAL_0:.*]]: tensor<4x8xf32, #sparse>) -> index { 12// CHECK-DAG: %[[VAL_1:.*]] = arith.constant 0 : index 13// CHECK-DAG: %[[VAL_2:.*]] = arith.constant 1 : index 14// CHECK: %[[VAL_3:.*]] = sparse_tensor.extract_iteration_space %[[VAL_0]] lvls = 0 to 2 15// CHECK: %[[VAL_4:.*]] = sparse_tensor.iterate %[[VAL_5:.*]] in %[[VAL_3]] iter_args(%[[VAL_6:.*]] = %[[VAL_1]]) 16// CHECK: %[[VAL_7:.*]] = arith.addi %[[VAL_6]], %[[VAL_2]] : index 17// CHECK: sparse_tensor.yield %[[VAL_7]] : index 18// CHECK: } 19// CHECK: return %[[VAL_4]] : index 20// CHECK: } 21func.func @sparse_sparse_collapse(%sp : tensor<4x8xf32, #COO>) -> index { 22 %i = arith.constant 0 : index 23 %c1 = arith.constant 1 : index 24 %l1 = sparse_tensor.extract_iteration_space %sp lvls = 0 25 : tensor<4x8xf32, #COO> -> !sparse_tensor.iter_space<#COO, lvls = 0> 26 %r1 = sparse_tensor.iterate %it1 in %l1 iter_args(%outer = %i): !sparse_tensor.iter_space<#COO, lvls = 0 to 1> -> index { 27 %l2 = sparse_tensor.extract_iteration_space %sp at %it1 lvls = 1 28 : tensor<4x8xf32, #COO>, !sparse_tensor.iterator<#COO, lvls = 0 to 1> -> !sparse_tensor.iter_space<#COO, lvls = 1> 29 %r2 = sparse_tensor.iterate %it2 in %l2 iter_args(%inner = %outer): !sparse_tensor.iter_space<#COO, lvls = 1 to 2> -> index { 30 %k = arith.addi %inner, %c1 : index 31 sparse_tensor.yield %k : index 32 } 33 sparse_tensor.yield %r2 : index 34 } 35 return %r1 : index 36} 37