[mlir][sparse] Let the SparseTensorReader handle the isSymmetric flag.

[bixia]

Fix FileCheck test.

[Peiming]

mlir/include/mlir/ExecutionEngine/SparseTensor/File.h

100

This looks scary, can we use a unique_ptr for this?

[wrengr]

mlir/include/mlir/ExecutionEngine/SparseTensor/File.h

98–100

We don't use the trailing underscore style anywhere else, so I don't think we should introduce it here.

[wrengr]

mlir/include/mlir/ExecutionEngine/SparseTensor/File.h

183–187

What's the point of having both hasPendingSymmetricValue and also pendingSymmetricValue? Given that the latter is a pointer, you can always check if pendingSymmetricValue is nullptr rather than storing an additional bool

[wrengr]

mlir/include/mlir/ExecutionEngine/SparseTensor/File.h

191–192

I'd rather avoid having this extra counter in the object. For the openSparseTensor function we already track this in the for-loop, so why doesn't the codegen-pass do the same?

[wrengr]

mlir/include/mlir/ExecutionEngine/SparseTensor/File.h

183–189

Akin to my previous comment, why does this need to be handled in the runtime library at all? Why doesn't the codegen-pass just allocate room for the pendingSymmetricValue and pendingSymmetricIndices and handle all this logic on the codegen side?

For my own work I'm trying to move more of the file processing stuff out of the library, so it seems like a step backward for the codegen-pass to want to add more stuff to the library

[bixia]

mlir/include/mlir/ExecutionEngine/SparseTensor/File.h

100

Need to discuss this, I actually have a bug here, as deleting void pointer has undefined behavior.
We may need to add type to SparseTensorReader, such as SparseTensorReader<V>, to resolve this.

183–187

pendingSymmetricValue is a buffer for storing a pending value that we encountered, and the same buffer is shared during the life time of the reader. That is, we always keep such a buffer, and the buffer only has a valid value right after reading an element and not valid right after we use the value.

183–189

Sorry for not making the context of this work clear.
The approach you talked about is https://reviews.llvm.org/D138214
We are trying to comparing these two approaches, and believe handling the isSymmetric flag as in this PR is better for performance because the loop in https://reviews.llvm.org/D138214 has two insert operators, which complicates the loop due to the complexity of the insert operator.

191–192

In openSparseTensor, we use a loop with two add-element function calls to handle symmetric values.
In codegen, we need a loop with two insert ops to do the same. We are trying to avoid two insert ops as explain above.

[aartbik]

Are we keeping codegen solution (controlled by expand_symmetry flag proposed in other pending revision)?

mlir/include/mlir/ExecutionEngine/SparseTensor/File.h

enum class ValueKind : uint8_t {
     kPattern = 1,
     kReal = 2,
     kInteger = 3,
     kComplex = 4,
     // The value set by `readExtFROSTTHeader`.
     kUndefined = 5
   };
 
