[mlir][sparse] Improving ExecutionEngine/SparseTensorUtils.h

[wrengr]

Summary

This change makes the public API of SparseTensorUtils.cpp explicit, whereas before the publicity of these functions was only implicit. Implicit publicity is sufficient for mlir-opt to generate calls to these functions, but it's not enough to enable C/C++ code to call them directly in the usual way (i.e., without going through codegen). Thus, leaving the publicity implicit prevents development of other tools (e.g., microbenchmarks).

In addition this change also marks the functions MLIR_CRUNNERUTILS_EXPORT, which is required by the JIT under certain configurations (albeit not for anything in our test suite).

[aartbik]

mlir/lib/ExecutionEngine/SparseTensorUtils.cpp

17–18

is this still needed now?

1482

looks like we miss the 8/64 case too

1564

strange, that usually is the case when a { or ( is not closed properly (as far as the much more context insensitive lint is concernced).
try to match closing brackets to see if you find something

[wrengr]

addressing comments

mlir/lib/ExecutionEngine/SparseTensorUtils.cpp

17–18

Nope. I'll also remove the <complex>, <cinttypes>, and <vector> includes, since they're also no longer necessary

1482

Bah, my random generator didn't catch that one ;) Also 32/64

I'll have to figure out a nice way to chain together x-macros to alleviate all the CASE and CASE_SECSAME situations. (Assuming I don't get things switched over to codegen first)

1564

After some further investigation, apparently the static keyword on the macro call for C32ABI somehow confuses clang-format and it assumes everything between the keyword and the next { or ; token must be part of a prototype. Luckily C++11 allows top-level extraneous semicolons (unlike earlier C++), so it's easy enough to resolve.

mlir/include/mlir/ExecutionEngine/SparseTensorUtils.h

 //===- SparseTensorUtils.h - Enums shared with the runtime ------*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
-// This header file defines several enums shared between
-// Transforms/SparseTensorConversion.cpp and ExecutionEngine/SparseUtils.cpp
+// This header file provides the enums and functions which comprise the
+// public API of `ExecutionEngine/SparseUtils.cpp`.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_EXECUTIONENGINE_SPARSETENSORUTILS_H_
 #define MLIR_EXECUTIONENGINE_SPARSETENSORUTILS_H_
 
+#include "mlir/ExecutionEngine/CRunnerUtils.h"
+
 #include <cinttypes>
+#include <complex>
+#include <vector>
 
 extern "C" {
 
+//===----------------------------------------------------------------------===//
+//
+// Typedefs and enums.  These are required to be public so that they
+// can be shared with `Transforms/SparseTensorConversion.cpp`, since
+// they define the arguments to the public functions declared later on.
+//
+// This section also defines x-macros <https://en.wikipedia.org/wiki/X_Macro>
+// so that we can generate variations of the public functions for each
+// supported primary- and/or overhead-type.
+//
+//===----------------------------------------------------------------------===//
+
 /// This type is used in the public API at all places where MLIR expects
 /// values with the built-in type "index". For now, we simply assume that
-/// type is 64-bit, but targets with different "index" bit widths should link
-/// with an alternatively built runtime support library.
+/// type is 64-bit, but targets with different "index" bit widths should
+/// link with an alternatively built runtime support library.
 // TODO: support such targets?
 using index_type = uint64_t;
 
 /// Encoding of overhead types (both pointer overhead and indices
 /// overhead), for "overloading" @newSparseTensor.
 enum class OverheadType : uint32_t {
   kIndex = 0,
   kU64 = 1,
   kU32 = 2,
   kU16 = 3,
   kU8 = 4
 };
 
+// This x-macro calls its argument on every overhead type which has
+// fixed-width.  It excludes `index_type` because that type is often
+// handled specially (e.g., by translating it into the architecture-dependent
+// equivalent fixed-width overhead type).
+#define FOREVERY_FIXED_O(DO)                                                   \
+  DO(64, uint64_t)                                                             \
+  DO(32, uint32_t)                                                             \
+  DO(16, uint16_t)                                                             \
+  DO(8, uint8_t)
+
+// This x-macro calls its argument on every overhead type, including
+// `index_type`.  Our naming convention uses an empty suffix for
+// `index_type`, so the missing first argument when we call `DO`
+// gets resolved to the empty token which can then be concatenated
+// as intended.  (This behavior is standard per C99 6.10.3/4 and
+// C++11 N3290 16.3/4; whereas in C++03 16.3/10 it was undefined behavior.)
+#define FOREVERY_O(DO)                                                         \
+  FOREVERY_FIXED_O(DO)                                                         \
+  DO(, index_type)
+
+// These are not just shorthands but indicate the particular
+// implementation used (e.g., as opposed to C99's `complex double`,
+// or MLIR's `ComplexType`).
+using complex64 = std::complex<double>;
+using complex32 = std::complex<float>;
+
 /// Encoding of the elemental type, for "overloading" @newSparseTensor.
 enum class PrimaryType : uint32_t {
   kF64 = 1,
   kF32 = 2,
   kI64 = 3,
   kI32 = 4,
   kI16 = 5,
   kI8 = 6,
   kC64 = 7,
   kC32 = 8
 };
 
+// This x-macro only specifies the non-complex `V` types, because the ABI
+// for complex types has compiler-/architecture-dependent details we need
+// to work around.  Namely, when a function takes a parameter of C/C++
+// type `complex32` (per se), then there is additional padding that causes
+// it not to match the LLVM type `!llvm.struct<(f32, f32)>`.  This only
+// happens with the `complex32` type itself, not with pointers/arrays
+// of complex values.  We also exclude `complex64` because it is in
+// principle susceptible to analogous ABI issues (even though we haven't
+// yet encountered them in practice).
+#define FOREVERY_SIMPLEX_V(DO)                                                 \
+  DO(F64, double)                                                              \
+  DO(F32, float)                                                               \
+  DO(I64, int64_t)                                                             \
+  DO(I32, int32_t)                                                             \
+  DO(I16, int16_t)                                                             \
+  DO(I8, int8_t)
+
+// This x-macro includes all `V` types, for when the aforementioned ABI
+// issues don't apply (e.g., because the functions take pointers/arrays).
+#define FOREVERY_V(DO)                                                         \
+  FOREVERY_SIMPLEX_V(DO)                                                       \
+  DO(C64, complex64)                                                           \
+  DO(C32, complex32)
+
 /// The actions performed by @newSparseTensor.
 enum class Action : uint32_t {
   kEmpty = 0,
   kFromFile = 1,
   kFromCOO = 2,
   kSparseToSparse = 3,
   kEmptyCOO = 4,
   kToCOO = 5,
   kToIterator = 6
 };
 
 /// This enum mimics `SparseTensorEncodingAttr::DimLevelType` for
 /// breaking dependency cycles.  `SparseTensorEncodingAttr::DimLevelType`
 /// is the source of truth and this enum should be kept consistent with it.
 enum class DimLevelType : uint8_t {
   kDense = 0,
   kCompressed = 1,
   kSingleton = 2
 };
 
+//===----------------------------------------------------------------------===//
+//
+// Public functions which operate on MLIR buffers (memrefs) to interact
+// with sparse tensors (which are only visible as opaque pointers externally).
+// Because these functions deal with memrefs, they should only be used
+// by MLIR compiler-generated code (or code similarly guaranteed to remain
+// in sync with MLIR; e.g., internal development tools like benchmarks).
+//
+// Where appropriate, we use macros to generate all variations of these
+// functions for each supported primary- and overhead-type.
+//
+//===----------------------------------------------------------------------===//
+
+/// The @newSparseTensor function for constructing a new sparse tensor.
+/// This is the "swiss army knife" method for materializing sparse
+/// tensors into the computation.  The types of the `ptr` argument and
+/// the result depend on the action, as explained in the following table
+/// (where "STS" means a sparse-tensor-storage object, and "COO" means
+/// a coordinate-scheme object).
+///
+/// Action:         `ptr`:          Returns:
+/// kEmpty          unused          STS, empty
+/// kEmptyCOO       unused          COO, empty
+/// kFromFile       char* filename  STS, read from the file
+/// kFromCOO        COO             STS, copied from the COO source
+/// kToCOO          STS             COO, copied from the STS source
+/// kSparseToSparse STS             STS, copied from the STS source
+/// kToIterator     STS             COO-Iterator, call @getNext to use
+MLIR_CRUNNERUTILS_EXPORT void *
+_mlir_ciface_newSparseTensor(StridedMemRefType<DimLevelType, 1> *aref, // NOLINT
+                             StridedMemRefType<index_type, 1> *sref,
+                             StridedMemRefType<index_type, 1> *pref,
+                             OverheadType ptrTp, OverheadType indTp,
+                             PrimaryType valTp, Action action, void *ptr);
+
+/// Tensor-storage method to obtain direct access to the values array.
+#define DECL_SPARSEVALUES(VNAME, V)                                            \
+  MLIR_CRUNNERUTILS_EXPORT void _mlir_ciface_sparseValues##VNAME(              \
+      StridedMemRefType<V, 1> *out, void *tensor);
+FOREVERY_V(DECL_SPARSEVALUES)
+#undef DECL_SPARSEVALUES
+
+/// Tensor-storage method to obtain direct access to the pointers array
+/// for the given dimension.
+#define DECL_SPARSEPOINTERS(PNAME, P)                                          \
+  MLIR_CRUNNERUTILS_EXPORT void _mlir_ciface_sparsePointers##PNAME(            \
+      StridedMemRefType<P, 1> *out, void *tensor, index_type d);
+FOREVERY_O(DECL_SPARSEPOINTERS)
+#undef DECL_SPARSEPOINTERS
+
+/// Tensor-storage method to obtain direct access to the indices array
+/// for the given dimension.
+#define DECL_SPARSEINDICES(INAME, I)                                           \
+  MLIR_CRUNNERUTILS_EXPORT void _mlir_ciface_sparseIndices##INAME(             \
+      StridedMemRefType<I, 1> *out, void *tensor, index_type d);
+FOREVERY_O(DECL_SPARSEINDICES)
+#undef DECL_SPARSEINDICES
+
+/// Coordinate-scheme method for adding a new element.
+#define DECL_ADDELT(VNAME, V)                                                  \
+  MLIR_CRUNNERUTILS_EXPORT void *_mlir_ciface_addElt##VNAME(                   \
+      void *coo, V value, StridedMemRefType<index_type, 1> *iref,              \
+      StridedMemRefType<index_type, 1> *pref);
+FOREVERY_SIMPLEX_V(DECL_ADDELT)
+DECL_ADDELT(C64, complex64)
+#undef DECL_ADDELT
+// Explicitly unpack the `complex32` into a pair of `float` arguments,
+// to work around ABI issues.
+// TODO: cleaner way to avoid ABI padding problem?
+MLIR_CRUNNERUTILS_EXPORT void *
+_mlir_ciface_addEltC32(void *coo, float r, float i,
+                       StridedMemRefType<index_type, 1> *iref,
+                       StridedMemRefType<index_type, 1> *pref);
+
+/// Coordinate-scheme method for getting the next element while iterating.
+#define DECL_GETNEXT(VNAME, V)                                                 \
+  MLIR_CRUNNERUTILS_EXPORT bool _mlir_ciface_getNext##VNAME(                   \
+      void *coo, StridedMemRefType<index_type, 1> *iref,                       \
+      StridedMemRefType<V, 0> *vref);
+FOREVERY_V(DECL_GETNEXT)
+#undef DECL_GETNEXT
+
+/// Tensor-storage method to insert elements in lexicographical index order.
+#define DECL_LEXINSERT(VNAME, V)                                               \
+  MLIR_CRUNNERUTILS_EXPORT void _mlir_ciface_lexInsert##VNAME(                 \
+      void *tensor, StridedMemRefType<index_type, 1> *cref, V val);
+FOREVERY_SIMPLEX_V(DECL_LEXINSERT)
+DECL_LEXINSERT(C64, complex64)
+#undef DECL_LEXINSERT
+// Explicitly unpack the `complex32` into a pair of `float` arguments,
+// to work around ABI issues.
+// TODO: cleaner way to avoid ABI padding problem?
+MLIR_CRUNNERUTILS_EXPORT void
+_mlir_ciface_lexInsertC32(void *tensor, StridedMemRefType<index_type, 1> *cref,
+                          float r, float i);
+
+/// Tensor-storage method to insert using expansion.
+#define DECL_EXPINSERT(VNAME, V)                                               \
+  MLIR_CRUNNERUTILS_EXPORT void _mlir_ciface_expInsert##VNAME(                 \
+      void *tensor, StridedMemRefType<index_type, 1> *cref,                    \
+      StridedMemRefType<V, 1> *vref, StridedMemRefType<bool, 1> *fref,         \
+      StridedMemRefType<index_type, 1> *aref, index_type count);
+FOREVERY_V(DECL_EXPINSERT)
+#undef DECL_EXPINSERT
+
+//===----------------------------------------------------------------------===//
+//
+// Public functions which accept only C-style data structures to interact
+// with sparse tensors (which are only visible as opaque pointers externally).
+// These functions can be used both by MLIR compiler-generated code
+// as well as by any external runtime that wants to interact with MLIR
+// compiler-generated code.
+//
+//===----------------------------------------------------------------------===//
+
+/// Tensor-storage method to get the size of the given dimension.
+MLIR_CRUNNERUTILS_EXPORT index_type sparseDimSize(void *tensor, index_type d);
+
+/// Tensor-storage method to finalize lexicographic insertions.
+MLIR_CRUNNERUTILS_EXPORT void endInsert(void *tensor);
+
+/// Coordinate-scheme method to write to file in extended FROSTT format.
+#define DECL_OUTSPARSETENSOR(VNAME, V)                                         \
+  MLIR_CRUNNERUTILS_EXPORT void outSparseTensor##VNAME(void *coo, void *dest,  \
+                                                       bool sort);
+FOREVERY_V(DECL_OUTSPARSETENSOR)
+#undef DECL_OUTSPARSETENSOR
+
+/// Releases the memory for the tensor-storage object.
+void delSparseTensor(void *tensor);
+
+/// Releases the memory for the coordinate-scheme object.
+#define DECL_DELCOO(VNAME, V)                                                  \
+  MLIR_CRUNNERUTILS_EXPORT void delSparseTensorCOO##VNAME(void *coo);
+FOREVERY_V(DECL_DELCOO)
+#undef DECL_DELCOO
+
+/// Helper function to read a sparse tensor filename from the environment,
+/// defined with the naming convention ${TENSOR0}, ${TENSOR1}, etc.
+MLIR_CRUNNERUTILS_EXPORT char *getTensorFilename(index_type id);
+
+/// Initializes sparse tensor from a COO-flavored format expressed using
+/// C-style data structures.  The expected parameters are:
+///
+///   rank:    rank of tensor
+///   nse:     number of specified elements (usually the nonzeros)
+///   shape:   array with dimension size for each rank
+///   values:  a "nse" array with values for all specified elements
+///   indices: a flat "nse * rank" array with indices for all specified elements
+///   perm:    the permutation of the dimensions in the storage
+///   sparse:  the sparsity for the dimensions
+///
+/// For example, the sparse matrix
+///     | 1.0 0.0 0.0 |
+///     | 0.0 5.0 3.0 |
+/// can be passed as
+///      rank    = 2
+///      nse     = 3
+///      shape   = [2, 3]
+///      values  = [1.0, 5.0, 3.0]
+///      indices = [ 0, 0,  1, 1,  1, 2]
+#define DECL_CONVERTTOMLIRSPARSETENSOR(VNAME, V)                               \
+  MLIR_CRUNNERUTILS_EXPORT void *convertToMLIRSparseTensor##VNAME(             \
+      uint64_t rank, uint64_t nse, uint64_t *shape, V *values,                 \
+      uint64_t *indices, uint64_t *perm, uint8_t *sparse);
+FOREVERY_V(DECL_CONVERTTOMLIRSPARSETENSOR)
+#undef DECL_CONVERTTOMLIRSPARSETENSOR
+
+/// Converts a sparse tensor to COO-flavored format expressed using
+/// C-style data structures.  The expected output parameters are pointers
+/// for these values:
+///
+///   rank:    rank of tensor
+///   nse:     number of specified elements (usually the nonzeros)
+///   shape:   array with dimension size for each rank
+///   values:  a "nse" array with values for all specified elements
+///   indices: a flat "nse * rank" array with indices for all specified elements
+///
+/// The input is a pointer to `SparseTensorStorage<P, I, V>`, typically
+/// returned from `convertToMLIRSparseTensor`.
+#define DECL_CONVERTFROMMLIRSPARSETENSOR(VNAME, V)                             \
+  MLIR_CRUNNERUTILS_EXPORT void convertFromMLIRSparseTensor##VNAME(            \
+      void *tensor, uint64_t *pRank, uint64_t *pNse, uint64_t **pShape,        \
+      V **pValues, uint64_t **pIndices);
+FOREVERY_V(DECL_CONVERTFROMMLIRSPARSETENSOR)
+#undef DECL_CONVERTFROMMLIRSPARSETENSOR
+
 } // extern "C"
 
 #endif // MLIR_EXECUTIONENGINE_SPARSETENSORUTILS_H_

