From ceda1ae9a7b061a98b1309949c73c9958e2a94d7 Mon Sep 17 00:00:00 2001
From: wren romano <2998727+wrengr@users.noreply.github.com>
Date: Wed, 5 Jan 2022 13:46:15 -0800
Subject: [PATCH] [mlir][sparse] Strengthening first arguments of fromCOO/toCOO

Better capturing of invariants

Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D116700
---
 mlir/lib/ExecutionEngine/SparseTensorUtils.cpp | 18 +++++++++---------
 1 file changed, 9 insertions(+), 9 deletions(-)

diff --git a/mlir/lib/ExecutionEngine/SparseTensorUtils.cpp b/mlir/lib/ExecutionEngine/SparseTensorUtils.cpp
index 3bd405738eeb..281016f785af 100644
--- a/mlir/lib/ExecutionEngine/SparseTensorUtils.cpp
+++ b/mlir/lib/ExecutionEngine/SparseTensorUtils.cpp
@@ -269,9 +269,10 @@ public:
         pointers[r].push_back(0);
     // Then assign contents from coordinate scheme tensor if provided.
     if (tensor) {
-      uint64_t nnz = tensor->getElements().size();
+      const std::vector<Element<V>> &elements = tensor->getElements();
+      uint64_t nnz = elements.size();
       values.reserve(nnz);
-      fromCOO(tensor, 0, nnz, 0);
+      fromCOO(elements, 0, nnz, 0);
     } else if (allDense) {
       values.resize(sz, 0);
     }
@@ -367,7 +368,7 @@ public:
     std::vector<uint64_t> reord(rank);
     for (uint64_t r = 0; r < rank; r++)
       reord[r] = perm[rev[r]];
-    toCOO(tensor, reord, 0, 0);
+    toCOO(*tensor, reord, 0, 0);
     assert(tensor->getElements().size() == values.size());
     return tensor;
   }
@@ -402,9 +403,8 @@ private:
   /// Initializes sparse tensor storage scheme from a memory-resident sparse
   /// tensor in coordinate scheme. This method prepares the pointers and
   /// indices arrays under the given per-dimension dense/sparse annotations.
-  void fromCOO(SparseTensorCOO<V> *tensor, uint64_t lo, uint64_t hi,
-               uint64_t d) {
-    const std::vector<Element<V>> &elements = tensor->getElements();
+  void fromCOO(const std::vector<Element<V>> &elements, uint64_t lo,
+               uint64_t hi, uint64_t d) {
     // Once dimensions are exhausted, insert the numerical values.
     assert(d <= getRank());
     if (d == getRank()) {
@@ -432,7 +432,7 @@ private:
           endDim(d + 1);
         full++;
       }
-      fromCOO(tensor, lo, seg, d + 1);
+      fromCOO(elements, lo, seg, d + 1);
       // And move on to next segment in interval.
       lo = seg;
     }
@@ -449,12 +449,12 @@ private:
 
   /// Stores the sparse tensor storage scheme into a memory-resident sparse
   /// tensor in coordinate scheme.
-  void toCOO(SparseTensorCOO<V> *tensor, std::vector<uint64_t> &reord,
+  void toCOO(SparseTensorCOO<V> &tensor, std::vector<uint64_t> &reord,
              uint64_t pos, uint64_t d) {
     assert(d <= getRank());
     if (d == getRank()) {
       assert(pos < values.size());
-      tensor->add(idx, values[pos]);
+      tensor.add(idx, values[pos]);
     } else if (isCompressedDim(d)) {
       // Sparse dimension.
       for (uint64_t ii = pointers[d][pos]; ii < pointers[d][pos + 1]; ii++) {