[mlir][sparse] generalize sparse storage format to many more types

[aartbik]

Summary

Rationale:
Narrower types for overhead storage yield a smaller memory footprint for
sparse tensors and thus needs to be supported. Also, more value types
need to be supported to deal with all kinds of kernels. Since the
"one-size-fits-all" sparse storage scheme implementation is used
instead of actual codegen, the library needs to be able to support
all combinations of desired types. With some crafty templating and
overloading, the actual code for this is kept reasonably sized though.

[bixia]

Thanks for your work! It looks good to me and I only have a few minor comments.

mlir/lib/ExecutionEngine/SparseUtils.cpp

77

Adds

84

Sorts

124

Instead of saying "unsupported", would it be more helpful to say something more specific, such as the requested the data type doesn't match how the data is stored in SparseTensorStorage?

159

We probably don't need this until it needs an implementation that is different from what is in the base class.

184–185

The previous word was fine. It is "pointer structure" not "pointers structure". Right?
Alternative, we may want to say "structure with pointers".

206

Similar to the above.

224

This is called sparsity in class SparseTensorStorage. Why don't we call it sparsity here?

[aartbik]

Thanks Bixia!

mlir/lib/ExecutionEngine/SparseUtils.cpp

124

I had that originally, then decided against it, but now I am convinced it is needed again ;-)
Thanks!

[aartbik]

addressed comments

mlir/integration_test/Sparse/CPU/sparse_sampled_matmul.mlir

+// RUN: mlir-opt %s \
+// RUN:   --test-sparsification="lower ptr-type=2 ind-type=2 fast-output" \
+// RUN:   --convert-linalg-to-loops \
+// RUN:   --func-bufferize --tensor-constant-bufferize --tensor-bufferize \
+// RUN:   --std-bufferize --finalizing-bufferize  \
+// RUN:   --convert-scf-to-std --convert-vector-to-llvm --convert-std-to-llvm | \
+// RUN: TENSOR0="%mlir_integration_test_dir/data/test.mtx" \
+// RUN: mlir-cpu-runner \
+// RUN:  -e entry -entry-point-result=void  \
+// RUN:  -shared-libs=%mlir_integration_test_dir/libmlir_c_runner_utils%shlibext | \
+// RUN: FileCheck %s
+
+//
+// Use descriptive names for opaque pointers.
+//
+!Filename     = type !llvm.ptr<i8>
+!SparseTensor = type !llvm.ptr<i8>
+
+#trait_sampled_dense_dense = {
+  indexing_maps = [
+    affine_map<(i,j,k) -> (i,j)>,  // S
+    affine_map<(i,j,k) -> (i,k)>,  // A
+    affine_map<(i,j,k) -> (k,j)>,  // B
+    affine_map<(i,j,k) -> (i,j)>   // X (out)
+  ],
+  sparse = [
+    [ "S", "S" ],  // S
+    [ "D", "D" ],  // A
+    [ "D", "D" ],  // B
+    [ "D", "D" ]   // X
+  ],
+  iterator_types = ["parallel", "parallel", "reduction"],
+  doc = "X(i,j) += S(i,j) SUM_k A(i,k) B(k,j)"
+}
+
+//
+// Integration test that lowers a kernel annotated as sparse to
+// actual sparse code, initializes a matching sparse storage scheme
+// from file, and runs the resulting code with the JIT compiler.
+//
+module {
+  //
+  // The kernel expressed as an annotated Linalg op. The kernel
+  // computes a sampled matrix matrix multiplication.
+  //
+  func @sampled_dense_dense(%argS: !SparseTensor,
+                            %arga: tensor<?x?xf32>,
+                            %argb: tensor<?x?xf32>,
+                            %argx: tensor<?x?xf32>) -> tensor<?x?xf32> {
+    %args = linalg.sparse_tensor %argS : !SparseTensor to tensor<?x?xf32>
+    %0 = linalg.generic #trait_sampled_dense_dense
+      ins(%args, %arga, %argb: tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>)
+      outs(%argx: tensor<?x?xf32>) {
+        ^bb(%s: f32, %a: f32, %b: f32, %x: f32):
+          %0 = mulf %a, %b : f32
+          %1 = mulf %s, %0 : f32
+          %2 = addf %x, %1 : f32
+          linalg.yield %2 : f32
+    } -> tensor<?x?xf32>
+    return %0 : tensor<?x?xf32>
+  }
+
+  //
+  // Runtime support library that is called directly from here.
+  //
+  func private @getTensorFilename(index) -> (!Filename)
+  func private @newSparseTensor(!Filename, memref<?xi1>, index, index, index) -> (!SparseTensor)
+  func private @delSparseTensor(!SparseTensor) -> ()
+  func private @print_memref_f32(%ptr : tensor<*xf32>)
+
+  //
+  // Main driver that reads matrix from file and calls the sparse kernel.
+  //
+  func @entry() {
+    %d0 = constant 0.0 : f32
+    %c0 = constant 0 : index
+    %c1 = constant 1 : index
+    %c2 = constant 2 : index
+    %c5 = constant 5 : index
+    %c10 = constant 10 : index
+
+    // Mark both dimensions of the matrix as sparse and encode the
+    // storage scheme types (this must match the metadata in the
+    // trait and compiler switches).
+    %annotations = alloc(%c2) : memref<?xi1>
+    %sparse = constant true
+    store %sparse, %annotations[%c0] : memref<?xi1>
+    store %sparse, %annotations[%c1] : memref<?xi1>
+    %i32 = constant 3 : index
+    %f32 = constant 1 : index
+
+    // Setup memory for the dense matrices and initialize.
+    %adata = alloc(%c5, %c10) : memref<?x?xf32>
+    %bdata = alloc(%c10, %c5) : memref<?x?xf32>
+    %xdata = alloc(%c5,  %c5) : memref<?x?xf32>
+    scf.for %i = %c0 to %c5 step %c1 {
+      scf.for %j = %c0 to %c5 step %c1 {
+        store %d0, %xdata[%i, %j] : memref<?x?xf32>
+      }
+      %p = addi %i, %c1 : index
+      %q = index_cast %p : index to i32
+      %d = sitofp %q : i32 to f32
+      scf.for %j = %c0 to %c10 step %c1 {
+        store %d, %adata[%i, %j] : memref<?x?xf32>
+        store %d, %bdata[%j, %i] : memref<?x?xf32>
+      }
+    }
+    %a = tensor_load %adata : memref<?x?xf32>
+    %b = tensor_load %bdata : memref<?x?xf32>
+    %x = tensor_load %xdata : memref<?x?xf32>
+
+    // Read the sparse matrix from file, construct sparse storage
+    // according to <sparse,sparse> in memory, and call the kernel.
+    %fileName = call @getTensorFilename(%c0) : (index) -> (!Filename)
+    %s = call @newSparseTensor(%fileName, %annotations, %i32, %i32, %f32)
+      : (!Filename, memref<?xi1>, index, index, index) -> (!SparseTensor)
+    %0 = call @sampled_dense_dense(%s, %a, %b, %x)
+       : (!SparseTensor, tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>) -> tensor<?x?xf32>
+
+    // Print the result for verification.
+    //
+    // CHECK: ( 10, 0, 0, 56, 0 )
+    // CHECK: ( 0, 80, 0, 0, 250 )
+    // CHECK: ( 0, 0, 270, 0, 0 )
+    // CHECK: ( 164, 0, 0, 640, 0 )
+    // CHECK: ( 0, 520, 0, 0, 1250 )
+    //
+    %r = tensor_to_memref %0 : memref<?x?xf32>
+    scf.for %i = %c0 to %c5 step %c1 {
+      %v = vector.transfer_read %r[%i, %c0], %d0: memref<?x?xf32>, vector<5xf32>
+      vector.print %v : vector<5xf32>
+    }
+
+    // Release the resources.
+    call @delSparseTensor(%s) : (!SparseTensor) -> ()
+    dealloc %adata : memref<?x?xf32>
+    dealloc %bdata : memref<?x?xf32>
+    dealloc %xdata : memref<?x?xf32>
+
+    return
+  }
+}

