[mlir][sparse] Add expand_symmetry attribute to the new operator.

[bixia]

Summary

The attribute tells the operator to handle symmetric structures for 2D tensors.
By default, the operator assumes the input tensor is not symmetric.

[Peiming]

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

42

we should use *operation according to Alex

[wrengr]

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

44–45

The default should be the other way around: assume non-symmetric unless $symmetric is passed.

[aartbik]

As discussed offline, until that time that we support symmetry explicitly in the sparse tensor type somehow, let's make an attribute "expand_symmetric" which will do the a_ij/aji insertion when a symmetric matrix is read from e.g. MM. In all other case, we just read the data provided, marked symmetric or not.

[bixia]

Rename the attribute and reverse the default value.

[bixia]

Change operator to operation.

[aartbik]

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

43

Rephrase: to make symmetry in external formats explicit in the storage. That is, ....
Also add a note on that this claims more storage than a pure symmetric storage, and thus may cause a bad performance hit.
True symmetric storage is planned for the future.

[bixia]

Update document to address review comment.

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

 class SparseTensor_Op<string mnemonic, list<Trait> traits = []>
   : Op<SparseTensor_Dialect, mnemonic, traits>;
 
 //===----------------------------------------------------------------------===//
 // Sparse Tensor Operations.
 //===----------------------------------------------------------------------===//
 
 def SparseTensor_NewOp : SparseTensor_Op<"new", [Pure]>,