-  explicit SparseTensorReader(const char *filename) : filename(filename) {
+  explicit SparseTensorReader(const char *filename)
+      : filename(filename), hasPendingSymmetricValue_(false),
+        pendingSymmetricValue_(nullptr), nnzCount_(0) {
     assert(filename && "Received nullptr for filename");
   }
 
   // Disallows copying, to avoid duplicating the `file` pointer.
   SparseTensorReader(const SparseTensorReader &) = delete;
   SparseTensorReader &operator=(const SparseTensorReader &) = delete;
 
   // This dtor tries to avoid leaking the `file`.  (Though it's better
   // to call `closeFile` explicitly when possible, since there are
   // circumstances where dtors are not called reliably.)
-  ~SparseTensorReader() { closeFile(); }
+  ~SparseTensorReader() {
+    closeFile();
+    if (pendingSymmetricValue_)
+      delete pendingSymmetricValue_;
+  }
 
   /// Opens the file for reading.
   void openFile();
 
   /// Closes the file.
   void closeFile();
 
   /// Reads and parses the file's header.
uint64_t getDimSize(uint64_t d) const {
     assert(d < getRank() && "Dimension out of bounds");
     return idata[2 + d];
   }
 
   /// Asserts the shape subsumes the actual dimension sizes.  Is only
   /// valid after parsing the header.
   void assertMatchesShape(uint64_t rank, const uint64_t *shape) const;
 
+  /// A wrapper of readCOOElement with the ability to buffer the value in the
+  /// symmetric structure for the next call of the routine. Returns false when
+  /// there is no more value to return.
+  template <typename V>
+  bool readCOOElementHandleSymmetric(uint64_t rank, uint64_t *indices,
+                                     V *value) {
+    if (hasPendingSymmetricValue_) {
+      hasPendingSymmetricValue_ = false;
+      indices[0] = pendingSymmetricIndices_[0];
+      indices[1] = pendingSymmetricIndices_[1];
+      *value = *static_cast<V *>(pendingSymmetricValue_);
+      return true;
+    }
+
+    if (nnzCount_++ == getNNZ())
+      return false;
+
+    *value = readCOOElement<V>(rank, indices);
+    if (isSymmetric() && indices[0] != indices[1]) {
+      hasPendingSymmetricValue_ = true;
+      if (!pendingSymmetricValue_)
+        pendingSymmetricValue_ = static_cast<void *>(new V);
+      *static_cast<V *>(pendingSymmetricValue_) = *value;
+      pendingSymmetricIndices_[0] = indices[1];
+      pendingSymmetricIndices_[1] = indices[0];
+    }
+    return true;
+  }
+
   /// Reads a sparse tensor element from the next line in the input file and
   /// returns the value of the element. Stores the coordinates of the element
   /// to the `indices` array.
   template <typename V>
   V readCOOElement(uint64_t rank, uint64_t *indices) {
     char *linePtr = readCOOIndices(rank, indices);
     return detail::readCOOValue<V>(&linePtr, isPattern());
   }
private:
 
   static constexpr int kColWidth = 1025;
   const char *const filename;
   FILE *file = nullptr;
   ValueKind valueKind_ = ValueKind::kInvalid;
   bool isSymmetric_ = false;
   uint64_t idata[512];
   char line[kColWidth];
+  bool hasPendingSymmetricValue_;
+  void *pendingSymmetricValue_;
+  uint64_t pendingSymmetricIndices_[2];
+  uint64_t nnzCount_;
 };
 
 //===----------------------------------------------------------------------===//
 
 /// Reads a sparse tensor with the given filename into a memory-resident
 /// sparse tensor.
 ///
 /// Preconditions:

mlir/include/mlir/ExecutionEngine/SparseTensorRuntime.h

MLIR_CRUNNERUTILS_EXPORT void _mlir_ciface_getSparseTensorReaderDimSizes(
     void *p, StridedMemRefType<index_type, 1> *dref);
 
 /// Releases the SparseTensorReader. This also closes the file associated with
 /// the reader.
 MLIR_CRUNNERUTILS_EXPORT void delSparseTensorReader(void *p);
 
 /// Returns the next element for the sparse tensor being read.
 #define DECL_GETNEXT(VNAME, V)                                                 \
-  MLIR_CRUNNERUTILS_EXPORT void _mlir_ciface_getSparseTensorReaderNext##VNAME( \
+  MLIR_CRUNNERUTILS_EXPORT bool _mlir_ciface_getSparseTensorReaderNext##VNAME( \
       void *p, StridedMemRefType<index_type, 1> *iref,                         \
       StridedMemRefType<V, 0> *vref);
 MLIR_SPARSETENSOR_FOREVERY_V(DECL_GETNEXT)
 #undef DECL_GETNEXT
 
 using SparseTensorWriter = std::ostream;
 
 /// Creates a SparseTensorWriter for outputing a sparse tensor to a file with

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

if (!dstTp.hasStaticShape()) {
               loc, dimSizes, constantIndex(rewriter, loc, d.index())));
         }
       }
     }
 
     // Implement the NewOp as follows:
     //   %tmp = bufferization.alloc_tensor : an unordered COO with identity
     //                                       storage ordering
