rebase

Author: Peiming Liu

mlir/lib/Dialect/SparseTensor/Transforms/SparseSpaceCollapse.cpp

@@ -53,7 +53,7 @@ bool legalToCollapse(SmallVectorImpl<CollapseSpaceInfo> &toCollapse,
                      ExtractIterSpaceOp curSpace) {
 
   auto getIterateOpOverSpace = [](ExtractIterSpaceOp space) -> IterateOp {
-    Value spaceVal = space.getResultSpace();
+    Value spaceVal = space.getExtractedSpace();
     if (spaceVal.hasOneUse())
       return llvm::dyn_cast<IterateOp>(*spaceVal.getUsers().begin());
     return nullptr;
@@ -116,7 +116,7 @@ void collapseSparseSpace(MutableArrayRef<CollapseSpaceInfo> toCollapse) {
   auto innermost = toCollapse.back().loop;
 
   IRMapping mapper;
-  mapper.map(leaf, collapsedSpace.getResultSpace());
+  mapper.map(leaf, collapsedSpace.getExtractedSpace());
   for (auto z : llvm::zip_equal(innermost.getInitArgs(), rItOp.getInitArgs()))
     mapper.map(std::get<0>(z), std::get<1>(z));
 

mlir/test/Dialect/SparseTensor/sparse_space_collapse.mlir

@@ -19,9 +19,13 @@
 // CHECK:           return
 // CHECK:         }
 func.func @sparse_sparse_collapse(%sp : tensor<4x8xf32, #COO>, %i : index) {
-  %l1 = sparse_tensor.extract_iteration_space %sp lvls = 0 : tensor<4x8xf32, #COO>
+  %l1 = sparse_tensor.extract_iteration_space %sp lvls = 0
+      : tensor<4x8xf32, #COO>
+     -> !sparse_tensor.iter_space<#COO, lvls = 0>
   %r1 = sparse_tensor.iterate %it1 in %l1 at(%crd0) iter_args(%outer = %i): !sparse_tensor.iter_space<#COO, lvls = 0 to 1> -> index {
-    %l2 = sparse_tensor.extract_iteration_space %sp at %it1 lvls = 1 : tensor<4x8xf32, #COO>, !sparse_tensor.iterator<#COO, lvls = 0 to 1>
+    %l2 = sparse_tensor.extract_iteration_space %sp at %it1 lvls = 1
+        : tensor<4x8xf32, #COO>, !sparse_tensor.iterator<#COO, lvls = 0 to 1>
+       -> !sparse_tensor.iter_space<#COO, lvls = 1>
     %r2 = sparse_tensor.iterate %it2 in %l2 iter_args(%inner = %outer): !sparse_tensor.iter_space<#COO, lvls = 1 to 2> -> index {
       %k ="test.op"(%inner) : (index) -> index
       sparse_tensor.yield %k : index