[mlir][sparse] add int64 storage type to sparse tensor runtime support library

[aartbik]

Summary

This format was missing from the support library. Although there are some
subtleties reading in an external format for int64 as double, there is no
good reason to omit support for this data type form the support library.

[gussmith23]

Just one question, otherwise looks good!

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

180–186

Is it possible to use the PrimaryTypeEnum here, or use some kind of enum?

[aartbik]

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

180–186

Yes, that is a good suggestion. AFAIK, we are not allowed to create shared dependences between runtime and compiler (as in including the same header file). But I agree this hardcoding is a bit too error prone, so I at least created a shadow enum with a comment in this file.

[aartbik]

added note on value consistency with support lib

[aartbik]

oops, removed debug code

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

 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/Transforms/DialectConversion.h"
 
 using namespace mlir;
 using namespace mlir::sparse_tensor;
 
 namespace {
 
-/// Returns internal type encoding for overhead storage.
+/// Internal encoding of primary storage. Keep this enum consistent
+/// with the equivalent enum in the sparse runtime support library.
+enum PrimaryTypeEnum : uint64_t {
+  kF64 = 1,
+  kF32 = 2,
+  kI64 = 3,
+  kI32 = 4,
+  kI16 = 5,
+  kI8 = 6
+};
+
+/// Returns internal type encoding for overhead storage. Keep these
+/// values consistent with the sparse runtime support library.
 static unsigned getOverheadTypeEncoding(unsigned width) {
   switch (width) {
   default:
     return 1;
   case 32:
     return 2;
   case 16:
     return 3;
   case 8:
     return 4;
   }
 }
 
-/// Returns internal dimension level type encoding.
+/// Returns internal dimension level type encoding. Keep these
+/// values consistent with the sparse runtime support library.
 static unsigned
 getDimLevelTypeEncoding(SparseTensorEncodingAttr::DimLevelType dlt) {
   switch (dlt) {
   case SparseTensorEncodingAttr::DimLevelType::Dense:
     return 0;
   case SparseTensorEncodingAttr::DimLevelType::Compressed:
     return 1;
   case SparseTensorEncodingAttr::DimLevelType::Singleton:
if (p) {
         perm[i] = APInt(64, i);
     }
     params.push_back(getTensor(rewriter, 64, loc, perm));
     // Secondary and primary types encoding.
     unsigned secPtr = getOverheadTypeEncoding(enc.getPointerBitWidth());
     unsigned secInd = getOverheadTypeEncoding(enc.getIndexBitWidth());
     unsigned primary;
     if (eltType.isF64())
-      primary = 1;
+      primary = kF64;
     else if (eltType.isF32())
-      primary = 2;
+      primary = kF32;
+    else if (eltType.isInteger(64))
+      primary = kI64;
     else if (eltType.isInteger(32))
-      primary = 3;
+      primary = kI32;
     else if (eltType.isInteger(16))
-      primary = 4;
+      primary = kI16;
     else if (eltType.isInteger(8))
-      primary = 5;
+      primary = kI8;
     else
       return failure();
     params.push_back(
         rewriter.create<ConstantOp>(loc, rewriter.getI64IntegerAttr(secPtr)));
     params.push_back(
         rewriter.create<ConstantOp>(loc, rewriter.getI64IntegerAttr(secInd)));
     params.push_back(
         rewriter.create<ConstantOp>(loc, rewriter.getI64IntegerAttr(primary)));
     // Generate the call to create new tensor.
matchAndRewrite(ToValuesOp op, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const override {
     Type resType = op.getType();
     Type eltType = resType.cast<ShapedType>().getElementType();
     StringRef name;
     if (eltType.isF64())
       name = "sparseValuesF64";
     else if (eltType.isF32())
       name = "sparseValuesF32";
+    else if (eltType.isInteger(64))
+      name = "sparseValuesI64";
     else if (eltType.isInteger(32))
       name = "sparseValuesI32";
     else if (eltType.isInteger(16))
       name = "sparseValuesI16";
     else if (eltType.isInteger(8))
       name = "sparseValuesI8";
     else
       return failure();

mlir/lib/ExecutionEngine/SparseUtils.cpp

public:
   virtual void getIndices(std::vector<uint64_t> **, uint64_t) { fatal("i64"); }
   virtual void getIndices(std::vector<uint32_t> **, uint64_t) { fatal("i32"); }
   virtual void getIndices(std::vector<uint16_t> **, uint64_t) { fatal("i16"); }
   virtual void getIndices(std::vector<uint8_t> **, uint64_t) { fatal("i8"); }
 
   // Primary storage.
   virtual void getValues(std::vector<double> **) { fatal("valf64"); }
   virtual void getValues(std::vector<float> **) { fatal("valf32"); }
+  virtual void getValues(std::vector<int64_t> **) { fatal("vali64"); }
   virtual void getValues(std::vector<int32_t> **) { fatal("vali32"); }
   virtual void getValues(std::vector<int16_t> **) { fatal("vali16"); }
   virtual void getValues(std::vector<int8_t> **) { fatal("vali8"); }
 
   virtual ~SparseTensorStorageBase() {}
 
 private:
   void fatal(const char *tp) {
RET NAME(void *tensor, uint64_t d) {                                         \
     static_cast<SparseTensorStorageBase *>(tensor)->LIB(&v, d);                \
     return {v->data(), v->data(), 0, {v->size()}, {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);
 
 enum OverheadTypeEnum : uint64_t { kU64 = 1, kU32 = 2, kU16 = 3, kU8 = 4 };
 
 enum PrimaryTypeEnum : uint64_t {
   kF64 = 1,
   kF32 = 2,
-  kI32 = 3,
-  kI16 = 4,
-  kI8 = 5
+  kI64 = 3,
+  kI32 = 4,
+  kI16 = 5,
+  kI8 = 6
 };
 
 void *newSparseTensor(char *filename, uint8_t *abase, uint8_t *adata,
                       uint64_t aoff, uint64_t asize, uint64_t astride,
                       uint64_t *pbase, uint64_t *pdata, uint64_t poff,
                       uint64_t psize, uint64_t pstride, uint64_t ptrTp,
                       uint64_t indTp, uint64_t valTp) {
   assert(astride == 1 && pstride == 1);
void *newSparseTensor(char *filename, uint8_t *abase, uint8_t *adata,
   CASE(kU16, kU16, kF32, uint16_t, uint16_t, float);
   CASE(kU16, kU8, kF32, uint16_t, uint8_t, float);
   CASE(kU8, kU64, kF32, uint8_t, uint64_t, float);
   CASE(kU8, kU32, kF32, uint8_t, uint32_t, float);
   CASE(kU8, kU16, kF32, uint8_t, uint16_t, float);
   CASE(kU8, kU8, kF32, uint8_t, uint8_t, float);
 
   // Integral matrices with same overhead storage.
+  CASE(kU64, kU64, kI64, uint64_t, uint64_t, int64_t);
   CASE(kU64, kU64, kI32, uint64_t, uint64_t, int32_t);
   CASE(kU64, kU64, kI16, uint64_t, uint64_t, int16_t);
   CASE(kU64, kU64, kI8, uint64_t, uint64_t, int8_t);
   CASE(kU32, kU32, kI32, uint32_t, uint32_t, int32_t);
   CASE(kU32, kU32, kI16, uint32_t, uint32_t, int16_t);
   CASE(kU32, kU32, kI8, uint32_t, uint32_t, int8_t);
   CASE(kU16, kU16, kI32, uint16_t, uint16_t, int32_t);
   CASE(kU16, kU16, kI16, uint16_t, uint16_t, int16_t);

 IMPL2(MemRef1DU8, sparsePointers8, uint8_t, getPointers)
 IMPL2(MemRef1DU64, sparseIndices, uint64_t, getIndices)
 IMPL2(MemRef1DU64, sparseIndices64, uint64_t, getIndices)
 IMPL2(MemRef1DU32, sparseIndices32, uint32_t, getIndices)
 IMPL2(MemRef1DU16, sparseIndices16, uint16_t, getIndices)
 IMPL2(MemRef1DU8, sparseIndices8, uint8_t, getIndices)
 IMPL1(MemRef1DF64, sparseValuesF64, double, getValues)
 IMPL1(MemRef1DF32, sparseValuesF32, float, getValues)
+IMPL1(MemRef1DI64, sparseValuesI64, int64_t, getValues)
 IMPL1(MemRef1DI32, sparseValuesI32, int32_t, getValues)
 IMPL1(MemRef1DI16, sparseValuesI16, int16_t, getValues)
 IMPL1(MemRef1DI8, sparseValuesI8, int8_t, getValues)
 
 void delSparseTensor(void *tensor) {
   delete static_cast<SparseTensorStorageBase *>(tensor);
 }
 
 #undef TEMPLATE
 #undef CASE
 #undef IMPL1
 #undef IMPL2
 
 } // extern "C"
 
 #endif // MLIR_CRUNNERUTILS_DEFINE_FUNCTIONS