-    //   for i = 0 to nnz
-    //     get the next element from the input file
+    //   while getSparseTensorReaderNext
     //     insert the element to %tmp
     //   %t = sparse_tensor.ConvertOp %tmp
     RankedTensorType cooTp = getUnorderedCOOFromType(dstTp);
     auto cooBuffer =
         rewriter.create<AllocTensorOp>(loc, cooTp, dynSizesArray).getResult();
 
-    Value c0 = constantIndex(rewriter, loc, 0);
-    Value c1 = constantIndex(rewriter, loc, 1);
-    Value nnz = createFuncCall(rewriter, loc, "getSparseTensorReaderNNZ",
-                               {indexTp}, {reader}, EmitCInterface::Off)
-                    .getResult(0);
-    Value symmetric;
-    // We assume only rank 2 tensors may have the isSymmetric flag set.
-    if (rank == 2) {
-      symmetric =
-          createFuncCall(rewriter, loc, "getSparseTensorReaderIsSymmetric",
-                         {rewriter.getI1Type()}, {reader}, EmitCInterface::Off)
-              .getResult(0);
-    } else {
-      symmetric = Value();
-    }
     Type eltTp = dstTp.getElementType();
     Value value = genAllocaScalar(rewriter, loc, eltTp);
-    scf::ForOp forOp = rewriter.create<scf::ForOp>(loc, c0, nnz, c1,
-                                                   ArrayRef<Value>(cooBuffer));
-    rewriter.setInsertionPointToStart(forOp.getBody());
-
+    SmallVector<Type> types{cooBuffer.getType()};
+    scf::WhileOp whileOp =
+        rewriter.create<scf::WhileOp>(loc, types, ValueRange{cooBuffer});
+
+    // The before-region of the WhileOp.
+    Block *before =
+        rewriter.createBlock(&whileOp.getBefore(), {}, types, {loc});
+    rewriter.setInsertionPointToEnd(before);
     SmallString<18> getNextFuncName{"getSparseTensorReaderNext",
                                     primaryTypeFunctionSuffix(eltTp)};
     Value indices = dimSizes; // Reuse the indices memref to store indices.
-    createFuncCall(rewriter, loc, getNextFuncName, {}, {reader, indices, value},
-                   EmitCInterface::On);
+    Value result =
+        createFuncCall(rewriter, loc, getNextFuncName, {rewriter.getI1Type()},
+                       {reader, indices, value}, EmitCInterface::On)
+            .getResult(0);
+    rewriter.create<scf::ConditionOp>(loc, result, before->getArguments());
+
+    // The after-region of the WhileOp.
+    Block *after = rewriter.createBlock(&whileOp.getAfter(), {}, types, {loc});
+    rewriter.setInsertionPointToEnd(after);
+
     SmallVector<Value> indicesArray;
     for (uint64_t i = 0; i < rank; i++) {
       indicesArray.push_back(rewriter.create<memref::LoadOp>(
           loc, indices, constantIndex(rewriter, loc, i)));
     }
     Value v = rewriter.create<memref::LoadOp>(loc, value);
-    Value t = rewriter.create<InsertOp>(loc, v, forOp.getRegionIterArg(0),
-                                        indicesArray);
-    if (symmetric) {
-      Value eq = rewriter.create<arith::CmpIOp>(
-          loc, arith::CmpIPredicate::ne, indicesArray[0], indicesArray[1]);
-      Value cond = rewriter.create<arith::AndIOp>(loc, symmetric, eq);
-      scf::IfOp ifOp =
-          rewriter.create<scf::IfOp>(loc, t.getType(), cond, /*else*/ true);
-      rewriter.setInsertionPointToStart(&ifOp.getThenRegion().front());
-      rewriter.create<scf::YieldOp>(
-          loc, Value(rewriter.create<InsertOp>(
-                   loc, v, t, ValueRange{indicesArray[1], indicesArray[0]})));
-      rewriter.setInsertionPointToStart(&ifOp.getElseRegion().front());
-      rewriter.create<scf::YieldOp>(loc, t);
-      t = ifOp.getResult(0);
-      rewriter.setInsertionPointAfter(ifOp);
-    }
+    Value t =
+        rewriter.create<InsertOp>(loc, v, after->getArgument(0), indicesArray);
+
     rewriter.create<scf::YieldOp>(loc, ArrayRef<Value>(t));
