rebase

Author: Peiming Liu

mlir/test/Dialect/SparseTensor/sparse_iteration_to_scf.mlir

@@ -18,6 +18,6 @@
 // CHECK:           %[[POS_HI:.*]] = memref.load %[[POS_MEM]]{{\[}}%[[C1]]] : memref<?xindex>
 // CHECK:           %[[ITER_SPACE:.*]] = builtin.unrealized_conversion_cast %[[POS_MEM]], %[[CRD_MEM]], %[[LVL_SIZE]], %[[POS_LO]], %[[POS_HI]]
 func.func @sparse_1D_space(%sp : tensor<?x?xf32, #COO>) -> !sparse_tensor.iter_space<#COO, lvls = 0> {
-  %l1 = sparse_tensor.extract_iteration_space %sp lvls = 0 : tensor<?x?xf32, #COO>
+  %l1 = sparse_tensor.extract_iteration_space %sp lvls = 0 : tensor<?x?xf32, #COO> -> !sparse_tensor.iter_space<#COO, lvls = 0>
   return %l1 : !sparse_tensor.iter_space<#COO, lvls = 0>
 }

mlir/test/Dialect/SparseTensor/sparse_space_collapse.mlir

@@ -8,15 +8,15 @@
 }>
 
 // CHECK-LABEL:   func.func @sparse_sparse_collapse(
-// CHECK-SAME:         %[[VAL_0:.*]]: tensor<4x8xf32, #sparse>,
-// CHECK-SAME:         %[[VAL_1:.*]]: index) {
-// CHECK:           %[[VAL_3:.*]] = sparse_tensor.extract_iteration_space %[[VAL_0]] lvls = 0 to 2 : tensor<4x8xf32, #sparse>
-// CHECK:           %[[VAL_4:.*]] = sparse_tensor.iterate %[[VAL_5:.*]] in %[[VAL_3]] at(%[[VAL_6:.*]], _) iter_args(%[[VAL_7:.*]] = %[[VAL_1]])
-// CHECK:             %[[VAL_8:.*]] = "test.op"(%[[VAL_7]]) : (index) -> index
-// CHECK:             sparse_tensor.yield %[[VAL_8]] : index
+// CHECK-SAME:        %[[VAL_0:.*]]: tensor<4x8xf32, #sparse>) -> index {
+// CHECK-DAG:       %[[VAL_1:.*]] = arith.constant 0 : index
+// CHECK-DAG:       %[[VAL_2:.*]] = arith.constant 1 : index
+// CHECK:           %[[VAL_3:.*]] = sparse_tensor.extract_iteration_space %[[VAL_0]] lvls = 0 to 2
+// CHECK:           %[[VAL_4:.*]] = sparse_tensor.iterate %[[VAL_5:.*]] in %[[VAL_3]] iter_args(%[[VAL_6:.*]] = %[[VAL_1]])
+// CHECK:             %[[VAL_7:.*]] = arith.addi %[[VAL_6]], %[[VAL_2]] : index
+// CHECK:             sparse_tensor.yield %[[VAL_7]] : index
 // CHECK:           }
-// CHECK:           "test.sink"(%[[VAL_4]]) : (index) -> ()
-// CHECK:           return
+// CHECK:           return %[[VAL_4]] : index
 // CHECK:         }
 func.func @sparse_sparse_collapse(%sp : tensor<4x8xf32, #COO>) -> index {
   %i = arith.constant 0 : index