[mlir][sparse] fix incorrect coordinates ordering computed by the foreach operation.

[Peiming]

small fixes.

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

Value val = enc ? rewriter.create<memref::LoadOp>(loc, vals, pidx)
                     : rewriter.create<memref::LoadOp>(loc, vals, coords);
 
     // 2. Inline the block in the foreach operator.
     Block *srcBlock = op.getBody();
 
     SmallVector<Value, 4> args;
     // Remap coordinates.
     for (int64_t i = 0; i < rank; i++) {
-      Value actual = coords[toOrigDim(enc, i)];
+      Value actual = coords[toStoredDim(enc, i)];
       args.push_back(actual);
     }
     // Remap value.
     args.push_back(val);
     // Remap reduction variables.
     args.append(reduc);
 
     // Remove sparse_tensor.yield.

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

#SortedCOO = #sparse_tensor.encoding<{
   dimLevelType = [ "compressed-nu", "singleton" ]
 }>
 
 #SortedCOOPerm = #sparse_tensor.encoding<{
   dimLevelType = [ "compressed-nu", "singleton" ],
   dimOrdering = affine_map<(i,j) -> (j,i)>
 }>
 
+#CCCPerm = #sparse_tensor.encoding<{
+  dimLevelType = [ "compressed", "compressed", "compressed"],
+  dimOrdering = affine_map<(d0, d1, d2) -> (d1, d2, d0)>
+}>
+
 module {
   /// uses foreach operator to print coords and values.
   func.func @foreach_print_const() {
     // Initialize a tensor.
     %0 = arith.constant sparse<[[0, 0], [1, 6]], [1.0, 5.0]> : tensor<8x7xf32>
     sparse_tensor.foreach in %0 : tensor<8x7xf32> do {
       ^bb0(%1: index, %2: index, %v: f32) :
         vector.print %1: index
sparse_tensor.foreach in %arg0 : tensor<2x2xf64, #SortedCOOPerm> do {
       ^bb0(%1: index, %2: index, %v: f64) :
         vector.print %1: index
         vector.print %2: index
         vector.print %v: f64
      }
      return
   }
 
+  func.func @foreach_print_3d(%arg0: tensor<7x8x9xf64, #CCCPerm>) {
+    sparse_tensor.foreach in %arg0 : tensor<7x8x9xf64, #CCCPerm> do {
+      ^bb0(%1: index, %2: index, %3: index, %v: f64) :
+        vector.print %1: index
+        vector.print %2: index
+        vector.print %3: index
+        vector.print %v: f64
+     }
+     return
+  }
+
+
   func.func @foreach_print_dense(%arg0: tensor<2x2xf64>) {
     sparse_tensor.foreach in %arg0 : tensor<2x2xf64> do {
     ^bb0(%1: index, %2: index, %v: f64) :
       vector.print %1: index
       vector.print %2: index
       vector.print %v: f64
    }
    return
   }
   
   //
   // Main driver.
   //
   func.func @entry() {
     //
     // Initialize a 3-dim dense tensor.
     //
     %src = arith.constant dense<
        [[  1.0,  2.0],
         [  5.0,  6.0]]
     > : tensor<2x2xf64>
     
+    %src3d = arith.constant sparse<
+       [[1, 2, 3], [4, 5, 6]], [1.0, 2.0] 
+    > : tensor<7x8x9xf64>
+    
     //
     // Convert dense tensor directly to various sparse tensors.
     //
     %s1 = sparse_tensor.convert %src : tensor<2x2xf64> to tensor<2x2xf64, #Row>
     %s2 = sparse_tensor.convert %src : tensor<2x2xf64> to tensor<2x2xf64, #CSR>
     %s3 = sparse_tensor.convert %src : tensor<2x2xf64> to tensor<2x2xf64, #DCSC>
     %s4 = sparse_tensor.convert %src : tensor<2x2xf64> to tensor<2x2xf64, #SortedCOO>
     %s5 = sparse_tensor.convert %src : tensor<2x2xf64> to tensor<2x2xf64, #SortedCOOPerm>
+    %s6 = sparse_tensor.convert %src3d : tensor<7x8x9xf64>  to tensor<7x8x9xf64, #CCCPerm>
     // CHECK: 0
     // CHECK-NEXT: 0
     // CHECK-NEXT: 1
     // CHECK-NEXT: 1
     // CHECK-NEXT: 6
     // CHECK-NEXT: 5
     call @foreach_print_const() : () -> ()
     // CHECK-NEXT: 0
func.func @entry() {
     // CHECK-NEXT: 5
     // CHECK-NEXT: 0
     // CHECK-NEXT: 1
     // CHECK-NEXT: 2
     // CHECK-NEXT: 1
     // CHECK-NEXT: 1
     // CHECK-NEXT: 6
     call @foreach_print_5(%s5) : (tensor<2x2xf64, #SortedCOOPerm>) -> ()
-    
+    // CHECK-NEXT: 1
+    // CHECK-NEXT: 2
+    // CHECK-NEXT: 3
+    // CHECK-NEXT: 1
+    // CHECK-NEXT: 4
+    // CHECK-NEXT: 5
+    // CHECK-NEXT: 6
+    // CHECK-NEXT: 2
+    call @foreach_print_3d(%s6): (tensor<7x8x9xf64, #CCCPerm>) -> ()
+    
     bufferization.dealloc_tensor %s1 : tensor<2x2xf64, #Row>
     bufferization.dealloc_tensor %s2 : tensor<2x2xf64, #CSR>
     bufferization.dealloc_tensor %s3 : tensor<2x2xf64, #DCSC>
     bufferization.dealloc_tensor %s4 : tensor<2x2xf64, #SortedCOO>
     bufferization.dealloc_tensor %s5 : tensor<2x2xf64, #SortedCOOPerm>
 
     return
   }
 }