-    rewriter.setInsertionPointAfter(forOp);
+    rewriter.setInsertionPointAfter(whileOp);
     // Link SSA chain.
-    cooBuffer = forOp.getResult(0);
+    cooBuffer = whileOp.getResult(0);
 
     // Release the sparse tensor reader.
     createFuncCall(rewriter, loc, "delSparseTensorReader", {}, {reader},
                    EmitCInterface::Off);
     cooBuffer = rewriter.create<LoadOp>(loc, cooBuffer, true);
     Value newOp = rewriter.replaceOpWithNewOp<ConvertOp>(op, dstTp, cooBuffer);
 
     // Release the unordered COO tensor buffer.

mlir/lib/ExecutionEngine/SparseTensorRuntime.cpp

void _mlir_ciface_getSparseTensorReaderDimSizes(
   SparseTensorReader &file = *static_cast<SparseTensorReader *>(p);
   const index_type *sizes = file.getDimSizes();
   index_type rank = file.getRank();
   for (index_type r = 0; r < rank; ++r)
     dimSizes[r] = sizes[r];
 }
 
 #define IMPL_GETNEXT(VNAME, V)                                                 \
-  void _mlir_ciface_getSparseTensorReaderNext##VNAME(                          \
+  bool _mlir_ciface_getSparseTensorReaderNext##VNAME(                          \
       void *p, StridedMemRefType<index_type, 1> *iref,                         \
       StridedMemRefType<V, 0> *vref) {                                         \
     assert(p &&vref);                                                          \
     ASSERT_NO_STRIDE(iref);                                                    \
     index_type *indices = MEMREF_GET_PAYLOAD(iref);                            \
     SparseTensorReader *stfile = static_cast<SparseTensorReader *>(p);         \
     index_type rank = stfile->getRank();                                       \
     V *value = MEMREF_GET_PAYLOAD(vref);                                       \
-    *value = stfile->readCOOElement<V>(rank, indices);                         \
+    return stfile->readCOOElementHandleSymmetric<V>(rank, indices, value);     \
   }
 MLIR_SPARSETENSOR_FOREVERY_V(IMPL_GETNEXT)
 #undef IMPL_GETNEXT
 
 void _mlir_ciface_outSparseTensorWriterMetaData(
     void *p, index_type rank, index_type nnz,
     StridedMemRefType<index_type, 1> *dref) {
   assert(p);

mlir/test/Dialect/SparseTensor/rewriting_for_codegen.mlir

 // CHECK-DAG:     %[[C0:.*]] = arith.constant 0 : index
 // CHECK-DAG:     %[[C1:.*]] = arith.constant 1 : index
 // CHECK:         %[[R:.*]] = call @createSparseTensorReader(%[[A]])
 // CHECK:         %[[DS:.*]] = memref.alloca(%[[C2]]) : memref<?xindex>
 // CHECK:         call @getSparseTensorReaderDimSizes(%[[R]], %[[DS]])
 // CHECK:         %[[D0:.*]] = memref.load %[[DS]]{{\[}}%[[C0]]]
 // CHECK:         %[[D1:.*]] = memref.load %[[DS]]{{\[}}%[[C1]]]
 // CHECK:         %[[T:.*]] = bufferization.alloc_tensor(%[[D0]], %[[D1]])
-// CHECK:         %[[N:.*]] = call @getSparseTensorReaderNNZ(%[[R]])
-// CHECK:         %[[S:.*]] = call @getSparseTensorReaderIsSymmetric(%[[R]])
 // CHECK:         %[[VB:.*]] = memref.alloca()
-// CHECK:         %[[T2:.*]] = scf.for %{{.*}} = %[[C0]] to %[[N]] step %[[C1]] iter_args(%[[A2:.*]] = %[[T]])
-// CHECK:           func.call @getSparseTensorReaderNextF32(%[[R]], %[[DS]], %[[VB]])
+// CHECK:         %[[T4:.*]] = scf.while (%[[T1:.*]] = %[[T]]) 
+// CHECK:           %[[R1:.*]] = func.call @getSparseTensorReaderNextF32(%[[R]], %[[DS]], %[[VB]])
+// CHECK:           scf.condition(%[[R1]]) %[[T1]]
+// CHECK:         } do {
+// CHECK:        ^bb0(%[[T2:.*]]: tensor
 // CHECK:           %[[E0:.*]] = memref.load %[[DS]]{{\[}}%[[C0]]]
 // CHECK:           %[[E1:.*]] = memref.load %[[DS]]{{\[}}%[[C1]]]
 // CHECK:           %[[V:.*]] = memref.load %[[VB]][]
-// CHECK:           %[[T1:.*]] = sparse_tensor.insert %[[V]] into %[[A2]]{{\[}}%[[E0]], %[[E1]]]
-// CHECK:           %[[NE:.*]] = arith.cmpi ne, %[[E0]], %[[E1]]
-// CHECK:           %[[COND:.*]] = arith.andi %[[S]], %[[NE]]
-// CHECK:           %[[T3:.*]] = scf.if %[[COND]]
-// CHECK:             %[[T4:.*]] = sparse_tensor.insert %[[V]] into %[[T1]]{{\[}}%[[E1]], %[[E0]]]
-// CHECK:             scf.yield %[[T4]]
-// CHECK:           else
-// CHECK:             scf.yield %[[T1]]
+// CHECK:           %[[T3:.*]] = sparse_tensor.insert %[[V]] into %[[T2]]{{\[}}%[[E0]], %[[E1]]]
 // CHECK:           scf.yield %[[T3]]
 // CHECK:         }
 // CHECK:         call @delSparseTensorReader(%[[R]])
-// CHECK:         %[[T5:.*]] = sparse_tensor.load %[[T2]] hasInserts
-// CHECK:         %[[R:.*]] = sparse_tensor.convert %[[T5]]
+// CHECK:         %[[T5:.*]] = sparse_tensor.load %[[T4]] hasInserts
+// CHECK:         %[[T6:.*]] = sparse_tensor.convert %[[T5]]
 // CHECK:         bufferization.dealloc_tensor %[[T5]]
-// CHECK:         return %[[R]]
+// CHECK:         return %[[T6]]
 func.func @sparse_new(%arg0: !llvm.ptr<i8>) -> tensor<?x?xf32, #CSR> {
   %0 = sparse_tensor.new %arg0 : !llvm.ptr<i8> to tensor<?x?xf32, #CSR>
   return %0 : tensor<?x?xf32, #CSR>
 }
 
 // CHECK-LABEL:   func.func @sparse_out(
 // CHECK-SAME:    %[[A:.*]]: tensor<10x20xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "compressed" ] }>>,
 // CHECK-SAME:    %[[B:.*]]: !llvm.ptr<i8>) {

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

