[mlir][sparse] cleanup ABI issues in C interface with memrefs

[aartbik]

Summary

This change adds automatic wrapper functoins with emit_c_interface
to all methods in the sparse support library that deal with MEMREFs.
The wrappers will take care of passing MEMREFs by value internally
and by pointer externally, thereby avoiding ABI issues across platforms.

[aartbik]

update

[mehdi_amini]

That was fast! Thanks :)

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

104

(naming convention: camelCase)

180

(same below)

[aartbik]

fix clang tidy issue

[aartbik]

addressed Mehdi's comments

[bondhugula]

mlir/lib/ExecutionEngine/SparseUtils.cpp

114

Unrelated side question: shouldn't this file be named SparseTensorUtils.cpp?

[aartbik]

mlir/lib/ExecutionEngine/SparseUtils.cpp

114

Yeah, I see this as a "sparse support library" but a "sparse tensor support library" would make sense too. Happy to rename (later). We will probably do some other cleanup too, such as adding a header.

[aartbik]

NOLINT on _mlir_ciface_XXX name

[aartbik]

NOLINT on right line this time ;-)

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

getDimLevelTypeEncoding(SparseTensorEncodingAttr::DimLevelType dlt) {
   case SparseTensorEncodingAttr::DimLevelType::Singleton:
     return 2;
   }
   llvm_unreachable("Unknown SparseTensorEncodingAttr::DimLevelType");
 }
 
 /// Returns integers of given width and values as a constant tensor.
 /// We cast the static shape into a dynamic shape to ensure that the
-/// method signature remains uniform accross different tensor dimensions.
+/// method signature remains uniform across different tensor dimensions.
 static Value getTensor(ConversionPatternRewriter &rewriter, unsigned width,
                        Location loc, ArrayRef<APInt> values) {
   Type etp = rewriter.getIntegerType(width);
   unsigned sz = values.size();
   RankedTensorType tt1 = RankedTensorType::get({sz}, etp);
   RankedTensorType tt2 = RankedTensorType::get({ShapedType::kDynamicSize}, etp);
   auto elts =
       rewriter.create<ConstantOp>(loc, DenseElementsAttr::get(tt1, values));
   return rewriter.create<tensor::CastOp>(loc, tt2, elts);
 }
 