mlir/lib/ExecutionEngine/SparseTensorUtils.cpp

 //===- SparseTensorUtils.cpp - Sparse Tensor Utils for MLIR execution -----===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This file implements a light-weight runtime support library that is useful
 // for sparse tensor manipulations. The functionality provided in this library
 // is meant to simplify benchmarking, testing, and debugging MLIR code that
 // operates on sparse tensors. The provided functionality is **not** part
 // of core MLIR, however.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/ExecutionEngine/SparseTensorUtils.h"
-#include "mlir/ExecutionEngine/CRunnerUtils.h"
 
 #ifdef MLIR_CRUNNERUTILS_DEFINE_FUNCTIONS
 
 #include <algorithm>
 #include <cassert>
-#include <complex>
 #include <cctype>
-#include <cinttypes>
 #include <cstdio>
 #include <cstdlib>
 #include <cstring>
 #include <fstream>
 #include <functional>
 #include <iostream>
 #include <limits>
 #include <numeric>
-#include <vector>
-
-using complex64 = std::complex<double>;
-using complex32 = std::complex<float>;
 
 //===----------------------------------------------------------------------===//
 //
 // 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

 private:
   const std::vector<uint64_t> dimSizes; // per-dimension sizes
   std::vector<Element<V>> elements;     // all COO elements
   std::vector<uint64_t> indices;        // shared index pool
   bool iteratorLocked = false;
   unsigned iteratorPos = 0;
 };
 