mlir/integration_test/Sparse/CPU/sparse_sum.mlir

func @kernel_sum_reduce(%argA: !SparseTensor,
     } -> tensor<f64>
     return %0 : tensor<f64>
   }
 
   //
   // Runtime support library that is called directly from here.
   //
   func private @getTensorFilename(index) -> (!Filename)
-  func private @newSparseTensor(!Filename, memref<?xi1>) -> (!SparseTensor)
+  func private @newSparseTensor(!Filename, memref<?xi1>, index, index, index) -> (!SparseTensor)
   func private @delSparseTensor(!SparseTensor) -> ()
   func private @print_memref_f64(%ptr : tensor<*xf64>)
 
   //
   // Main driver that reads matrix from file and calls the sparse kernel.
   //
   func @entry() {
     %d0 = constant 0.0 : f64
     %c0 = constant 0 : index
     %c1 = constant 1 : index
     %c2 = constant 2 : index
 
-    // Mark both dimensions of the matrix as sparse
-    // (this must match the annotation in the trait).
+    // Mark both dimensions of the matrix as sparse and encode the
+    // storage scheme types (this must match the metadata in the
+    // trait and compiler switches).
     %annotations = alloc(%c2) : memref<?xi1>
     %sparse = constant true
     store %sparse, %annotations[%c0] : memref<?xi1>
     store %sparse, %annotations[%c1] : memref<?xi1>
+    %i64 = constant 2 : index
+    %f64 = constant 0 : index
 
     // Setup memory for a single reduction scalar,
     // initialized to zero.
     %xdata = alloc() : memref<f64>
     store %d0, %xdata[] : memref<f64>
     %x = tensor_load %xdata : memref<f64>
 
     // Read the sparse matrix from file, construct sparse storage
     // according to <sparse,sparse> in memory, and call the kernel.
     %fileName = call @getTensorFilename(%c0) : (index) -> (!Filename)
-    %a = call @newSparseTensor(%fileName, %annotations)
-      : (!Filename, memref<?xi1>) -> (!SparseTensor)
+    %a = call @newSparseTensor(%fileName, %annotations, %i64, %i64, %f64)
+      : (!Filename, memref<?xi1>, index, index, index) -> (!SparseTensor)
     %0 = call @kernel_sum_reduce(%a, %x)
       : (!SparseTensor, tensor<f64>) -> tensor<f64>
 
     // Print the result for verification.
     //
     // CHECK: 28.2
     //
     %m = tensor_to_memref %0 : memref<f64>

mlir/lib/Dialect/Linalg/Transforms/SparseLowering.cpp

public:
   LogicalResult
   matchAndRewrite(linalg::SparseTensorToPointersMemRefOp op,
                   ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const override {
     Type resType = op.getType();
     Type eltType = resType.cast<ShapedType>().getElementType();
     StringRef name;
     if (eltType.isIndex() || eltType.isInteger(64))
-      name = "sparsePtrsI64";
+      name = "sparsePointers64";
+    else if (eltType.isInteger(32))
+      name = "sparsePointers32";
     else
       return failure();
     rewriter.replaceOpWithNewOp<CallOp>(
         op, resType, getFunc(op, name, resType, operands), operands);
     return success();
   }
 };
 
 /// Sparse conversion rule for index accesses.
 class TensorToIndicesConverter
     : public OpConversionPattern<linalg::SparseTensorToIndicesMemRefOp> {
 public:
   using OpConversionPattern::OpConversionPattern;
   LogicalResult
   matchAndRewrite(linalg::SparseTensorToIndicesMemRefOp op,
                   ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const override {
     Type resType = op.getType();
     Type eltType = resType.cast<ShapedType>().getElementType();
     StringRef name;
     if (eltType.isIndex() || eltType.isInteger(64))
-      name = "sparseIndxsI64";
+      name = "sparseIndices64";
+    else if (eltType.isInteger(32))
+      name = "sparseIndices32";
     else
       return failure();
     rewriter.replaceOpWithNewOp<CallOp>(
         op, resType, getFunc(op, name, resType, operands), operands);
     return success();
   }
 };
 
 /// Sparse conversion rule for value accesses.
 class TensorToValuesConverter
     : public OpConversionPattern<linalg::SparseTensorToValuesMemRefOp> {
 public:
   using OpConversionPattern::OpConversionPattern;
   LogicalResult
   matchAndRewrite(linalg::SparseTensorToValuesMemRefOp op,
                   ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const override {
     Type resType = op.getType();
     Type eltType = resType.cast<ShapedType>().getElementType();
     StringRef name;
     if (eltType.isF64())
-      name = "sparseValsF64";
+      name = "sparseValuesF64";
+    else if (eltType.isF32())
+      name = "sparseValuesF32";
     else
       return failure();
     rewriter.replaceOpWithNewOp<CallOp>(
         op, resType, getFunc(op, name, resType, operands), operands);
     return success();
   }
 };
 

mlir/lib/ExecutionEngine/SparseUtils.cpp

 #include <cinttypes>
 #include <cstdio>
 #include <cstdlib>
 #include <cstring>
 #include <vector>
 
 //===----------------------------------------------------------------------===//
 //
-// Internal support for reading sparse tensors in one of the following
-// external file formats:
+// Internal support for storing and reading sparse tensors.
+//
+// The following memory-resident sparse storage schemes are supported:
+//
+// (a) A coordinate scheme for temporarily storing and lexicographically
+//     sorting a sparse tensor by index.
+//
+// (b) A "one-size-fits-all" sparse storage scheme defined by per-rank
+//     sparse/dense annnotations to be used by generated MLIR code.
+//
+// The following external formats are supported:
 //
 // (1) Matrix Market Exchange (MME): *.mtx
 //     https://math.nist.gov/MatrixMarket/formats.html
 //
 // (2) Formidable Repository of Open Sparse Tensors and Tools (FROSTT): *.tns
 //     http://frostt.io/tensors/file-formats.html
 //
 //===----------------------------------------------------------------------===//

 /// by indices before passing it back to the client (most packed storage
 /// formats require the elements to appear in lexicographic index order).
 struct SparseTensor {
 public:
   SparseTensor(const std::vector<uint64_t> &szs, uint64_t capacity)
       : sizes(szs), pos(0) {
     elements.reserve(capacity);
   }
-  // Add element as indices and value.
+  /// Adds element as indices and value.
   void add(const std::vector<uint64_t> &ind, double val) {
     assert(sizes.size() == ind.size());
     for (int64_t r = 0, rank = sizes.size(); r < rank; r++)
       assert(ind[r] < sizes[r]); // within bounds
     elements.emplace_back(Element(ind, val));
   }
-  // Sort elements lexicographically by index.
+  /// Sorts elements lexicographically by index.
   void sort() { std::sort(elements.begin(), elements.end(), lexOrder); }
-  // Primitive one-time iteration.
+  /// Primitive one-time iteration.
   const Element &next() { return elements[pos++]; }
+  /// Getter for sizes array.
+  const std::vector<uint64_t> &getSizes() const { return sizes; }
+  /// Getter for elements array.
+  const std::vector<Element> &getElements() const { return elements; }
 
 private:
-  // Returns true if indices of e1 < indices of e2.
+  /// Returns true if indices of e1 < indices of e2.
   static bool lexOrder(const Element &e1, const Element &e2) {
     assert(e1.indices.size() == e2.indices.size());
     for (int64_t r = 0, rank = e1.indices.size(); r < rank; r++) {
       if (e1.indices[r] == e2.indices[r])
         continue;
       return e1.indices[r] < e2.indices[r];
     }
     return false;
   }
-
-public:
   std::vector<uint64_t> sizes; // per-rank dimension sizes
   std::vector<Element> elements;
   uint64_t pos;
 };
 
+/// Abstract base class of sparse tensor storage. Note that we use
+/// function overloading to implement "partial" method specialization.
+class SparseTensorStorageBase {
+public:
+  virtual uint64_t getDimSize(uint64_t) = 0;
+  virtual void getPointers(std::vector<uint64_t> **, uint64_t) { fatal("p64"); }
+  virtual void getPointers(std::vector<uint32_t> **, uint64_t) { fatal("p32"); }
+  virtual void getIndices(std::vector<uint64_t> **, uint64_t) { fatal("i64"); }
+  virtual void getIndices(std::vector<uint32_t> **, uint64_t) { fatal("i32"); }
+  virtual void getValues(std::vector<double> **) { fatal("valf64"); }
+  virtual void getValues(std::vector<float> **) { fatal("valf32"); }
+  virtual ~SparseTensorStorageBase() {}
+
+private:
+  void fatal(const char *tp) {
+    fprintf(stderr, "unsupported %s\n", tp);
+    exit(1);
+  }
+};
+
 /// A memory-resident sparse tensor using a storage scheme based on per-rank
 /// annotations on dense/sparse. This data structure provides a bufferized
 /// form of an imaginary SparseTensorType, until such a type becomes a
 /// first-class citizen of MLIR. In contrast to generating setup methods for
 /// each differently annotated sparse tensor, this method provides a convenient
 /// "one-size-fits-all" solution that simply takes an input tensor and
 /// annotations to implement all required setup in a general manner.
 template <typename P, typename I, typename V>
-class SparseTensorStorage {
+class SparseTensorStorage : public SparseTensorStorageBase {
 public:
   /// Constructs sparse tensor storage scheme following the given
   /// per-rank dimension dense/sparse annotations.
   SparseTensorStorage(SparseTensor *tensor, bool *sparsity)
-      : sizes(tensor->sizes), positions(sizes.size()), indices(sizes.size()) {
+      : sizes(tensor->getSizes()), pointers(sizes.size()),
+        indices(sizes.size()) {
     // Provide hints on capacity.
     // TODO: needs fine-tuning based on sparsity
-    values.reserve(tensor->elements.size());
+    values.reserve(tensor->getElements().size());
     for (uint64_t d = 0, s = 1, rank = sizes.size(); d < rank; d++) {
-      s *= tensor->sizes[d];
+      s *= tensor->getSizes()[d];
       if (sparsity[d]) {
-        positions[d].reserve(s + 1);
+        pointers[d].reserve(s + 1);
         indices[d].reserve(s);
         s = 1;
       }
     }
     // Then setup the tensor.
-    traverse(tensor, sparsity, 0, tensor->elements.size(), 0);
+    traverse(tensor, sparsity, 0, tensor->getElements().size(), 0);
+  }
+
+  virtual ~SparseTensorStorage() {}
+
+  uint64_t getDimSize(uint64_t d) override { return sizes[d]; }
+  void getPointers(std::vector<P> **out, uint64_t d) override {
+    *out = &pointers[d];
+  }
+  void getIndices(std::vector<I> **out, uint64_t d) override {
+    *out = &indices[d];
   }
+  void getValues(std::vector<V> **out) override { *out = &values; }
 
 private:
   /// Initializes sparse tensor storage scheme from a memory-resident
   /// representation of an external sparse tensor. This method prepares
   /// the pointers and indices arrays under the given per-rank dimension
   /// dense/sparse annotations.
   void traverse(SparseTensor *tensor, bool *sparsity, uint64_t lo, uint64_t hi,
                 uint64_t d) {
-    const std::vector<Element> &elements = tensor->elements;
+    const std::vector<Element> &elements = tensor->getElements();
     // Once dimensions are exhausted, insert the numerical values.
     if (d == sizes.size()) {
       values.push_back(lo < hi ? elements[lo].value : 0.0);
       return;
     }
     // Prepare a sparse pointer structure at this dimension.
-    if (sparsity[d] && positions[d].empty())
-      positions[d].push_back(0);
+    if (sparsity[d] && pointers[d].empty())
+      pointers[d].push_back(0);
     // Visit all elements in this interval.
     uint64_t full = 0;
     while (lo < hi) {
       // Find segment in interval with same index elements in this dimension.
       unsigned idx = elements[lo].indices[d];
       unsigned seg = lo + 1;
       while (seg < hi && elements[seg].indices[d] == idx)
         seg++;
       // Handle segment in interval for sparse or dense dimension.
       if (sparsity[d]) {
         indices[d].push_back(idx);
       } else {
         for (; full < idx; full++)
           traverse(tensor, sparsity, 0, 0, d + 1); // pass empty
         full++;
       }
       traverse(tensor, sparsity, lo, seg, d + 1);
       // And move on to next segment in interval.
       lo = seg;
     }
     // Finalize the sparse pointer structure at this dimension.
     if (sparsity[d]) {
-      positions[d].push_back(indices[d].size());
+      pointers[d].push_back(indices[d].size());
     } else {
-      for (uint64_t sz = tensor->sizes[d]; full < sz; full++)
+      for (uint64_t sz = tensor->getSizes()[d]; full < sz; full++)
         traverse(tensor, sparsity, 0, 0, d + 1); // pass empty
     }
   }
 
-public:
+private:
   std::vector<uint64_t> sizes; // per-rank dimension sizes
-  std::vector<std::vector<P>> positions;
+  std::vector<std::vector<P>> pointers;
   std::vector<std::vector<I>> indices;
   std::vector<V> values;
 };
 
-typedef SparseTensorStorage<uint64_t, uint64_t, double>
-    SparseTensorStorageU64U64F64;
+/// Templated reader.
+template <typename P, typename I, typename V>
+void *newSparseTensor(char *filename, bool *sparsity) {
+  uint64_t idata[64];
+  SparseTensor *t = static_cast<SparseTensor *>(openTensorC(filename, idata));
+  SparseTensorStorageBase *tensor =
+      new SparseTensorStorage<P, I, V>(t, sparsity);
+  delete t;
+  return tensor;
+}
 
 /// Helper to convert string to lower case.
 static char *toLower(char *token) {
   for (char *c = token; *c; c++)
     *c = tolower(*c);
   return token;
 }
 

 // Note that input parameters in the "MLIRized" version of a function mimic
 // the data layout of a MemRef<?xT> (but cannot use a direct struct). The
 // output parameter uses a direct struct.
 //
 //===----------------------------------------------------------------------===//
 
 extern "C" {
 
-/// Cannot use templates with C linkage.
-
-struct MemRef1DU64 {
-  const uint64_t *base;
-  const uint64_t *data;
-  uint64_t off;
-  uint64_t sizes[1];
-  uint64_t strides[1];
-};
-
-struct MemRef1DF64 {
-  const double *base;
-  const double *data;
-  uint64_t off;
-  uint64_t sizes[1];
-  uint64_t strides[1];
-};
-
 /// Reads in a sparse tensor with the given filename. The call yields a
 /// pointer to an opaque memory-resident sparse tensor object that is only
 /// understood by other methods in the sparse runtime support library. An
 /// array parameter is used to pass the rank, the number of nonzero elements,
 /// and the dimension sizes (one per rank).
 void *openTensorC(char *filename, uint64_t *idata) {
   // Open the file.
   FILE *file = fopen(filename, "r");

 /// defined with the naming convention ${TENSOR0}, ${TENSOR1}, etc.
 char *getTensorFilename(uint64_t id) {
   char var[80];
   sprintf(var, "TENSOR%" PRIu64, id);
   char *env = getenv(var);
   return env;
 }
 
-///
-/// Sparse primitives that support an opaque implementation of a bufferized
-/// SparseTensor in MLIR. This could be replaced by actual codegen in MLIR.
-///
+//===----------------------------------------------------------------------===//
+//
+// Public API of the sparse runtime support library that support an opaque
+// implementation of a bufferized SparseTensor in MLIR. This could be replaced
+// by actual codegen in MLIR.
+//
+//===----------------------------------------------------------------------===//
+
+// Cannot use templates with C linkage.
 
-void *newSparseTensorC(char *filename, bool *annotations) {
-  uint64_t idata[64];
-  SparseTensor *t = static_cast<SparseTensor *>(openTensorC(filename, idata));
-  SparseTensorStorageU64U64F64 *tensor =
-      new SparseTensorStorageU64U64F64(t, annotations);
-  delete t;
-  return tensor;
-}
+struct MemRef1DU64 {
+  const uint64_t *base;
+  const uint64_t *data;
+  uint64_t off;
+  uint64_t sizes[1];
+  uint64_t strides[1];
+};
+
+struct MemRef1DU32 {
+  const uint32_t *base;
+  const uint32_t *data;
+  uint64_t off;
+  uint64_t sizes[1];
+  uint64_t strides[1];
+};
+
+struct MemRef1DF64 {
+  const double *base;
+  const double *data;
+  uint64_t off;
+  uint64_t sizes[1];
+  uint64_t strides[1];
+};
+
+struct MemRef1DF32 {
+  const float *base;
+  const float *data;
+  uint64_t off;
+  uint64_t sizes[1];
+  uint64_t strides[1];
+};
+
+enum TypeEnum : uint64_t { kF64 = 0, kF32 = 1, kU64 = 2, kU32 = 3 };
 
-/// "MLIRized" version.
 void *newSparseTensor(char *filename, bool *abase, bool *adata, uint64_t aoff,
-                      uint64_t asize, uint64_t astride) {
+                      uint64_t asize, uint64_t astride, uint64_t ptrTp,
+                      uint64_t indTp, uint64_t valTp) {
   assert(astride == 1);
-  return newSparseTensorC(filename, abase + aoff);
+  bool *sparsity = abase + aoff;
+  if (ptrTp == kU64 && indTp == kU64 && valTp == kF64)
+    return newSparseTensor<uint64_t, uint64_t, double>(filename, sparsity);
+  if (ptrTp == kU64 && indTp == kU64 && valTp == kF32)
+    return newSparseTensor<uint64_t, uint64_t, float>(filename, sparsity);
+  if (ptrTp == kU64 && indTp == kU32 && valTp == kF64)
+    return newSparseTensor<uint64_t, uint32_t, double>(filename, sparsity);
+  if (ptrTp == kU64 && indTp == kU32 && valTp == kF32)
+    return newSparseTensor<uint64_t, uint32_t, float>(filename, sparsity);
+  if (ptrTp == kU32 && indTp == kU64 && valTp == kF64)
+    return newSparseTensor<uint32_t, uint64_t, double>(filename, sparsity);
+  if (ptrTp == kU32 && indTp == kU64 && valTp == kF32)
+    return newSparseTensor<uint32_t, uint64_t, float>(filename, sparsity);
+  if (ptrTp == kU32 && indTp == kU32 && valTp == kF64)
+    return newSparseTensor<uint32_t, uint32_t, double>(filename, sparsity);
+  if (ptrTp == kU32 && indTp == kU32 && valTp == kF32)
+    return newSparseTensor<uint32_t, uint32_t, float>(filename, sparsity);
+  fputs("unsupported combination of types\n", stderr);
+  exit(1);
 }
 
 uint64_t sparseDimSize(void *tensor, uint64_t d) {
-  return static_cast<SparseTensorStorageU64U64F64 *>(tensor)->sizes[d];
+  return static_cast<SparseTensorStorageBase *>(tensor)->getDimSize(d);
+}
+
+MemRef1DU64 sparsePointers64(void *tensor, uint64_t d) {
+  std::vector<uint64_t> *v;
+  static_cast<SparseTensorStorageBase *>(tensor)->getPointers(&v, d);
+  return {v->data(), v->data(), 0, {v->size()}, {1}};
+}
+
+MemRef1DU32 sparsePointers32(void *tensor, uint64_t d) {
+  std::vector<uint32_t> *v;
+  static_cast<SparseTensorStorageBase *>(tensor)->getPointers(&v, d);
+  return {v->data(), v->data(), 0, {v->size()}, {1}};
+}
+
+MemRef1DU64 sparseIndices64(void *tensor, uint64_t d) {
+  std::vector<uint64_t> *v;
+  static_cast<SparseTensorStorageBase *>(tensor)->getIndices(&v, d);
+  return {v->data(), v->data(), 0, {v->size()}, {1}};
 }
 
-MemRef1DU64 sparsePtrsI64(void *tensor, uint64_t d) {
-  const std::vector<uint64_t> &v =
-      static_cast<SparseTensorStorageU64U64F64 *>(tensor)->positions[d];
-  return {v.data(), v.data(), 0, {v.size()}, {1}};
+MemRef1DU32 sparseIndices32(void *tensor, uint64_t d) {
+  std::vector<uint32_t> *v;
+  static_cast<SparseTensorStorageBase *>(tensor)->getIndices(&v, d);
+  return {v->data(), v->data(), 0, {v->size()}, {1}};
 }
 
-MemRef1DU64 sparseIndxsI64(void *tensor, uint64_t d) {
-  const std::vector<uint64_t> &v =
-      static_cast<SparseTensorStorageU64U64F64 *>(tensor)->indices[d];
-  return {v.data(), v.data(), 0, {v.size()}, {1}};
+MemRef1DF64 sparseValuesF64(void *tensor) {
+  std::vector<double> *v;
+  static_cast<SparseTensorStorageBase *>(tensor)->getValues(&v);
+  return {v->data(), v->data(), 0, {v->size()}, {1}};
 }
 
-MemRef1DF64 sparseValsF64(void *tensor) {
-  const std::vector<double> &v =
-      static_cast<SparseTensorStorageU64U64F64 *>(tensor)->values;
-  return {v.data(), v.data(), 0, {v.size()}, {1}};
+MemRef1DF32 sparseValuesF32(void *tensor) {
+  std::vector<float> *v;
+  static_cast<SparseTensorStorageBase *>(tensor)->getValues(&v);
+  return {v->data(), v->data(), 0, {v->size()}, {1}};
 }
 
 void delSparseTensor(void *tensor) {
-  delete static_cast<SparseTensorStorageU64U64F64 *>(tensor);
+  delete static_cast<SparseTensorStorageBase *>(tensor);
 }
 
 } // extern "C"
 
 #endif // MLIR_CRUNNERUTILS_DEFINE_FUNCTIONS

mlir/test/Dialect/Linalg/sparse_lower.mlir

 
 // CHECK-MIR-LABEL:   func @matvec(
 // CHECK-MIR-SAME:                 %[[VAL_0:.*]]: !llvm.ptr<i8>,
 // CHECK-MIR-SAME:                 %[[VAL_1:.*]]: tensor<64xf64>,
 // CHECK-MIR-SAME:                 %[[VAL_2:.*]]: tensor<64xf64>) -> tensor<64xf64> {
 // CHECK-MIR:           %[[VAL_3:.*]] = constant 64 : index
 // CHECK-MIR:           %[[VAL_4:.*]] = constant 0 : index
 // CHECK-MIR:           %[[VAL_5:.*]] = constant 1 : index
-// CHECK-MIR:           %[[VAL_6:.*]] = call @sparsePtrsI64(%[[VAL_0]], %[[VAL_5]]) : (!llvm.ptr<i8>, index) -> memref<?xindex>
-// CHECK-MIR:           %[[VAL_7:.*]] = call @sparseIndxsI64(%[[VAL_0]], %[[VAL_5]]) : (!llvm.ptr<i8>, index) -> memref<?xindex>
-// CHECK-MIR:           %[[VAL_8:.*]] = call @sparseValsF64(%[[VAL_0]]) : (!llvm.ptr<i8>) -> memref<?xf64>
+// CHECK-MIR:           %[[VAL_6:.*]] = call @sparsePointers64(%[[VAL_0]], %[[VAL_5]]) : (!llvm.ptr<i8>, index) -> memref<?xindex>
+// CHECK-MIR:           %[[VAL_7:.*]] = call @sparseIndices64(%[[VAL_0]], %[[VAL_5]]) : (!llvm.ptr<i8>, index) -> memref<?xindex>
+// CHECK-MIR:           %[[VAL_8:.*]] = call @sparseValuesF64(%[[VAL_0]]) : (!llvm.ptr<i8>) -> memref<?xf64>
 // CHECK-MIR:           %[[VAL_9:.*]] = tensor_to_memref %[[VAL_1]] : memref<64xf64>
 // CHECK-MIR:           %[[VAL_10:.*]] = tensor_to_memref %[[VAL_2]] : memref<64xf64>
 // CHECK-MIR:           %[[VAL_11:.*]] = alloc() : memref<64xf64>
 // CHECK-MIR:           scf.for %[[VAL_12:.*]] = %[[VAL_4]] to %[[VAL_3]] step %[[VAL_5]] {
 // CHECK-MIR:             %[[VAL_13:.*]] = load %[[VAL_10]]{{\[}}%[[VAL_12]]] : memref<64xf64>
 // CHECK-MIR:             store %[[VAL_13]], %[[VAL_11]]{{\[}}%[[VAL_12]]] : memref<64xf64>
 // CHECK-MIR:           }
 // CHECK-MIR:           scf.for %[[VAL_14:.*]] = %[[VAL_4]] to %[[VAL_3]] step %[[VAL_5]] {

 
 // CHECK-LIR-LABEL:   func @matvec(
 // CHECK-LIR-SAME:                 %[[VAL_0:.*]]: !llvm.ptr<i8>,
 // CHECK-LIR-SAME:                 %[[VAL_1:.*]]: memref<64xf64>,
 // CHECK-LIR-SAME:                 %[[VAL_2:.*]]: memref<64xf64>) -> memref<64xf64> {
 // CHECK-LIR:           %[[VAL_3:.*]] = constant 64 : index
 // CHECK-LIR:           %[[VAL_4:.*]] = constant 0 : index
 // CHECK-LIR:           %[[VAL_5:.*]] = constant 1 : index
-// CHECK-LIR:           %[[VAL_6:.*]] = call @sparsePtrsI64(%[[VAL_0]], %[[VAL_5]]) : (!llvm.ptr<i8>, index) -> memref<?xindex>
-// CHECK-LIR:           %[[VAL_7:.*]] = call @sparseIndxsI64(%[[VAL_0]], %[[VAL_5]]) : (!llvm.ptr<i8>, index) -> memref<?xindex>
-// CHECK-LIR:           %[[VAL_8:.*]] = call @sparseValsF64(%[[VAL_0]]) : (!llvm.ptr<i8>) -> memref<?xf64>
+// CHECK-LIR:           %[[VAL_6:.*]] = call @sparsePointers64(%[[VAL_0]], %[[VAL_5]]) : (!llvm.ptr<i8>, index) -> memref<?xindex>
+// CHECK-LIR:           %[[VAL_7:.*]] = call @sparseIndices64(%[[VAL_0]], %[[VAL_5]]) : (!llvm.ptr<i8>, index) -> memref<?xindex>
+// CHECK-LIR:           %[[VAL_8:.*]] = call @sparseValuesF64(%[[VAL_0]]) : (!llvm.ptr<i8>) -> memref<?xf64>
 // CHECK-LIR:           %[[VAL_9:.*]] = alloc() : memref<64xf64>
 // CHECK-LIR:           scf.for %[[VAL_10:.*]] = %[[VAL_4]] to %[[VAL_3]] step %[[VAL_5]] {
 // CHECK-LIR:             %[[VAL_11:.*]] = load %[[VAL_2]]{{\[}}%[[VAL_10]]] : memref<64xf64>
 // CHECK-LIR:             store %[[VAL_11]], %[[VAL_9]]{{\[}}%[[VAL_10]]] : memref<64xf64>
 // CHECK-LIR:           }
 // CHECK-LIR:           scf.for %[[VAL_12:.*]] = %[[VAL_4]] to %[[VAL_3]] step %[[VAL_5]] {
 // CHECK-LIR:             %[[VAL_13:.*]] = load %[[VAL_6]]{{\[}}%[[VAL_12]]] : memref<?xindex>
 // CHECK-LIR:             %[[VAL_14:.*]] = addi %[[VAL_12]], %[[VAL_5]] : index

 
 // CHECK-FAST-LABEL:   func @matvec(
 // CHECK-FAST-SAME:                 %[[VAL_0:.*]]: !llvm.ptr<i8>,
 // CHECK-FAST-SAME:                 %[[VAL_1:.*]]: memref<64xf64>,
 // CHECK-FAST-SAME:                 %[[VAL_2:.*]]: memref<64xf64>) -> memref<64xf64> {
 // CHECK-FAST:           %[[VAL_3:.*]] = constant 64 : index
 // CHECK-FAST:           %[[VAL_4:.*]] = constant 0 : index
 // CHECK-FAST:           %[[VAL_5:.*]] = constant 1 : index
-// CHECK-FAST:           %[[VAL_6:.*]] = call @sparsePtrsI64(%[[VAL_0]], %[[VAL_5]]) : (!llvm.ptr<i8>, index) -> memref<?xindex>
-// CHECK-FAST:           %[[VAL_7:.*]] = call @sparseIndxsI64(%[[VAL_0]], %[[VAL_5]]) : (!llvm.ptr<i8>, index) -> memref<?xindex>
-// CHECK-FAST:           %[[VAL_8:.*]] = call @sparseValsF64(%[[VAL_0]]) : (!llvm.ptr<i8>) -> memref<?xf64>
+// CHECK-FAST:           %[[VAL_6:.*]] = call @sparsePointers64(%[[VAL_0]], %[[VAL_5]]) : (!llvm.ptr<i8>, index) -> memref<?xindex>
+// CHECK-FAST:           %[[VAL_7:.*]] = call @sparseIndices64(%[[VAL_0]], %[[VAL_5]]) : (!llvm.ptr<i8>, index) -> memref<?xindex>
+// CHECK-FAST:           %[[VAL_8:.*]] = call @sparseValuesF64(%[[VAL_0]]) : (!llvm.ptr<i8>) -> memref<?xf64>
 // CHECK-FAST:           scf.for %[[VAL_9:.*]] = %[[VAL_4]] to %[[VAL_3]] step %[[VAL_5]] {
 // CHECK-FAST:             %[[VAL_10:.*]] = load %[[VAL_6]]{{\[}}%[[VAL_9]]] : memref<?xindex>
 // CHECK-FAST:             %[[VAL_11:.*]] = addi %[[VAL_9]], %[[VAL_5]] : index
 // CHECK-FAST:             %[[VAL_12:.*]] = load %[[VAL_6]]{{\[}}%[[VAL_11]]] : memref<?xindex>
 // CHECK-FAST:             %[[VAL_13:.*]] = load %[[VAL_2]]{{\[}}%[[VAL_9]]] : memref<64xf64>
 // CHECK-FAST:             %[[VAL_14:.*]] = scf.for %[[VAL_15:.*]] = %[[VAL_10]] to %[[VAL_12]] step %[[VAL_5]] iter_args(%[[VAL_16:.*]] = %[[VAL_13]]) -> (f64) {
 // CHECK-FAST:               %[[VAL_17:.*]] = load %[[VAL_7]]{{\[}}%[[VAL_15]]] : memref<?xindex>
 // CHECK-FAST:               %[[VAL_18:.*]] = load %[[VAL_8]]{{\[}}%[[VAL_15]]] : memref<?xf64>

mlir/test/Dialect/Linalg/sparse_lower_calls.mlir

 // RUN: mlir-opt --test-sparsification="lower" %s | FileCheck %s
 
 !SparseTensor = type !llvm.ptr<i8>
 
 // CHECK-LABEL: func @sparse_pointers(
 //  CHECK-SAME: %[[A:.*]]: !llvm.ptr<i8>)
 //       CHECK: %[[C:.*]] = constant 1 : index
-//       CHECK: %[[T:.*]] = call @sparsePtrsI64(%[[A]], %[[C]]) : (!llvm.ptr<i8>, index) -> memref<?xindex>
+//       CHECK: %[[T:.*]] = call @sparsePointers64(%[[A]], %[[C]]) : (!llvm.ptr<i8>, index) -> memref<?xindex>
 //       CHECK: return %[[T]] : memref<?xindex>
 func @sparse_pointers(%arg0: !SparseTensor) -> memref<?xindex> {
   %a = linalg.sparse_tensor %arg0 : !SparseTensor to tensor<128xf64>
   %c = constant 1 : index
   %0 = linalg.sparse_pointers %a, %c : tensor<128xf64> to memref<?xindex>
   return %0 : memref<?xindex>
 }
 
+// CHECK-LABEL: func @sparse_pointers32(
+//  CHECK-SAME: %[[A:.*]]: !llvm.ptr<i8>)
+//       CHECK: %[[C:.*]] = constant 1 : index
+//       CHECK: %[[T:.*]] = call @sparsePointers32(%[[A]], %[[C]]) : (!llvm.ptr<i8>, index) -> memref<?xi32>
+//       CHECK: return %[[T]] : memref<?xi32>
+func @sparse_pointers32(%arg0: !SparseTensor) -> memref<?xi32> {
+  %a = linalg.sparse_tensor %arg0 : !SparseTensor to tensor<128xf64>
+  %c = constant 1 : index
+  %0 = linalg.sparse_pointers %a, %c : tensor<128xf64> to memref<?xi32>
+  return %0 : memref<?xi32>
+}
+
 // CHECK-LABEL: func @sparse_indices(
 //  CHECK-SAME: %[[A:.*]]: !llvm.ptr<i8>)
 //       CHECK: %[[C:.*]] = constant 1 : index
-//       CHECK: %[[T:.*]] = call @sparseIndxsI64(%[[A]], %[[C]]) : (!llvm.ptr<i8>, index) -> memref<?xindex>
+//       CHECK: %[[T:.*]] = call @sparseIndices64(%[[A]], %[[C]]) : (!llvm.ptr<i8>, index) -> memref<?xindex>
 //       CHECK: return %[[T]] : memref<?xindex>
 func @sparse_indices(%arg0: !SparseTensor) -> memref<?xindex> {
   %a = linalg.sparse_tensor %arg0 : !SparseTensor to tensor<128xf64>
   %c = constant 1 : index
   %0 = linalg.sparse_indices %a, %c : tensor<128xf64> to memref<?xindex>
   return %0 : memref<?xindex>
 }
 
-// CHECK-LABEL: func @sparse_values(
+// CHECK-LABEL: func @sparse_indices32(
 //  CHECK-SAME: %[[A:.*]]: !llvm.ptr<i8>)
-//       CHECK: %[[T:.*]] = call @sparseValsF64(%[[A]]) : (!llvm.ptr<i8>) -> memref<?xf64>
+//       CHECK: %[[C:.*]] = constant 1 : index
+//       CHECK: %[[T:.*]] = call @sparseIndices32(%[[A]], %[[C]]) : (!llvm.ptr<i8>, index) -> memref<?xi32>
+//       CHECK: return %[[T]] : memref<?xi32>
+func @sparse_indices32(%arg0: !SparseTensor) -> memref<?xi32> {
+  %a = linalg.sparse_tensor %arg0 : !SparseTensor to tensor<128xf64>
+  %c = constant 1 : index
+  %0 = linalg.sparse_indices %a, %c : tensor<128xf64> to memref<?xi32>
+  return %0 : memref<?xi32>
+}
+
+// CHECK-LABEL: func @sparse_valuesf64(
+//  CHECK-SAME: %[[A:.*]]: !llvm.ptr<i8>)
+//       CHECK: %[[T:.*]] = call @sparseValuesF64(%[[A]]) : (!llvm.ptr<i8>) -> memref<?xf64>
 //       CHECK: return %[[T]] : memref<?xf64>
-func @sparse_values(%arg0: !SparseTensor) -> memref<?xf64> {
+func @sparse_valuesf64(%arg0: !SparseTensor) -> memref<?xf64> {
   %a = linalg.sparse_tensor %arg0 : !SparseTensor to tensor<128xf64>
   %0 = linalg.sparse_values %a : tensor<128xf64> to memref<?xf64>
   return %0 : memref<?xf64>
 }
+
+// CHECK-LABEL: func @sparse_valuesf32(
+//  CHECK-SAME: %[[A:.*]]: !llvm.ptr<i8>)
+//       CHECK: %[[T:.*]] = call @sparseValuesF32(%[[A]]) : (!llvm.ptr<i8>) -> memref<?xf32>
+//       CHECK: return %[[T]] : memref<?xf32>
+func @sparse_valuesf32(%arg0: !SparseTensor) -> memref<?xf32> {
+  %a = linalg.sparse_tensor %arg0 : !SparseTensor to tensor<128xf32>
+  %0 = linalg.sparse_values %a : tensor<128xf32> to memref<?xf32>
+  return %0 : memref<?xf32>
+}