[mlir][sparse] fix bugs when computing the memory size when lowering pack op.

[Peiming]

fix check tests.

[Peiming]

fix check tests.

[aartbik]

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

1276

this assert already encapsulates L1281 assert

(but I am okay pedentically checking, but just in case you want to reduce it a bit)

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

foreachFieldAndTypeInSparseTensor(
               mem = rewriter.create<memref::CastOp>(loc, fType, mem);
             }
             fields.push_back(mem);
           }
           return true;
         });
 
     MutSparseTensorDescriptor desc(stt, fields);
+    Value c0 = constantIndex(rewriter, loc, 0);
     Value c1 = constantIndex(rewriter, loc, 1);
     Value c2 = constantIndex(rewriter, loc, 2);
-    Value posBack = c1; // index to the last value in the postion array
-    Value memSize = c2; // memory size for current array
+    Value posBack = c0; // index to the last value in the postion array
+    Value memSize = c1; // memory size for current array
 
     Level trailCOOStart = getCOOStart(stt.getEncoding());
     Level trailCOORank = stt.getLvlRank() - trailCOOStart;
     // Sets up SparseTensorSpecifier.
     for (Level lvl = 0, lvlRank = stt.getLvlRank(); lvl < lvlRank; lvl++) {
       assert(!ShapedType::isDynamic(stt.getDimShape()[lvl]));
 
       // FIXME: dim/lvl confusion!
       // Sets up the level size.
       auto lvlSize = constantIndex(rewriter, loc, stt.getDimShape()[lvl]);
       desc.setLvlSize(rewriter, loc, lvl, lvlSize);
       // We use a single AOS array to store the trailing COO, so there is only
       // one memory size to set for the entire COO section.
       if (lvl > trailCOOStart)
         continue;
 
       // Sets up the memory size by reading the last value in position array.
       DimLevelType dlt = stt.getLvlType(lvl);
       // Simply forwards the position index when this is a dense level.
       if (isDenseDLT(dlt)) {
-        memSize = rewriter.create<arith::MulIOp>(loc, lvlSize, posBack);
+        memSize = rewriter.create<arith::MulIOp>(loc, lvlSize, memSize);
         posBack = rewriter.create<arith::SubIOp>(loc, memSize, c1);
         continue;
       }
 
       if (isDLTWithPos(dlt)) {
         assert(isCompressedDLT(dlt) || isCompressedWithHiDLT(dlt));
         if (isCompressedWithHiDLT(dlt)) {
           memSize = rewriter.create<arith::MulIOp>(loc, memSize, c2);
           posBack = rewriter.create<arith::SubIOp>(loc, memSize, c1);
+        } else {
+          assert(isCompressedDLT(dlt));
+          posBack = memSize;
+          memSize = rewriter.create<arith::AddIOp>(loc, memSize, c1);
         }
         desc.setPosMemSize(rewriter, loc, lvl, memSize);
         // The last value in position array is the memory size for next level.
         memSize = genIndexLoad(rewriter, loc, desc.getPosMemRef(lvl), posBack);
         posBack = rewriter.create<arith::SubIOp>(loc, posBack, c1);
       }
       assert(isDLTWithCrd(dlt) && lvl <= trailCOOStart);
       // FIXME: This seems to be unnecessarily complex, can we simplify it?

mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_pack.mlir

func.func @entry() {
     // CHECK-NEXT:3
     sparse_tensor.foreach in %csr : tensor<2x2xf64, #CSR> do {
       ^bb0(%1: index, %2: index, %v: f64) :
         vector.print %1: index
         vector.print %2: index
         vector.print %v: f64
      }
 
+    %d_csr = tensor.empty() : tensor<3xf64>
+    %p_csr = tensor.empty() : tensor<3xi32>
+    %i_csr = tensor.empty() : tensor<3xi32>
+    %rd_csr, %rp_csr, %ri_csr = sparse_tensor.unpack %csr : tensor<2x2xf64, #CSR>
+                 outs(%d_csr, %p_csr, %i_csr : tensor<3xf64>, tensor<3xi32>, tensor<3xi32>)
+                 -> tensor<3xf64>, tensor<3xi32>, tensor<3xi32>
+
+    // CHECK-NEXT: ( 1, 2, 3 )
+    %vd_csr = vector.transfer_read %rd_csr[%c0], %f0 : tensor<3xf64>, vector<3xf64>
+    vector.print %vd_csr : vector<3xf64>
 
     // CHECK-NEXT:1
     // CHECK-NEXT:2
     // CHECK-NEXT:3
     //
     // CHECK-NEXT:4
     // CHECK-NEXT:5
     //