-// See <https://en.wikipedia.org/wiki/X_Macro>
-//
-// `FOREVERY_SIMPLEX_V` only specifies the non-complex `V` types, because
-// the ABI for complex types has compiler/architecture dependent complexities
-// we need to work around.  Namely, when a function takes a parameter of
-// C/C++ type `complex32` (per se), then there is additional padding that
-// causes it not to match the LLVM type `!llvm.struct<(f32, f32)>`.  This
-// only happens with the `complex32` type itself, not with pointers/arrays
-// of complex values.  So far `complex64` doesn't exhibit this ABI
-// incompatibility, but we exclude it anyways just to be safe.
-#define FOREVERY_SIMPLEX_V(DO)                                                 \
-  DO(F64, double)                                                              \
-  DO(F32, float)                                                               \
-  DO(I64, int64_t)                                                             \
-  DO(I32, int32_t)                                                             \
-  DO(I16, int16_t)                                                             \
-  DO(I8, int8_t)
-
-#define FOREVERY_V(DO)                                                         \
-  FOREVERY_SIMPLEX_V(DO)                                                       \
-  DO(C64, complex64)                                                           \
-  DO(C32, complex32)
-
-// This x-macro calls its argument on every overhead type which has
-// fixed-width.  It excludes `index_type` because that type is often
-// handled specially (e.g., by translating it into the architecture-dependent
-// equivalent fixed-width overhead type).
-#define FOREVERY_FIXED_O(DO)                                                   \
-  DO(64, uint64_t)                                                             \
-  DO(32, uint32_t)                                                             \
-  DO(16, uint16_t)                                                             \
-  DO(8, uint8_t)
-
-// This x-macro calls its argument on every overhead type, including
-// `index_type`.  Our naming convention uses an empty suffix for
-// `index_type`, so the missing first argument when we call `DO`
-// gets resolved to the empty token which can then be concatenated
-// as intended.  (This behavior is standard per C99 6.10.3/4 and
-// C++11 N3290 16.3/4; whereas in C++03 16.3/10 it was undefined behavior.)
-#define FOREVERY_O(DO)                                                         \
-  FOREVERY_FIXED_O(DO)                                                         \
-  DO(, index_type)
-
 // Forward.
 template <typename V>
 class SparseTensorEnumeratorBase;
 
 // Helper macro for generating error messages when some
 // `SparseTensorStorage<P,I,V>` is cast to `SparseTensorStorageBase`
 // and then the wrong "partial method specialization" is called.
 #define FATAL_PIV(NAME) FATAL("<P,I,V> type mismatch for: " #NAME);
