[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

165–172

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

[aartbik]

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

165–172

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

matchAndRewrite(NewOp op, ArrayRef<Value> operands,
     // Secondary and primary types encoding.
     unsigned secPtr = getOverheadTypeEncoding(enc.getPointerBitWidth());
     unsigned secInd = getOverheadTypeEncoding(enc.getIndexBitWidth());
     unsigned primary;
     if (eltType.isF64())
       primary = 1;
     else if (eltType.isF32())
       primary = 2;
-    else if (eltType.isInteger(32))
+    else if (eltType.isInteger(64))
       primary = 3;
-    else if (eltType.isInteger(16))
+    else if (eltType.isInteger(32))
       primary = 4;
-    else if (eltType.isInteger(8))
+    else if (eltType.isInteger(16))
       primary = 5;
+    else if (eltType.isInteger(8))
+      primary = 6;
     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)));
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