module {
   func.func private @createSparseTensorReader(!Filename) -> (!TensorReader)
   func.func private @delSparseTensorReader(!TensorReader) -> ()
   func.func private @getSparseTensorReaderRank(!TensorReader) -> (index)
   func.func private @getSparseTensorReaderNNZ(!TensorReader) -> (index)
   func.func private @getSparseTensorReaderIsSymmetric(!TensorReader) -> (i1)
   func.func private @getSparseTensorReaderDimSizes(!TensorReader,
     memref<?xindex>) -> () attributes { llvm.emit_c_interface }
   func.func private @getSparseTensorReaderNextF32(!TensorReader,
-    memref<?xindex>, memref<f32>) -> () attributes { llvm.emit_c_interface }
+    memref<?xindex>, memref<f32>) -> (i1) attributes { llvm.emit_c_interface }
 
   func.func private @createSparseTensorWriter(!Filename) -> (!TensorWriter)
   func.func private @delSparseTensorWriter(!TensorWriter)
   func.func private @outSparseTensorWriterMetaData(!TensorWriter, index, index,
     memref<?xindex>) -> () attributes { llvm.emit_c_interface }
   func.func private @outSparseTensorWriterNextF32(!TensorWriter, index,
     memref<?xindex>, memref<f32>) -> () attributes { llvm.emit_c_interface }
 