for (uint64_t k = 0; k < nnz; k++) {
     if (isSymmetric && indices[0] != indices[1])
       tensor->add({indices[1], indices[0]}, value);
   }
   // Close the file and return tensor.
   fclose(file);
   return tensor;
 }
 
-/// Writes the sparse tensor to extended FROSTT format.
+/// Writes the sparse tensor to `dest` in extended FROSTT format.
 template <typename V>
 static void outSparseTensor(void *tensor, void *dest, bool sort) {
   assert(tensor && dest);
   auto coo = static_cast<SparseTensorCOO<V> *>(tensor);
   if (sort)
     coo->sort();
   char *filename = static_cast<char *>(dest);
   auto &dimSizes = coo->getDimSizes();
static void fromMLIRSparseTensor(void *tensor, uint64_t *pRank, uint64_t *pNse,
   delete coo;
   *pRank = rank;
   *pNse = nse;
   *pShape = shape;
   *pValues = values;
   *pIndices = indices;
 }
 
-} // namespace
+} // anonymous namespace
 
 extern "C" {
 
 //===----------------------------------------------------------------------===//
 //
-// Public API with methods that operate on MLIR buffers (memrefs) to interact
-// with sparse tensors, which are only visible as opaque pointers externally.
-// These methods should be used exclusively by MLIR compiler-generated code.
-//
-// Some macro magic is used to generate implementations for all required type
-// combinations that can be called from MLIR compiler-generated code.
+// Public functions which operate on MLIR buffers (memrefs) to interact
+// with sparse tensors (which are only visible as opaque pointers externally).
 //
 //===----------------------------------------------------------------------===//
 
 #define CASE(p, i, v, P, I, V)                                                 \
   if (ptrTp == (p) && indTp == (i) && valTp == (v)) {                          \
     SparseTensorCOO<V> *coo = nullptr;                                         \
     if (action <= Action::kFromCOO) {                                          \
       if (action == Action::kFromFile) {                                       \

 #define CASE_SECSAME(p, v, P, V) CASE(p, p, v, P, P, V)
 
 // Assume index_type is in fact uint64_t, so that _mlir_ciface_newSparseTensor
 // can safely rewrite kIndex to kU64.  We make this assertion to guarantee
 // that this file cannot get out of sync with its header.
 static_assert(std::is_same<index_type, uint64_t>::value,
               "Expected index_type == uint64_t");
 
-/// Constructs a new sparse tensor. This is the "swiss army knife"
-/// method for materializing sparse tensors into the computation.
-///
-/// Action:
-/// kEmpty = returns empty storage to fill later
-/// kFromFile = returns storage, where ptr contains filename to read
-/// kFromCOO = returns storage, where ptr contains coordinate scheme to assign
-/// kEmptyCOO = returns empty coordinate scheme to fill and use with kFromCOO
-/// kToCOO = returns coordinate scheme from storage in ptr to use with kFromCOO
-/// kToIterator = returns iterator from storage in ptr (call getNext() to use)
 void *
 _mlir_ciface_newSparseTensor(StridedMemRefType<DimLevelType, 1> *aref, // NOLINT
                              StridedMemRefType<index_type, 1> *sref,
                              StridedMemRefType<index_type, 1> *pref,
                              OverheadType ptrTp, OverheadType indTp,
                              PrimaryType valTp, Action action, void *ptr) {
   assert(aref && sref && pref);
   assert(aref->strides[0] == 1 && sref->strides[0] == 1 &&
_mlir_ciface_newSparseTensor(StridedMemRefType<DimLevelType, 1> *aref, // NOLINT
   CASE(OverheadType::kU8, OverheadType::kU8, PrimaryType::kF32, uint8_t,
        uint8_t, float);
 
   // Integral matrices with both overheads of the same type.
   CASE_SECSAME(OverheadType::kU64, PrimaryType::kI64, uint64_t, int64_t);
   CASE_SECSAME(OverheadType::kU64, PrimaryType::kI32, uint64_t, int32_t);
   CASE_SECSAME(OverheadType::kU64, PrimaryType::kI16, uint64_t, int16_t);
   CASE_SECSAME(OverheadType::kU64, PrimaryType::kI8, uint64_t, int8_t);
+  CASE_SECSAME(OverheadType::kU32, PrimaryType::kI64, uint32_t, int64_t);
   CASE_SECSAME(OverheadType::kU32, PrimaryType::kI32, uint32_t, int32_t);
   CASE_SECSAME(OverheadType::kU32, PrimaryType::kI16, uint32_t, int16_t);
   CASE_SECSAME(OverheadType::kU32, PrimaryType::kI8, uint32_t, int8_t);
+  CASE_SECSAME(OverheadType::kU16, PrimaryType::kI64, uint16_t, int64_t);
   CASE_SECSAME(OverheadType::kU16, PrimaryType::kI32, uint16_t, int32_t);
   CASE_SECSAME(OverheadType::kU16, PrimaryType::kI16, uint16_t, int16_t);
   CASE_SECSAME(OverheadType::kU16, PrimaryType::kI8, uint16_t, int8_t);
+  CASE_SECSAME(OverheadType::kU8, PrimaryType::kI64, uint8_t, int64_t);
   CASE_SECSAME(OverheadType::kU8, PrimaryType::kI32, uint8_t, int32_t);
   CASE_SECSAME(OverheadType::kU8, PrimaryType::kI16, uint8_t, int16_t);
   CASE_SECSAME(OverheadType::kU8, PrimaryType::kI8, uint8_t, int8_t);
 
   // Complex matrices with wide overhead.
   CASE_SECSAME(OverheadType::kU64, PrimaryType::kC64, uint64_t, complex64);
   CASE_SECSAME(OverheadType::kU64, PrimaryType::kC32, uint64_t, complex32);
 
   // Unsupported case (add above if needed).
   // TODO: better pretty-printing of enum values!
   FATAL("unsupported combination of types: <P=%d, I=%d, V=%d>\n",
         static_cast<int>(ptrTp), static_cast<int>(indTp),
         static_cast<int>(valTp));
 }
 #undef CASE
 #undef CASE_SECSAME
 
-/// Methods that provide direct access to values.
 #define IMPL_SPARSEVALUES(VNAME, V)                                            \
   void _mlir_ciface_sparseValues##VNAME(StridedMemRefType<V, 1> *ref,          \
                                         void *tensor) {                        \
     assert(ref &&tensor);                                                      \
     std::vector<V> *v;                                                         \
     static_cast<SparseTensorStorageBase *>(tensor)->getValues(&v);             \
     ref->basePtr = ref->data = v->data();                                      \
     ref->offset = 0;                                                           \
void _mlir_ciface_##NAME(StridedMemRefType<TYPE, 1> *ref, void *tensor,      \
     assert(ref &&tensor);                                                      \
     std::vector<TYPE> *v;                                                      \
     static_cast<SparseTensorStorageBase *>(tensor)->LIB(&v, d);                \
     ref->basePtr = ref->data = v->data();                                      \
     ref->offset = 0;                                                           \
     ref->sizes[0] = v->size();                                                 \
     ref->strides[0] = 1;                                                       \
   }
-/// Methods that provide direct access to pointers.
 #define IMPL_SPARSEPOINTERS(PNAME, P)                                          \
   IMPL_GETOVERHEAD(sparsePointers##PNAME, P, getPointers)
 FOREVERY_O(IMPL_SPARSEPOINTERS)
 #undef IMPL_SPARSEPOINTERS
 
-/// Methods that provide direct access to indices.
 #define IMPL_SPARSEINDICES(INAME, I)                                           \
   IMPL_GETOVERHEAD(sparseIndices##INAME, I, getIndices)
 FOREVERY_O(IMPL_SPARSEINDICES)
 #undef IMPL_SPARSEINDICES
 #undef IMPL_GETOVERHEAD
 
-/// Helper to add value to coordinate scheme, one per value type.
 #define IMPL_ADDELT(VNAME, V)                                                  \
   void *_mlir_ciface_addElt##VNAME(void *coo, V value,                         \
                                    StridedMemRefType<index_type, 1> *iref,     \
                                    StridedMemRefType<index_type, 1> *pref) {   \
     assert(coo &&iref &&pref);                                                 \
     assert(iref->strides[0] == 1 && pref->strides[0] == 1);                    \
     assert(iref->sizes[0] == pref->sizes[0]);                                  \
     const index_type *indx = iref->data + iref->offset;                        \
     const index_type *perm = pref->data + pref->offset;                        \
     uint64_t isize = iref->sizes[0];                                           \
     std::vector<index_type> indices(isize);                                    \
     for (uint64_t r = 0; r < isize; r++)                                       \
       indices[perm[r]] = indx[r];                                              \
     static_cast<SparseTensorCOO<V> *>(coo)->add(indices, value);               \
     return coo;                                                                \
   }
 FOREVERY_SIMPLEX_V(IMPL_ADDELT)
-// `complex64` apparently doesn't encounter any ABI issues (yet).
 IMPL_ADDELT(C64, complex64)
-// TODO: cleaner way to avoid ABI padding problem?
-IMPL_ADDELT(C32ABI, complex32)
+// Marked static because it's not part of the public API.
+// (The `static` keyword confuses clang-format, so the extraneous trailing
+// semicolon is required to teach clang-format not to indent the prototype
+// of `_mlir_ciface_addEltC32` below.)
+static IMPL_ADDELT(C32ABI, complex32);
+#undef IMPL_ADDELT
 void *_mlir_ciface_addEltC32(void *coo, float r, float i,
                              StridedMemRefType<index_type, 1> *iref,
                              StridedMemRefType<index_type, 1> *pref) {
   return _mlir_ciface_addEltC32ABI(coo, complex32(r, i), iref, pref);
 }
-#undef IMPL_ADDELT
 
-/// Helper to enumerate elements of coordinate scheme, one per value type.
 #define IMPL_GETNEXT(VNAME, V)                                                 \
   bool _mlir_ciface_getNext##VNAME(void *coo,                                  \
                                    StridedMemRefType<index_type, 1> *iref,     \
                                    StridedMemRefType<V, 0> *vref) {            \
     assert(coo &&iref &&vref);                                                 \
     assert(iref->strides[0] == 1);                                             \
     index_type *indx = iref->data + iref->offset;                              \
     V *value = vref->data + vref->offset;                                      \
     const uint64_t isize = iref->sizes[0];                                     \
     const Element<V> *elem =                                                   \
         static_cast<SparseTensorCOO<V> *>(coo)->getNext();                     \
     if (elem == nullptr)                                                       \
       return false;                                                            \
     for (uint64_t r = 0; r < isize; r++)                                       \
       indx[r] = elem->indices[r];                                              \
     *value = elem->value;                                                      \
     return true;                                                               \
   }
 FOREVERY_V(IMPL_GETNEXT)
 #undef IMPL_GETNEXT
 
-/// Insert elements in lexicographical index order, one per value type.
 #define IMPL_LEXINSERT(VNAME, V)                                               \
   void _mlir_ciface_lexInsert##VNAME(                                          \
       void *tensor, StridedMemRefType<index_type, 1> *cref, V val) {           \
     assert(tensor &&cref);                                                     \
     assert(cref->strides[0] == 1);                                             \
     index_type *cursor = cref->data + cref->offset;                            \
     assert(cursor);                                                            \
     static_cast<SparseTensorStorageBase *>(tensor)->lexInsert(cursor, val);    \
   }
 FOREVERY_SIMPLEX_V(IMPL_LEXINSERT)
-// `complex64` apparently doesn't encounter any ABI issues (yet).
 IMPL_LEXINSERT(C64, complex64)
-// TODO: cleaner way to avoid ABI padding problem?
-IMPL_LEXINSERT(C32ABI, complex32)
+// Marked static because it's not part of the public API.
+// (The `static` keyword confuses clang-format, so the extraneous trailing
+// semicolon is required to teach clang-format not to indent the prototype
+// of `_mlir_ciface_lexInsertC32` below.)
+static IMPL_LEXINSERT(C32ABI, complex32);
+#undef IMPL_LEXINSERT
 void _mlir_ciface_lexInsertC32(void *tensor,
                                StridedMemRefType<index_type, 1> *cref, float r,
                                float i) {
   _mlir_ciface_lexInsertC32ABI(tensor, cref, complex32(r, i));
 }
-#undef IMPL_LEXINSERT
 
-/// Insert using expansion, one per value type.
 #define IMPL_EXPINSERT(VNAME, V)                                               \
   void _mlir_ciface_expInsert##VNAME(                                          \
       void *tensor, StridedMemRefType<index_type, 1> *cref,                    \
       StridedMemRefType<V, 1> *vref, StridedMemRefType<bool, 1> *fref,         \
       StridedMemRefType<index_type, 1> *aref, index_type count) {              \
     assert(tensor &&cref &&vref &&fref &&aref);                                \
     assert(cref->strides[0] == 1);                                             \
     assert(vref->strides[0] == 1);                                             \
     assert(fref->strides[0] == 1);                                             \
     assert(aref->strides[0] == 1);                                             \
     assert(vref->sizes[0] == fref->sizes[0]);                                  \
     index_type *cursor = cref->data + cref->offset;                            \
     V *values = vref->data + vref->offset;                                     \
     bool *filled = fref->data + fref->offset;                                  \
     index_type *added = aref->data + aref->offset;                             \
     static_cast<SparseTensorStorageBase *>(tensor)->expInsert(                 \
         cursor, values, filled, added, count);                                 \
   }
 FOREVERY_V(IMPL_EXPINSERT)
 #undef IMPL_EXPINSERT
 
-/// Output a sparse tensor, one per value type.
+// TODO: move this into the next section, since it doesn't depend on memrefs.
 #define IMPL_OUTSPARSETENSOR(VNAME, V)                                         \
   void outSparseTensor##VNAME(void *coo, void *dest, bool sort) {              \
     return outSparseTensor<V>(coo, dest, sort);                                \
   }
 FOREVERY_V(IMPL_OUTSPARSETENSOR)
 #undef IMPL_OUTSPARSETENSOR
 
 //===----------------------------------------------------------------------===//
 //
-// Public API with methods that accept C-style data structures to interact
-// with sparse tensors, which are only visible as opaque pointers externally.
-// These methods can be used both by MLIR compiler-generated code as well as by
-// an external runtime that wants to interact with MLIR compiler-generated code.
+// Public functions which accept only C-style data structures to interact
+// with sparse tensors (which are only visible as opaque pointers externally).
 //
 //===----------------------------------------------------------------------===//
 
-/// Helper method to read a sparse tensor filename from the environment,
-/// defined with the naming convention ${TENSOR0}, ${TENSOR1}, etc.
+// TODO: move this lower down (after `delSparseTensorCOO`) since it's
+// independent of our sparse-tensor formats.
 char *getTensorFilename(index_type id) {
   char var[80];
   sprintf(var, "TENSOR%" PRIu64, id);
   char *env = getenv(var);
   if (!env)
     FATAL("Environment variable %s is not set\n", var);
   return env;
 }
 
-/// Returns size of sparse tensor in given dimension.
 index_type sparseDimSize(void *tensor, index_type d) {
   return static_cast<SparseTensorStorageBase *>(tensor)->getDimSize(d);
 }
 
-/// Finalizes lexicographic insertions.
 void endInsert(void *tensor) {
   return static_cast<SparseTensorStorageBase *>(tensor)->endInsert();
 }
 
-/// Releases sparse tensor storage.
 void delSparseTensor(void *tensor) {
   delete static_cast<SparseTensorStorageBase *>(tensor);
 }
 
-/// Releases sparse tensor coordinate scheme.
 #define IMPL_DELCOO(VNAME, V)                                                  \
   void delSparseTensorCOO##VNAME(void *coo) {                                  \
     delete static_cast<SparseTensorCOO<V> *>(coo);                             \
   }
 FOREVERY_V(IMPL_DELCOO)
 #undef IMPL_DELCOO
 
-/// Initializes sparse tensor from a COO-flavored format expressed using C-style
-/// data structures. The expected parameters are:
-///
-///   rank:    rank of tensor
-///   nse:     number of specified elements (usually the nonzeros)
-///   shape:   array with dimension size for each rank
-///   values:  a "nse" array with values for all specified elements
-///   indices: a flat "nse x rank" array with indices for all specified elements
-///   perm:    the permutation of the dimensions in the storage
-///   sparse:  the sparsity for the dimensions
-///
-/// For example, the sparse matrix
-///     | 1.0 0.0 0.0 |
-///     | 0.0 5.0 3.0 |
-/// can be passed as
-///      rank    = 2
-///      nse     = 3
-///      shape   = [2, 3]
-///      values  = [1.0, 5.0, 3.0]
-///      indices = [ 0, 0,  1, 1,  1, 2]
-//
 // TODO: generalize beyond 64-bit indices.
-//
 #define IMPL_CONVERTTOMLIRSPARSETENSOR(VNAME, V)                               \
   void *convertToMLIRSparseTensor##VNAME(                                      \
       uint64_t rank, uint64_t nse, uint64_t *shape, V *values,                 \
       uint64_t *indices, uint64_t *perm, uint8_t *sparse) {                    \
     return toMLIRSparseTensor<V>(rank, nse, shape, values, indices, perm,      \
                                  sparse);                                      \
   }
 FOREVERY_V(IMPL_CONVERTTOMLIRSPARSETENSOR)
 #undef IMPL_CONVERTTOMLIRSPARSETENSOR
 
