[mlir][sparse] support foreach on dense tensor.

[bixia]

mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorOps.td

837–838

This is not correct anymore

839–845

need to add description about dense tensor

[Peiming]

update doc.

[Peiming]

fix typo.

[wrengr]

mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorOps.td

839–840

Should say instead that this iterates over the stored elements/values (with a parenthetical remark that these are typically the non-zero elements but can also include any stored zeros; à la the new NOE op).

[Peiming]

update doc.

[aartbik]

mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorOps.td

838

Perhaps: Iterates over stored elements in a tensor
as the shorthand?

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

494

val = rewriter.create<memref::LoadOp>(op.getLoc(), vals, enc ? loopEmitter.getLastLevelTensorPointerIndex(0) : coords)

(you can use a local, but move the largely shared part of creating the LoadOp out of the if, assuming the types allow)

[Peiming]

address comments.

[Peiming]

cleanup

mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorOps.td

def SparseTensor_YieldOp : SparseTensor_Op<"yield", [Pure, Terminator]>,
   let assemblyFormat = [{
         $result attr-dict `:` type($result)
   }];
   let hasVerifier = 1;
 }
 
 def SparseTensor_ForeachOp : SparseTensor_Op<"foreach",
     [SingleBlockImplicitTerminator<"YieldOp">]>,
-    Arguments<(ins AnySparseTensor:$tensor)>{
-  let summary = "Iterates over non-zero elements in a sparse tensor";
+    Arguments<(ins AnyTensor:$tensor)>{
+  let summary = "Iterates over elements in a tensor";
   let description = [{
-     Iterates over every non-zero element in the given sparse tensor and executes
-     the block.
+     Iterates over stored elements in a tensor (which are typically, but not always,
+     non-zero for sparse tensors) and executes the block.
 
-     For a input sparse tensor with rank n, the block must take n + 1 arguments. The
+     For an input tensor with rank n, the block must take n + 1 arguments. The
      first n arguments must be Index type, together indicating the current coordinates
      of the element being visited. The last argument must have the same type as the
-     sparse tensor's element type, representing the actual value loaded from the input
+     tensor's element type, representing the actual value loaded from the input
      tensor at the given coordinates.
 
      Example:
 
      ```mlir
      sparse_tensor.foreach in %0 : tensor<?x?xf64, #DCSR> do {
       ^bb0(%arg1: index, %arg2: index, %arg3: f64):
         do something...

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

LogicalResult matchAndRewrite(ForeachOp op,
     auto enc = getSparseTensorEncoding(rtp);
 
     // 1. Generates loop for the sparse input.
     SparseTensorLoopEmitter loopEmitter(ValueRange{input});
     loopEmitter.initializeLoopEmit(rewriter, loc);
     for (int64_t i = 0; i < rank; i++)
       loopEmitter.enterLoopOverTensorAtDim(rewriter, loc, 0, i);
 
-    Value vals = loopEmitter.getTensorValueBuffer(0);
-    Value idx = loopEmitter.getLastLevelTensorPointerIndex(0);
-    Value val = rewriter.create<memref::LoadOp>(op.getLoc(), vals, idx);
-
     SmallVector<Value, 4> coords;
     coords.reserve(rank);
     loopEmitter.getCoordinateArray(coords);
 
+    Value vals = loopEmitter.getTensorValueBuffer(0);
+    Value pidx = loopEmitter.getLastLevelTensorPointerIndex(0);
+    // Loads the value from sparse tensor using pointer index;
+    // loads the value from dense tensor using coordinate array.
+    Value val = enc ? rewriter.create<memref::LoadOp>(loc, vals, pidx)
+                    : rewriter.create<memref::LoadOp>(loc, vals, coords);
+
     for (int64_t i = 0; i < rank; i++)
       loopEmitter.exitCurrentLoop();
 
     // 2. Inline the block in the foreach operator.
     Block::iterator inlinePos = rewriter.getInsertionPoint();
     Block *srcBlock = op.getBody();
     // Remove sparse_tensor.yield.
     rewriter.eraseOp(srcBlock->getTerminator());
 

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

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_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<
func.func @entry() {
     // CHECK-NEXT: 0
     // CHECK-NEXT: 1
     // CHECK-NEXT: 2
     // CHECK-NEXT: 1
     // CHECK-NEXT: 0
     // CHECK-NEXT: 5
     // CHECK-NEXT: 1
     // CHECK-NEXT: 1
-    // CHECK-NEXT: 6
+    // CHECK-NEXT: 6    
+    call @foreach_print_dense(%src) : (tensor<2x2xf64>) -> ()
+    // CHECK-NEXT: 0
+    // CHECK-NEXT: 0
+    // CHECK-NEXT: 1
+    // CHECK-NEXT: 0
+    // CHECK-NEXT: 1
+    // CHECK-NEXT: 2
+    // CHECK-NEXT: 1
+    // CHECK-NEXT: 0
+    // CHECK-NEXT: 5
+    // CHECK-NEXT: 1
+    // CHECK-NEXT: 1
+    // CHECK-NEXT: 6
     call @foreach_print_1(%s1) : (tensor<2x2xf64, #Row>) -> ()
     // CHECK-NEXT: 0
     // CHECK-NEXT: 0
     // CHECK-NEXT: 1
     // CHECK-NEXT: 0
     // CHECK-NEXT: 1
     // CHECK-NEXT: 2
     // CHECK-NEXT: 1