[mlir][sparse] Extend the operator new rewriter to handle isSymmetric flag.

[bixia]

Fix a FileCheck test.

[aartbik]

mlir/test/Dialect/SparseTensor/rewriting_for_codegen.mlir

30 ↗

(On Diff #476141)

For the general case, this is unavoidable for now, since we don't know a priori is the matrix is symmetric.
But since this is executed O(noe) time, let's try to also provide a way that keep this innermost loop clean of branching

as discussed offline, in a follow up

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

LogicalResult matchAndRewrite(NewOp op,
     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());
 
     SmallString<18> getNextFuncName{"getSparseTensorReaderNext",
                                     primaryTypeFunctionSuffix(eltTp)};
     Value indices = dimSizes; // Reuse the indices memref to store indices.
     createFuncCall(rewriter, loc, getNextFuncName, {eltTp},
                    {reader, indices, value}, EmitCInterface::On)
         .getResult(0);
     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);
-    auto t = rewriter.create<InsertOp>(loc, v, forOp.getRegionIterArg(0),
+    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);
+    }
     rewriter.create<scf::YieldOp>(loc, ArrayRef<Value>(t));
     rewriter.setInsertionPointAfter(forOp);
     // Link SSA chain.
     cooBuffer = forOp.getResult(0);
 
     // Release the sparse tensor reader.
     createFuncCall(rewriter, loc, "delSparseTensorReader", {}, {reader},
                    EmitCInterface::Off);

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

-// RUN: mlir-opt %s --sparse-compiler | \
-// RUN: TENSOR0="%mlir_src_dir/test/Integration/data/test_symmetric.mtx" \
-// RUN: mlir-cpu-runner \
-// RUN:  -e entry -entry-point-result=void  \
-// RUN:  -shared-libs=%mlir_lib_dir/libmlir_c_runner_utils%shlibext | \
-// RUN: FileCheck %s
+// DEFINE: %{option} = enable-runtime-library=true
+// DEFINE: %{command} = mlir-opt %s --sparse-compiler=%{option} | \
+// DEFINE: TENSOR0="%mlir_src_dir/test/Integration/data/test_symmetric.mtx" \
+// DEFINE: mlir-cpu-runner \
+// DEFINE:  -e entry -entry-point-result=void  \
+// DEFINE:  -shared-libs=%mlir_lib_dir/libmlir_c_runner_utils%shlibext | \
+// DEFINE: FileCheck %s
+//
+// RUN: %{command}
+//
+// Do the same run, but now with direct IR generation.
+// REDEFINE: %{option} = enable-runtime-library=false
+// RUN: %{command}
 
 !Filename = !llvm.ptr<i8>
 
 #SparseMatrix = #sparse_tensor.encoding<{
   dimLevelType = [ "compressed", "compressed" ]
 }>
 
 #trait_sum_reduce = {