-/// Converts a sparse tensor to COO-flavored format expressed using C-style
-/// data structures. The expected output parameters are pointers for these
-/// values:
-///
-///   rank:    rank of tensor
-///   nse:     number of specified elements (usually the nonzeros)
-///   shape:   array with dimension size for each rank
-///   values:  a "nse" array with values for all specified elements
-///   indices: a flat "nse x rank" array with indices for all specified elements
-///
-/// The input is a pointer to SparseTensorStorage<P, I, V>, typically returned
-/// from convertToMLIRSparseTensor.
-///
-//  TODO: Currently, values are copied from SparseTensorStorage to
-//  SparseTensorCOO, then to the output. We may want to reduce the number of
-//  copies.
+// TODO: Currently, values are copied from SparseTensorStorage to
+// SparseTensorCOO, then to the output.  We may want to reduce the number
+// of copies.
 //
 // TODO: generalize beyond 64-bit indices, no dim ordering, all dimensions
 // compressed
-//
 #define IMPL_CONVERTFROMMLIRSPARSETENSOR(VNAME, V)                             \
   void convertFromMLIRSparseTensor##VNAME(void *tensor, uint64_t *pRank,       \
                                           uint64_t *pNse, uint64_t **pShape,   \
                                           V **pValues, uint64_t **pIndices) {  \
     fromMLIRSparseTensor<V>(tensor, pRank, pNse, pShape, pValues, pIndices);   \
   }
 FOREVERY_V(IMPL_CONVERTFROMMLIRSPARSETENSOR)
 #undef IMPL_CONVERTFROMMLIRSPARSETENSOR
 
 } // extern "C"
 
 #endif // MLIR_CRUNNERUTILS_DEFINE_FUNCTIONS

utils/bazel/llvm-project-overlay/mlir/BUILD.bazel

deps = [
         ":SCFDialect",
         ":SCFTransforms",
         ":SparseTensor",
         ":SparseTensorPassIncGen",
         ":SparseTensorUtils",
         ":TensorDialect",
         ":Transforms",
         ":VectorOps",
+        ":mlir_c_runner_utils",
         "//llvm:Support",
     ],
 )
 
 cc_library(
     name = "SparseTensorPipelines",
     srcs = glob(["lib/Dialect/SparseTensor/Pipelines/*.cpp"]),
     hdrs = ["include/mlir/Dialect/SparseTensor/Pipelines/Passes.h"],