-/// Returns function reference (first hit also inserts into module).
+/// Returns a function reference (first hit also inserts into module). Sets
+/// the "_emit_c_interface" on the function declaration when requested,
+/// so that LLVM lowering generates a wrapper function that takes care
+/// of ABI complications with passing in and returning MemRefs to C functions.
 static FlatSymbolRefAttr getFunc(Operation *op, StringRef name, Type resultType,
-                                 ValueRange operands) {
+                                 ValueRange operands,
+                                 bool emitCInterface = false) {
   MLIRContext *context = op->getContext();
   auto module = op->getParentOfType<ModuleOp>();
   auto result = SymbolRefAttr::get(context, name);
   auto func = module.lookupSymbol<FuncOp>(result.getAttr());
   if (!func) {
     OpBuilder moduleBuilder(module.getBodyRegion());
-    moduleBuilder
-        .create<FuncOp>(
-            op->getLoc(), name,
-            FunctionType::get(context, operands.getTypes(), resultType))
-        .setPrivate();
+    func = moduleBuilder.create<FuncOp>(
+        FunctionType::get(context, operands.getTypes(), resultType));
+    func.setPrivate();
+    if (emitCInterface)
+      func->setAttr("llvm.emit_c_interface", UnitAttr::get(context));
   }
   return result;
 }
 
 /// Generates a call into the "swiss army knife" method of the sparse runtime
 /// support library for materializing sparse tensors into the computation. The
 /// method returns the call value and assigns the permutation to 'perm'.
 static Value genNewCall(ConversionPatternRewriter &rewriter, Operation *op,
static Value genNewCall(ConversionPatternRewriter &rewriter, Operation *op,
   Type pTp = LLVM::LLVMPointerType::get(IntegerType::get(op->getContext(), 8));
   if (!ptr)
     ptr = rewriter.create<LLVM::NullOp>(loc, pTp);
   params.push_back(
       rewriter.create<ConstantOp>(loc, rewriter.getI32IntegerAttr(action)));
   params.push_back(ptr);
   // Generate the call to create new tensor.
   StringRef name = "newSparseTensor";
-  auto call =
-      rewriter.create<CallOp>(loc, pTp, getFunc(op, name, pTp, params), params);
+  auto call = rewriter.create<CallOp>(
+      loc, pTp, getFunc(op, name, pTp, params, /*emitCInterface=*/true),
+      params);
   return call.getResult(0);
 }
 
 /// Generates a call that adds one element to a coordinate scheme.
 static void genAddEltCall(ConversionPatternRewriter &rewriter, Operation *op,
                           Value ptr, Value tensor, Value ind, Value perm,
                           ValueRange ivs) {
   Location loc = op->getLoc();
for (auto iv : ivs) {
     rewriter.create<memref::StoreOp>(loc, iv, ind, idx);
   }
   SmallVector<Value, 8> params;
   params.push_back(ptr);
   params.push_back(val);
   params.push_back(ind);
   params.push_back(perm);
   Type pTp = LLVM::LLVMPointerType::get(IntegerType::get(op->getContext(), 8));
-  rewriter.create<CallOp>(loc, pTp, getFunc(op, name, pTp, params), params);
+  rewriter.create<CallOp>(
+      loc, pTp, getFunc(op, name, pTp, params, /*emitCInterface=*/true),
+      params);
 }
 
 //===----------------------------------------------------------------------===//
 // Conversion rules.
 //===----------------------------------------------------------------------===//
 
 /// Sparse conversion rule for returns.
 class SparseReturnConverter : public OpConversionPattern<ReturnOp> {
else if (eltType.isInteger(32))
       name = "sparsePointers32";
     else if (eltType.isInteger(16))
       name = "sparsePointers16";
     else if (eltType.isInteger(8))
       name = "sparsePointers8";
     else
       return failure();
     rewriter.replaceOpWithNewOp<CallOp>(
-        op, resType, getFunc(op, name, resType, operands), operands);
+        op, resType,
+        getFunc(op, name, resType, operands, /*emitCInterface=*/true),
+        operands);
     return success();
   }
 };
 
 /// Sparse conversion rule for index accesses.
 class SparseTensorToIndicesConverter : public OpConversionPattern<ToIndicesOp> {
 public:
   using OpConversionPattern::OpConversionPattern;
else if (eltType.isInteger(32))
       name = "sparseIndices32";
     else if (eltType.isInteger(16))
       name = "sparseIndices16";
     else if (eltType.isInteger(8))
       name = "sparseIndices8";
     else
       return failure();
     rewriter.replaceOpWithNewOp<CallOp>(
-        op, resType, getFunc(op, name, resType, operands), operands);
+        op, resType,
+        getFunc(op, name, resType, operands, /*emitCInterface=*/true),
+        operands);
     return success();
   }
 };
 
 /// Sparse conversion rule for value accesses.
 class SparseTensorToValuesConverter : public OpConversionPattern<ToValuesOp> {
 public:
   using OpConversionPattern::OpConversionPattern;
else if (eltType.isInteger(32))
       name = "sparseValuesI32";
     else if (eltType.isInteger(16))
       name = "sparseValuesI16";
     else if (eltType.isInteger(8))
       name = "sparseValuesI8";
     else
       return failure();
     rewriter.replaceOpWithNewOp<CallOp>(
-        op, resType, getFunc(op, name, resType, operands), operands);
+        op, resType,
+        getFunc(op, name, resType, operands, /*emitCInterface=*/true),
+        operands);
     return success();
   }
 };
 
 /// Sparse conversion rule for tensor reconstruction.
 class SparseTensorToTensorConverter : public OpConversionPattern<ToTensorOp> {
 public:
   using OpConversionPattern::OpConversionPattern;

mlir/lib/ExecutionEngine/SparseUtils.cpp

public:
   const std::vector<uint64_t> &getSizes() const { return sizes; }
   /// Getter for elements array.
   const std::vector<Element<V>> &getElements() const { return elements; }
 
   /// Factory method. Permutes the original dimensions according to
   /// the given ordering and expects subsequent add() calls to honor
   /// that same ordering for the given indices. The result is a
   /// fully permuted coordinate scheme.
-  static SparseTensorCOO<V> *newSparseTensorCOO(uint64_t size, uint64_t *sizes,
-                                                uint64_t *perm,
+  static SparseTensorCOO<V> *newSparseTensorCOO(uint64_t size,
+                                                const uint64_t *sizes,
+                                                const uint64_t *perm,
                                                 uint64_t capacity = 0) {
     std::vector<uint64_t> permsz(size);
     for (uint64_t r = 0; r < size; r++)
       permsz[perm[r]] = sizes[r];
     return new SparseTensorCOO<V>(permsz, capacity);
   }
 
 private:

 /// "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 : public SparseTensorStorageBase {
 public:
   /// Constructs a sparse tensor storage scheme from the given sparse
   /// tensor in coordinate scheme following the given per-rank dimension
   /// dense/sparse annotations.
-  SparseTensorStorage(SparseTensorCOO<V> *tensor, uint8_t *sparsity,
-                      uint64_t *perm)
+  SparseTensorStorage(SparseTensorCOO<V> *tensor, const uint8_t *sparsity,
+                      const uint64_t *perm)
       : sizes(tensor->getSizes()), rev(getRank()), pointers(getRank()),
         indices(getRank()) {
     // Store "reverse" permutation.
     for (uint64_t d = 0, rank = getRank(); d < rank; d++)
       rev[perm[d]] = d;
     // Provide hints on capacity.
     // TODO: needs fine-tuning based on sparsity
     uint64_t nnz = tensor->getElements().size();
public:
   void getIndices(std::vector<I> **out, uint64_t d) override {
     assert(d < getRank());
     *out = &indices[d];
   }
   void getValues(std::vector<V> **out) override { *out = &values; }
 
   /// Returns this sparse tensor storage scheme as a new memory-resident
   /// sparse tensor in coordinate scheme with the given dimension order.
-  SparseTensorCOO<V> *toCOO(uint64_t *perm) {
+  SparseTensorCOO<V> *toCOO(const uint64_t *perm) {
     // Restore original order of the dimension sizes and allocate coordinate
     // scheme with desired new ordering specified in perm.
     uint64_t size = getRank();
     std::vector<uint64_t> orgsz(size);
     for (uint64_t r = 0; r < size; r++)
       orgsz[rev[r]] = sizes[r];
     SparseTensorCOO<V> *tensor = SparseTensorCOO<V>::newSparseTensorCOO(
         size, orgsz.data(), perm, values.size());
     // Populate coordinate scheme restored from old ordering and changed with
     // new ordering. Rather than applying both reorderings during the recursion,
     // we compute the combine permutation in advance.
     std::vector<uint64_t> reord(size);
     for (uint64_t r = 0; r < size; r++)
       reord[r] = perm[rev[r]];
     std::vector<uint64_t> idx(size);
     toCOO(tensor, reord, idx, 0, 0);
     assert(tensor->getElements().size() == values.size());
     return tensor;
   }
 
   /// Factory method. Expects a coordinate scheme that respects the same
   /// permutation as is desired for the new sparse storage scheme.
-  static SparseTensorStorage<P, I, V> *
-  newSparseTensor(SparseTensorCOO<V> *t, uint8_t *sparsity, uint64_t *perm) {
+  static SparseTensorStorage<P, I, V> *newSparseTensor(SparseTensorCOO<V> *t,
+                                                       const uint8_t *sparsity,
+                                                       const uint64_t *perm) {
     t->sort(); // sort lexicographically
     SparseTensorStorage<P, I, V> *n =
         new SparseTensorStorage<P, I, V>(t, sparsity, perm);
     delete t;
     return n;
   }
 
 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-rank dimension dense/sparse annotations.
-  void fromCOO(SparseTensorCOO<V> *tensor, uint8_t *sparsity, uint64_t lo,
+  void fromCOO(SparseTensorCOO<V> *tensor, const uint8_t *sparsity, uint64_t lo,
                uint64_t hi, uint64_t d) {
     const std::vector<Element<V>> &elements = tensor->getElements();
     // Once dimensions are exhausted, insert the numerical values.
     if (d == getRank()) {
       assert(lo >= hi || lo < elements.size());
       values.push_back(lo < hi ? elements[lo].value : 0);
       return;
     }
for (uint64_t r = 0; r < idata[0]; r++) {
     }
   }
 }
 
 /// Reads a sparse tensor with the given filename into a memory-resident
 /// sparse tensor in coordinate scheme.
 template <typename V>
 static SparseTensorCOO<V> *openSparseTensorCOO(char *filename, uint64_t size,
-                                               uint64_t *sizes,
-                                               uint64_t *perm) {
+                                               const uint64_t *sizes,
+                                               const uint64_t *perm) {
   // Open the file.
   FILE *file = fopen(filename, "r");
   if (!file) {
     fprintf(stderr, "Cannot find %s\n", filename);
     exit(1);
   }
   // Perform some file format dependent set up.
   uint64_t idata[512];
#define TEMPLATE(NAME, TYPE)                                                   \
   struct NAME {                                                                \
     const TYPE *base;                                                          \
     const TYPE *data;                                                          \
     uint64_t off;                                                              \
     uint64_t sizes[1];                                                         \
     uint64_t strides[1];                                                       \
   }
 
+TEMPLATE(MemRef1DU64, uint64_t);
+TEMPLATE(MemRef1DU32, uint32_t);
+TEMPLATE(MemRef1DU16, uint16_t);
+TEMPLATE(MemRef1DU8, uint8_t);
+TEMPLATE(MemRef1DI64, int64_t);
+TEMPLATE(MemRef1DI32, int32_t);
+TEMPLATE(MemRef1DI16, int16_t);
+TEMPLATE(MemRef1DI8, int8_t);
+TEMPLATE(MemRef1DF64, double);
+TEMPLATE(MemRef1DF32, float);
+
 #define CASE(p, i, v, P, I, V)                                                 \
   if (ptrTp == (p) && indTp == (i) && valTp == (v)) {                          \
     SparseTensorCOO<V> *tensor = nullptr;                                      \
     if (action == 0)                                                           \
       tensor =                                                                 \
           openSparseTensorCOO<V>(static_cast<char *>(ptr), size, sizes, perm); \
     else if (action == 1)                                                      \
       tensor = static_cast<SparseTensorCOO<V> *>(ptr);                         \
     else if (action == 2)                                                      \
       return SparseTensorCOO<V>::newSparseTensorCOO(size, sizes, perm);        \
     else                                                                       \
       return static_cast<SparseTensorStorage<P, I, V> *>(ptr)->toCOO(perm);    \
     return SparseTensorStorage<P, I, V>::newSparseTensor(tensor, sparsity,     \
                                                          perm);                \
   }
 
-#define IMPL1(RET, NAME, TYPE, LIB)                                            \
-  RET NAME(void *tensor) {                                                     \
+#define IMPL1(REF, NAME, TYPE, LIB)                                            \
+  void _mlir_ciface_##NAME(REF *ref, void *tensor) {                           \
     std::vector<TYPE> *v;                                                      \
     static_cast<SparseTensorStorageBase *>(tensor)->LIB(&v);                   \
-    return {v->data(), v->data(), 0, {v->size()}, {1}};                        \
+    ref->base = ref->data = v->data();                                         \
+    ref->off = 0;                                                              \
+    ref->sizes[0] = v->size();                                                 \
+    ref->strides[0] = 1;                                                       \
   }
 
-#define IMPL2(RET, NAME, TYPE, LIB)                                            \
-  RET NAME(void *tensor, uint64_t d) {                                         \
+#define IMPL2(REF, NAME, TYPE, LIB)                                            \
+  void _mlir_ciface_##NAME(REF *ref, void *tensor, uint64_t d) {               \
     std::vector<TYPE> *v;                                                      \
     static_cast<SparseTensorStorageBase *>(tensor)->LIB(&v, d);                \
-    return {v->data(), v->data(), 0, {v->size()}, {1}};                        \
+    ref->base = ref->data = v->data();                                         \
+    ref->off = 0;                                                              \
+    ref->sizes[0] = v->size();                                                 \
+    ref->strides[0] = 1;                                                       \
   }
 
 #define IMPL3(NAME, TYPE)                                                      \
-  void *NAME(void *tensor, TYPE value, uint64_t *ibase, uint64_t *idata,       \
-             uint64_t ioff, uint64_t isize, uint64_t istride, uint64_t *pbase, \
-             uint64_t *pdata, uint64_t poff, uint64_t psize,                   \
-             uint64_t pstride) {                                               \
-    assert(istride == 1 && pstride == 1 && isize == psize);                    \
-    uint64_t *indx = idata + ioff;                                             \
+  void *_mlir_ciface_##NAME(void *tensor, TYPE value, MemRef1DU64 *iref,       \
+                            MemRef1DU64 *pref) {                               \
     if (!value)                                                                \
       return tensor;                                                           \
-    uint64_t *perm = pdata + poff;                                             \
+    assert(iref->strides[0] == 1 && pref->strides[0] == 1);                    \
+    assert(iref->sizes[0] == pref->sizes[0]);                                  \
+    const uint64_t *indx = iref->data + iref->off;                             \
+    const uint64_t *perm = pref->data + pref->off;                             \
+    uint64_t isize = iref->sizes[0];                                           \
     std::vector<uint64_t> indices(isize);                                      \
     for (uint64_t r = 0; r < isize; r++)                                       \
       indices[perm[r]] = indx[r];                                              \
     static_cast<SparseTensorCOO<TYPE> *>(tensor)->add(indices, value);         \
     return tensor;                                                             \
   }
 
-TEMPLATE(MemRef1DU64, uint64_t);
-TEMPLATE(MemRef1DU32, uint32_t);
-TEMPLATE(MemRef1DU16, uint16_t);
-TEMPLATE(MemRef1DU8, uint8_t);
-TEMPLATE(MemRef1DI64, int64_t);
-TEMPLATE(MemRef1DI32, int32_t);
-TEMPLATE(MemRef1DI16, int16_t);
-TEMPLATE(MemRef1DI8, int8_t);
-TEMPLATE(MemRef1DF64, double);
-TEMPLATE(MemRef1DF32, float);
-
 enum OverheadTypeEnum : uint64_t { kU64 = 1, kU32 = 2, kU16 = 3, kU8 = 4 };
 
 enum PrimaryTypeEnum : uint64_t {
   kF64 = 1,
   kF32 = 2,
   kI64 = 3,
   kI32 = 4,
   kI16 = 5,
   kI8 = 6
 };
 
 /// Constructs a new sparse tensor. This is the "swiss army knife"
 /// method for materializing sparse tensors into the computation.
 ///  action
 ///  0 : ptr contains filename to read into storage
 ///  1 : ptr contains coordinate scheme to assign to new storage
 ///  2 : returns empty coordinate scheme to fill (call back 1 to setup)
 ///  3 : returns coordinate scheme from storage in ptr (call back 1 to convert)
-void *newSparseTensor(uint8_t *abase, uint8_t *adata, uint64_t aoff,
-                      uint64_t asize, uint64_t astride, uint64_t *sbase,
-                      uint64_t *sdata, uint64_t soff, uint64_t ssize,
-                      uint64_t sstride, uint64_t *pbase, uint64_t *pdata,
-                      uint64_t poff, uint64_t psize, uint64_t pstride,
-                      uint64_t ptrTp, uint64_t indTp, uint64_t valTp,
+void *_mlir_ciface_newSparseTensor(MemRef1DU8 *aref, MemRef1DU64 *sref,
+                                   MemRef1DU64 *pref, uint64_t ptrTp,
+                                   uint64_t indTp, uint64_t valTp,
                                    uint32_t action, void *ptr) {
-  assert(astride == 1 && sstride == 1 && pstride == 1);
-  assert(asize == ssize && ssize == psize);
-  uint8_t *sparsity = adata + aoff;
-  uint64_t *sizes = sdata + soff;
-  uint64_t *perm = pdata + poff;
-  uint64_t size = asize;
+  assert(aref->strides[0] == 1 && sref->strides[0] == 1 &&
+         pref->strides[0] == 1);
+  assert(aref->sizes[0] == sref->sizes[0] && sref->sizes[0] == pref->sizes[0]);
+  const uint8_t *sparsity = aref->data + aref->off;
+  const uint64_t *sizes = sref->data + sref->off;
+  const uint64_t *perm = pref->data + pref->off;
+  uint64_t size = aref->sizes[0];
 
   // Double matrices with all combinations of overhead storage.
   CASE(kU64, kU64, kF64, uint64_t, uint64_t, double);
   CASE(kU64, kU32, kF64, uint64_t, uint32_t, double);
   CASE(kU64, kU16, kF64, uint64_t, uint16_t, double);
   CASE(kU64, kU8, kF64, uint64_t, uint8_t, double);
   CASE(kU32, kU64, kF64, uint32_t, uint64_t, double);
   CASE(kU32, kU32, kF64, uint32_t, uint32_t, double);