diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorConversion.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorConversion.cpp --- a/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorConversion.cpp +++ b/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorConversion.cpp @@ -27,7 +27,19 @@ 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: @@ -41,7 +53,8 @@ } } -/// 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) { @@ -159,15 +172,17 @@ 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( @@ -256,6 +271,8 @@ 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)) diff --git a/mlir/lib/ExecutionEngine/SparseUtils.cpp b/mlir/lib/ExecutionEngine/SparseUtils.cpp --- a/mlir/lib/ExecutionEngine/SparseUtils.cpp +++ b/mlir/lib/ExecutionEngine/SparseUtils.cpp @@ -18,6 +18,8 @@ #ifdef MLIR_CRUNNERUTILS_DEFINE_FUNCTIONS +#define AART + #include #include #include @@ -129,6 +131,7 @@ // Primary storage. virtual void getValues(std::vector **) { fatal("valf64"); } virtual void getValues(std::vector **) { fatal("valf32"); } + virtual void getValues(std::vector **) { fatal("vali64"); } virtual void getValues(std::vector **) { fatal("vali32"); } virtual void getValues(std::vector **) { fatal("vali16"); } virtual void getValues(std::vector **) { fatal("vali8"); } @@ -172,10 +175,86 @@ } // Then setup the tensor. traverse(tensor, sparsity, 0, nnz, 0); +#ifdef AART + dump(); +#endif } virtual ~SparseTensorStorage() {} +#ifdef AART + void dump() const { + fprintf(stderr, "++++++++++ rank=%lu +++++++++++\n", sizes.size()); + if constexpr (std::is_same_v) + fprintf(stderr, "PTR64 "); + else if constexpr (std::is_same_v) + fprintf(stderr, "PTR32 "); + else if constexpr (std::is_same_v) + fprintf(stderr, "PTR16 "); + else if constexpr (std::is_same_v) + fprintf(stderr, "PTR8 "); + if constexpr (std::is_same_v) + fprintf(stderr, "INDX64 "); + else if constexpr (std::is_same_v) + fprintf(stderr, "INDX32 "); + else if constexpr (std::is_same_v) + fprintf(stderr, "INDX16 "); + else if constexpr (std::is_same_v) + fprintf(stderr, "INDX8 "); + if constexpr (std::is_same_v) + fprintf(stderr, "VALF64\n"); + else if constexpr (std::is_same_v) + fprintf(stderr, "VALF32\n"); + else if constexpr (std::is_same_v) + fprintf(stderr, "VALI64\n"); + else if constexpr (std::is_same_v) + fprintf(stderr, "VALI32\n"); + else if constexpr (std::is_same_v) + fprintf(stderr, "VALI16\n"); + else if constexpr (std::is_same_v) + fprintf(stderr, "VALI8\n"); + for (uint64_t r = 0; r < sizes.size(); r++) { + fprintf(stderr, "dim %lu #%lu\n", r, sizes[r]); + fprintf(stderr, " positions[%lu] #%lu :", r, pointers[r].size()); + for (uint64_t i = 0; i < pointers[r].size(); i++) + if constexpr (std::is_same_v) + fprintf(stderr, " %lu", pointers[r][i]); + else if constexpr (std::is_same_v) + fprintf(stderr, " %u", pointers[r][i]); + else if constexpr (std::is_same_v) + fprintf(stderr, " %u", pointers[r][i]); + else if constexpr (std::is_same_v) + fprintf(stderr, " %u", pointers[r][i]); + fprintf(stderr, "\n indices[%lu] #%lu :", r, indices[r].size()); + for (uint64_t i = 0; i < indices[r].size(); i++) + if constexpr (std::is_same_v) + fprintf(stderr, " %lu", indices[r][i]); + else if constexpr (std::is_same_v) + fprintf(stderr, " %u", indices[r][i]); + else if constexpr (std::is_same_v) + fprintf(stderr, " %u", indices[r][i]); + else if constexpr (std::is_same_v) + fprintf(stderr, " %u", indices[r][i]); + fprintf(stderr, "\n"); + } + fprintf(stderr, "values #%lu :", values.size()); + for (uint64_t i = 0; i < values.size(); i++) + if constexpr (std::is_same_v) + fprintf(stderr, " %lf", values[i]); + else if constexpr (std::is_same_v) + fprintf(stderr, " %f", values[i]); + else if constexpr (std::is_same_v) + fprintf(stderr, " %ld", values[i]); + else if constexpr (std::is_same_v) + fprintf(stderr, " %d", values[i]); + else if constexpr (std::is_same_v) + fprintf(stderr, " %d", values[i]); + else if constexpr (std::is_same_v) + fprintf(stderr, " %d", values[i]); + fprintf(stderr, "\n+++++++++++++++++++++++++++++\n"); + } +#endif + uint64_t getRank() const { return sizes.size(); } uint64_t getDimSize(uint64_t d) override { return sizes[d]; } @@ -366,6 +445,16 @@ // Close the file and return sorted tensor. fclose(file); tensor->sort(); // sort lexicographically +#ifdef AART + const std::vector &elements = tensor->getElements(); + for (uint64_t k = 1; k < nnz; k++) { + uint64_t same = 0; + for (uint64_t r = 0; r < rank; r++) + if (elements[k].indices[r] == elements[k - 1].indices[r]) + same++; + assert(same < rank && "duplicate element"); + } +#endif return tensor; } @@ -373,6 +462,9 @@ template void *newSparseTensor(char *filename, uint8_t *sparsity, uint64_t *perm, uint64_t size) { +#ifdef AART + fprintf(stderr, "SPARSE SUPPORT LIB: OPEN %s\n", filename); +#endif SparseTensor *t = openTensor(filename, perm); assert(size == t->getRank()); // sparsity array must match rank SparseTensorStorageBase *tensor = @@ -437,6 +529,7 @@ 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); @@ -448,9 +541,10 @@ 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, @@ -499,6 +593,7 @@ 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); @@ -535,6 +630,7 @@ 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)