-    Arguments<(ins AnyType:$source)>,
+    Arguments<(ins AnyType:$source, UnitAttr:$expandSymmetry)>,
     Results<(outs AnySparseTensor:$result)> {
   string summary = "Materializes a new sparse tensor from given source";
   string description = [{
     Materializes a sparse tensor with contents taken from an opaque pointer
     provided by `source`. For targets that have access to a file system,
     for example, this pointer may be a filename (or file) of a sparse
     tensor in a particular external storage format. The form of the operation
     is kept deliberately very general to allow for alternative implementations
     in the future, such as pointers to buffers or runnable initialization
     code. The operation is provided as an anchor that materializes a properly
     typed sparse tensor with inital contents into a computation.
 
+    An optional attribute `expandSymmetry` can be used to extend this operator
+    to support symmetric structures for 2D tensors. That is, when the attribute
+    presents and a non-zero value is discovered at (i, j) where i!=j, we add the
+    same value to (j, i).
+
     Example:
 
     ```mlir
     sparse_tensor.new %source : !Source to tensor<1024x1024xf64, #CSR>
     ```
   }];
-  let assemblyFormat = "$source attr-dict `:` type($source) `to` type($result)";
+  let assemblyFormat = "(`expand_symmetry` $expandSymmetry^)? $source attr-dict"
+                       "`:` type($source) `to` type($result)";
+  let hasVerifier = 1;
 }
 
 def SparseTensor_ConvertOp : SparseTensor_Op<"convert",
   [Pure, SameOperandsAndResultElementType]>,
     Arguments<(ins AnyTensor:$source)>,
     Results<(outs AnyTensor:$dest)> {
   string summary = "Converts between different tensor types";
   string description = [{

mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp

 
 static LogicalResult isMatchingWidth(Value result, unsigned width) {
   Type etp = result.getType().cast<MemRefType>().getElementType();
   if ((width == 0 && etp.isIndex()) || (width > 0 && etp.isInteger(width)))
     return success();
   return failure();
 }
 
+LogicalResult NewOp::verify() {
+  if (getExpandSymmetry() &&
+      getResult().getType().cast<RankedTensorType>().getRank() != 2)
+    return emitOpError("expand_symmetry can only be used for 2D tensors");
+  return success();
+}
+
 LogicalResult ConvertOp::verify() {
   if (auto tp1 = getSource().getType().dyn_cast<RankedTensorType>()) {
     if (auto tp2 = getDest().getType().dyn_cast<RankedTensorType>()) {
       if (tp1.getRank() != tp2.getRank())
         return emitError("unexpected conversion mismatch in rank");
       auto shape1 = tp1.getShape();
       auto shape2 = tp2.getShape();
       // Accept size matches between the source and the destination type

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

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) {
+    // The verifier ensures only 2D tensors can have the expandSymmetry flag.
+    if (rank == 2 && op.getExpandSymmetry()) {
       symmetric =
           createFuncCall(rewriter, loc, "getSparseTensorReaderIsSymmetric",
                          {rewriter.getI1Type()}, {reader}, EmitCInterface::Off)
               .getResult(0);
     } else {
       symmetric = Value();
     }
     Type eltTp = dstTp.getElementType();

mlir/test/Dialect/SparseTensor/invalid.mlir

func.func @sparse_wrong_arity_compression(%arg0: memref<?xf64>,
   // expected-error@+1 {{'sparse_tensor.compress' op incorrect number of indices}}
   sparse_tensor.compress %arg0, %arg1, %arg2, %arg3 into %arg4[%arg5,%arg5]
     : memref<?xf64>, memref<?xi1>, memref<?xindex>, tensor<8x8xf64, #CSR>
   return
 }
 
 // -----
 
+#SparseVector = #sparse_tensor.encoding<{dimLevelType = ["compressed"]}>
+
+func.func @sparse_new(%arg0: !llvm.ptr<i8>) {
+  // expected-error@+1 {{expand_symmetry can only be used for 2D tensors}}
+  %0 = sparse_tensor.new expand_symmetry %arg0 : !llvm.ptr<i8> to tensor<128xf64, #SparseVector>
+  return
+}
+
+// -----
+
 func.func @sparse_convert_unranked(%arg0: tensor<*xf32>) -> tensor<10xf32> {
   // expected-error@+1 {{unexpected type in convert}}
   %0 = sparse_tensor.convert %arg0 : tensor<*xf32> to tensor<10xf32>
   return %0 : tensor<10xf32>
 }
 
 // -----
 

mlir/test/Dialect/SparseTensor/rewriting_for_codegen.mlir

 // RUN: mlir-opt %s -post-sparsification-rewrite="enable-runtime-library=false enable-convert=false" | \
 // RUN: FileCheck %s
 
 #CSR = #sparse_tensor.encoding<{
   dimLevelType = ["dense", "compressed"]
 }>
 
-// CHECK-LABEL:   func.func @sparse_new(
+// CHECK-LABEL:   func.func @sparse_new_symmetry(
 // CHECK-SAME:    %[[A:.*]]: !llvm.ptr<i8>) -> tensor<?x?xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "compressed" ] }>> {
 // CHECK-DAG:     %[[C2:.*]] = arith.constant 2 : index
 // 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:             scf.yield %[[T1]]
 // CHECK:           scf.yield %[[T3]]
 // CHECK:         }
 // CHECK:         call @delSparseTensorReader(%[[R]])
 // CHECK:         %[[T5:.*]] = sparse_tensor.load %[[T2]] hasInserts
 // CHECK:         %[[R:.*]] = sparse_tensor.convert %[[T5]]
 // CHECK:         bufferization.dealloc_tensor %[[T5]]
 // CHECK:         return %[[R]]
+func.func @sparse_new_symmetry(%arg0: !llvm.ptr<i8>) -> tensor<?x?xf32, #CSR> {
+  %0 = sparse_tensor.new expand_symmetry %arg0 : !llvm.ptr<i8> to tensor<?x?xf32, #CSR>
+  return %0 : tensor<?x?xf32, #CSR>
+}
+
+// CHECK-LABEL:   func.func @sparse_new(
+// CHECK-SAME:    %[[A:.*]]: !llvm.ptr<i8>) -> tensor<?x?xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "compressed" ] }>> {
+// CHECK-DAG:     %[[C2:.*]] = arith.constant 2 : index
+// 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:         %[[VB:.*]] = memref.alloca()
+// CHECK:         %[[T2:.*]] = scf.for %{{.*}} = %[[C0]] to %[[N]] step %[[C1]] iter_args(%[[A2:.*]] = %[[T]])
+// CHECK:           func.call @getSparseTensorReaderNextF32(%[[R]], %[[DS]], %[[VB]])
+// 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:           scf.yield %[[T1]]
+// CHECK:         }
+// CHECK:         call @delSparseTensorReader(%[[R]])
+// CHECK:         %[[T4:.*]] = sparse_tensor.load %[[T2]] hasInserts
+// CHECK:         %[[R:.*]] = sparse_tensor.convert %[[T4]]
+// CHECK:         bufferization.dealloc_tensor %[[T4]]
+// CHECK:         return %[[R]]
 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/Dialect/SparseTensor/roundtrip.mlir

 // RUN: mlir-opt %s -split-input-file | mlir-opt -split-input-file | FileCheck %s
 
 #SparseVector = #sparse_tensor.encoding<{dimLevelType = ["compressed"]}>
 
 // CHECK-LABEL: func @sparse_new(
 // CHECK-SAME: %[[A:.*]]: !llvm.ptr<i8>)
 //       CHECK: %[[T:.*]] = sparse_tensor.new %[[A]] : !llvm.ptr<i8> to tensor<128xf64, #{{.*}}>
 //       CHECK: return %[[T]] : tensor<128xf64, #{{.*}}>
 func.func @sparse_new(%arg0: !llvm.ptr<i8>) -> tensor<128xf64, #SparseVector> {
   %0 = sparse_tensor.new %arg0 : !llvm.ptr<i8> to tensor<128xf64, #SparseVector>
   return %0 : tensor<128xf64, #SparseVector>
 }
 
 // -----
 
+#SparseMatrix = #sparse_tensor.encoding<{dimLevelType = ["compressed", "compressed"]}>
+
+// CHECK-LABEL: func @sparse_new_symmetry(
+// CHECK-SAME: %[[A:.*]]: !llvm.ptr<i8>)
+//       CHECK: %[[T:.*]] = sparse_tensor.new expand_symmetry %[[A]] : !llvm.ptr<i8> to tensor<?x?xf64, #{{.*}}>
+//       CHECK: return %[[T]] : tensor<?x?xf64, #{{.*}}>
+func.func @sparse_new_symmetry(%arg0: !llvm.ptr<i8>) -> tensor<?x?xf64, #SparseMatrix> {
+  %0 = sparse_tensor.new expand_symmetry %arg0 : !llvm.ptr<i8> to tensor<?x?xf64, #SparseMatrix>
+  return %0 : tensor<?x?xf64, #SparseMatrix>
+}
+
+// -----
+
 #SparseVector = #sparse_tensor.encoding<{dimLevelType = ["compressed"]}>
 
 // CHECK-LABEL: func @sparse_dealloc(
 // CHECK-SAME: %[[A:.*]]: tensor<128xf64, #{{.*}}>
 //       CHECK: bufferization.dealloc_tensor %[[A]] : tensor<128xf64, #{{.*}}>
 //       CHECK: return
 func.func @sparse_dealloc(%arg0: tensor<128xf64, #SparseVector>) {
   bufferization.dealloc_tensor %arg0 : tensor<128xf64, #SparseVector>

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

func.func @entry() {
     %c0 = arith.constant 0 : index
 
     // Setup memory for a single reduction scalar,
     // initialized to zero.
     %x = tensor.from_elements %d0 : tensor<f64>
 
     // Read the sparse matrix from file, construct sparse storage.
     %fileName = call @getTensorFilename(%c0) : (index) -> (!Filename)
-    %a = sparse_tensor.new %fileName : !Filename to tensor<?x?xf64, #SparseMatrix>
+    %a = sparse_tensor.new expand_symmetry %fileName : !Filename to tensor<?x?xf64, #SparseMatrix>
 
     // Call the kernel.
     %0 = call @kernel_sum_reduce(%a, %x)
       : (tensor<?x?xf64, #SparseMatrix>, tensor<f64>) -> tensor<f64>
 
     // Print the result for verification.
     //
     // CHECK: 30.2

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

func.func @entry() {
     // Setup memory for a single reduction scalar,
     // initialized to zero.
     // TODO: tensor.from_elements does not support complex.
     %alloc = bufferization.alloc_tensor() : tensor<complex<f64>>
     %x = tensor.insert %d0 into %alloc[] : tensor<complex<f64>>
 
     // Read the sparse matrix from file, construct sparse storage.
     %fileName = call @getTensorFilename(%c0) : (index) -> (!Filename)
-    %a = sparse_tensor.new %fileName : !Filename to tensor<?x?xcomplex<f64>, #SparseMatrix>
+    %a = sparse_tensor.new expand_symmetry %fileName : !Filename to tensor<?x?xcomplex<f64>, #SparseMatrix>
 
     // Call the kernel.
     %0 = call @kernel_sum_reduce(%a, %x)
       : (tensor<?x?xcomplex<f64>, #SparseMatrix>, tensor<complex<f64>>) -> tensor<complex<f64>>
 
     // Print the result for verification.
     //
     // CHECK: 30.2