func.func @readTensorFile(%tensor: !TensorReader)
     %nnz = call @getSparseTensorReaderNNZ(%tensor) : (!TensorReader) -> index
 
     // Assume rank == 2.
     %x0s = memref.alloc(%nnz) : memref<?xindex>
     %x1s = memref.alloc(%nnz) : memref<?xindex>
     %vs = memref.alloc(%nnz) : memref<?xf32>
     %indices = memref.alloc(%rank) : memref<?xindex>
     %value = memref.alloca() : memref<f32>
-    scf.for %i = %c0 to %nnz step %c1 {
-      func.call @getSparseTensorReaderNextF32(%tensor, %indices, %value)
-        : (!TensorReader, memref<?xindex>, memref<f32>) -> ()
+    %count = scf.while (%i = %c0) : (index) -> (index) {
+      %r = func.call @getSparseTensorReaderNextF32(%tensor, %indices, %value)
+        : (!TensorReader, memref<?xindex>, memref<f32>) -> (i1)
+      scf.condition(%r) %i : index
+    } do {
+      ^bb0(%i : index):
       // TODO: can we use memref.subview to avoid the need for the %value
       //       buffer?
       %v = memref.load %value[] : memref<f32>
       memref.store %v, %vs[%i] : memref<?xf32>
       %i0 = memref.load %indices[%c0] : memref<?xindex>
       memref.store %i0, %x0s[%i] : memref<?xindex>
       %i1 = memref.load %indices[%c1] : memref<?xindex>
       memref.store %i1, %x1s[%i] : memref<?xindex>
+      %n = arith.addi %i, %c1 : index
+      scf.yield %n : index
     }
 
     // Release the resource for the indices.
     memref.dealloc %indices : memref<?xindex>
     return %x0s, %x1s, %vs : memref<?xindex>, memref<?xindex>, memref<?xf32>
   }
 
   // Reads a COO tensor from the given file name and prints its content.
func.func @createTensorFileFrom(%fileName0: !Filename, %fileName1: !Filename) {
     %rank = call @getSparseTensorReaderRank(%tensor0) : (!TensorReader) -> index
     %nnz = call @getSparseTensorReaderNNZ(%tensor0) : (!TensorReader) -> index
     %dimSizes = memref.alloc(%rank) : memref<?xindex>
     func.call @getSparseTensorReaderDimSizes(%tensor0,%dimSizes)
       : (!TensorReader, memref<?xindex>) -> ()
     call @outSparseTensorWriterMetaData(%tensor1, %rank, %nnz, %dimSizes)
       : (!TensorWriter, index, index, memref<?xindex>) -> ()
 
-    //TODO: handle isSymmetric.
     // Assume rank == 2.
     %indices = memref.alloc(%rank) : memref<?xindex>
     %value = memref.alloca() : memref<f32>
-    scf.for %i = %c0 to %nnz step %c1 {
-      func.call @getSparseTensorReaderNextF32(%tensor0, %indices, %value)
-        : (!TensorReader, memref<?xindex>, memref<f32>) -> ()
+    scf.while  : () -> () {
+      %r = func.call @getSparseTensorReaderNextF32(%tensor0, %indices, %value)
+        : (!TensorReader, memref<?xindex>, memref<f32>) -> (i1)
+      scf.condition(%r)
+    } do {
       func.call @outSparseTensorWriterNextF32(%tensor1, %rank, %indices, %value)
         : (!TensorWriter, index, memref<?xindex>, memref<f32>) -> ()
+        scf.yield
     }
 
     // Release the resources.
     call @delSparseTensorReader(%tensor0) : (!TensorReader) -> ()
     call @delSparseTensorWriter(%tensor1) : (!TensorWriter) -> ()
     memref.dealloc %dimSizes : memref<?xindex>
     memref.dealloc %indices : memref<?xindex>