[mlir][sparse] Factoring magic numbers into a header

[wrengr]

Summary

Addresses https://bugs.llvm.org/show_bug.cgi?id=52303

[wrengr]

Removing notes to self

[aartbik]

Wren, please note that bug https://bugs.llvm.org/show_bug.cgi?id=52303 was already taken by somebody else. I will reassign and coordinate with them, but please let's make sure we do not cause too many conflicts with external work while that is ongoing.

[aartbik]

I more or less expected a CMakeLists.txt change too, but I guess ${MLIR_MAIN_INCLUDE_DIR}/mlir/ExecutionEngine takes care of that already

mlir/include/mlir/ExecutionEngine/SparseUtils.h

11 ↗

(On Diff #383886)

Although I don't disagree, let's not put comments like this in the code.
The good part is that we share the enums ;-)

This situation is hard to solve since we are *generating* code that adheres to a certain API.
So it is very hard to enforce full API consistency with just headers.

16 ↗

(On Diff #383886)

move this line inside the macro ifdef

20 ↗

(On Diff #383886)

empty line before this line

45 ↗

(On Diff #383886)

Rather than saying "copy of", let's say that this enum mimics the SparseTensorEncodingAttr::DimLevelType (which is the source of truth) and its values should be kept consistent with that enum

53 ↗

(On Diff #383886)

empty line after this line

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

79

I assume that without the (uint32_t) cast is does not compile?

90

I found the original way of testing a bit more readable than having this cast branch just to safe a few (hidden) casts.

mlir/lib/ExecutionEngine/SparseUtils.cpp

662 ↗

(On Diff #383886)

it was never pretty, but now with the extra lines this has become a lot larger
anyway to avoid that?

[aartbik]

mlir/lib/ExecutionEngine/SparseUtils.cpp

645 ↗

(On Diff #383886)

wasn't this change to getNext() in one of your previous revisions?
is that lost again? or do we need to rebase this change to main?

[aartbik]

Heads up Wren.

It looks like https://reviews.llvm.org/D113043 will go in first, so you will need to (trivially) rebase against that name change

[wrengr]

Addressing comments

[wrengr]

mlir/include/mlir/ExecutionEngine/SparseUtils.h

20 ↗

(On Diff #383886)

Do we have a tool for enforcing this? I've just been following the usual mlir formatting/linting tool, but it didn't complain here.

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

79

Yes, unfortunately. If one of the other casting styles is preferred for this circumstance, let me know

mlir/lib/ExecutionEngine/SparseUtils.cpp

645 ↗

(On Diff #383886)

I think it's still just shrapnel from the original rebase/landing glitch

662 ↗

(On Diff #383886)

C++20 adds a using enum statement that would fix it; but for C++14 I'll have to think a bit more...

[wrengr]

Crossing off my todo

[aartbik]

mlir/include/mlir/ExecutionEngine/SparseTensorUtils.h

1

SparseTensorUtils.h now ;-)

14

SPARSETENSORUTILS, here and below

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

79

Yes, in that case I believe that static_cast<> is preferred over the hold style C cast syntax.

[wrengr]

Updating casting style. And fixing SparseUtils => SparseTensorUtils renaming.

[aartbik]

Nice progress! Ship it!

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

92

nit: LLVM style prefer no-braces if all branches are single statement

[wrengr]

LLVM style

[wrengr]

rebased

mlir/include/mlir/ExecutionEngine/SparseTensorUtils.h

+//===- SparseUtils.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
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_EXECUTIONENGINE_SPARSEUTILS_H_
+#define MLIR_EXECUTIONENGINE_SPARSEUTILS_H_
+
+#include <cinttypes>
+
+extern "C" {
+
+/// Encoding of the elemental type, for "overloading" @newSparseTensor.
+enum class OverheadType : uint32_t { kU64 = 1, kU32 = 2, kU16 = 3, kU8 = 4 };
+
+/// 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
+};
+
+/// The actions performed by @newSparseTensor.
+enum class Action : uint32_t {
+  kEmpty = 0,
+  kFromFile = 1,
+  kFromCOO = 2,
+  kEmptyCOO = 3,
+  kToCOO = 4,
+  kToIterator = 5
+};
+
+/// 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
+};
+
+} // extern "C"
+
+#endif // MLIR_EXECUTIONENGINE_SPARSEUTILS_H_

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

 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
 #include "mlir/Dialect/Linalg/Utils/Utils.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/SCF/SCF.h"
 #include "mlir/Dialect/SparseTensor/IR/SparseTensor.h"
 #include "mlir/Dialect/SparseTensor/Transforms/Passes.h"
 #include "mlir/Dialect/StandardOps/IR/Ops.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/ExecutionEngine/SparseTensorUtils.h"
 #include "mlir/Transforms/DialectConversion.h"
 
 using namespace mlir;
 using namespace mlir::sparse_tensor;
 
 namespace {
 
-/// New tensor storage action. Keep these values consistent with
-/// the sparse runtime support library.
-enum Action : uint32_t {
-  kEmpty = 0,
-  kFromFile = 1,
-  kFromCOO = 2,
-  kEmptyCOO = 3,
-  kToCOO = 4,
-  kToIter = 5
-};
-
 //===----------------------------------------------------------------------===//
 // Helper methods.
 //===----------------------------------------------------------------------===//
 
-/// Returns internal type encoding for primary storage. Keep these
-/// values consistent with the sparse runtime support library.
-static uint32_t getPrimaryTypeEncoding(Type tp) {
-  if (tp.isF64())
-    return 1;
-  if (tp.isF32())
-    return 2;
-  if (tp.isInteger(64))
-    return 3;
-  if (tp.isInteger(32))
-    return 4;
-  if (tp.isInteger(16))
-    return 5;
-  if (tp.isInteger(8))
-    return 6;
-  return 0;
-}
-
-/// Returns internal type encoding for overhead storage. Keep these
-/// values consistent with the sparse runtime support library.
-static uint32_t 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. Keep these
-/// values consistent with the sparse runtime support library.
-static uint32_t
-getDimLevelTypeEncoding(SparseTensorEncodingAttr::DimLevelType dlt) {
-  switch (dlt) {
-  case SparseTensorEncodingAttr::DimLevelType::Dense:
-    return 0;
-  case SparseTensorEncodingAttr::DimLevelType::Compressed:
-    return 1;
-  case SparseTensorEncodingAttr::DimLevelType::Singleton:
-    return 2;
-  }
-  llvm_unreachable("Unknown SparseTensorEncodingAttr::DimLevelType");
-}
-
 /// Generates a constant zero of the given type.
 inline static Value constantZero(ConversionPatternRewriter &rewriter,
                                  Location loc, Type t) {
   return rewriter.create<arith::ConstantOp>(loc, t, rewriter.getZeroAttr(t));
 }
 
 /// Generates a constant of `index` type.
 inline static Value constantIndex(ConversionPatternRewriter &rewriter,
                                   Location loc, int64_t i) {
   return rewriter.create<arith::ConstantIndexOp>(loc, i);
 }
 
 /// Generates a constant of `i32` type.
 inline static Value constantI32(ConversionPatternRewriter &rewriter,
                                 Location loc, int32_t i) {
   return rewriter.create<arith::ConstantIntOp>(loc, i, 32);
 }
 
 /// Generates a constant of `i8` type.
 inline static Value constantI8(ConversionPatternRewriter &rewriter,
                                Location loc, int8_t i) {
   return rewriter.create<arith::ConstantIntOp>(loc, i, 8);
 }
 
+/// Returns a constant of the internal type encoding for overhead storage.
+static Value constantOverheadTypeEncoding(ConversionPatternRewriter &rewriter,
+                                          Location loc, unsigned width) {
+  OverheadType sec;
+  switch (width) {
+  default:
+    sec = OverheadType::kU64;
+    break;
+  case 32:
+    sec = OverheadType::kU32;
+    break;
+  case 16:
+    sec = OverheadType::kU16;
+    break;
+  case 8:
+    sec = OverheadType::kU8;
+    break;
+  }
+  return constantI32(rewriter, loc, (uint32_t)sec);
+}
+
+/// Returns a constant of the internal type encoding for primary storage.
+static Value constantPrimaryTypeEncoding(ConversionPatternRewriter &rewriter,
+                                         Location loc, Type tp) {
+  PrimaryType primary;
+  if (tp.isF64()) {
+    primary = PrimaryType::kF64;
+  } else if (tp.isF32()) {
+    primary = PrimaryType::kF32;
+  } else if (tp.isInteger(64)) {
+    primary = PrimaryType::kI64;
+  } else if (tp.isInteger(32)) {
+    primary = PrimaryType::kI32;
+  } else if (tp.isInteger(16)) {
+    primary = PrimaryType::kI16;
+  } else if (tp.isInteger(8)) {
+    primary = PrimaryType::kI8;
+  } else {
+    llvm_unreachable("Unknown element type");
+  }
+  return constantI32(rewriter, loc, (uint32_t)primary);
+}
+
+/// Returns a constant of the internal dimension level type encoding.
+static Value
+constantDimLevelTypeEncoding(ConversionPatternRewriter &rewriter, Location loc,
+                             SparseTensorEncodingAttr::DimLevelType dlt) {
+  DimLevelType dlt2;
+  switch (dlt) {
+  case SparseTensorEncodingAttr::DimLevelType::Dense:
+    dlt2 = DimLevelType::kDense;
+    break;
+  case SparseTensorEncodingAttr::DimLevelType::Compressed:
+    dlt2 = DimLevelType::kCompressed;
+    break;
+  case SparseTensorEncodingAttr::DimLevelType::Singleton:
+    dlt2 = DimLevelType::kSingleton;
+    break;
+  default:
+    llvm_unreachable("Unknown SparseTensorEncodingAttr::DimLevelType");
+  }
+  return constantI8(rewriter, loc, (uint8_t)dlt2);
+}
+
 /// 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,
                                  TypeRange resultType, ValueRange operands,
                                  bool emitCInterface = false) {
   MLIRContext *context = op->getContext();
static Value genBuffer(ConversionPatternRewriter &rewriter, Location loc,
   return buffer;
 }
 
 /// Populates parameters required to call the "swiss army knife" method of the
 /// sparse runtime support library for materializing sparse tensors into the
 /// computation.
 static void newParams(ConversionPatternRewriter &rewriter,
                       SmallVector<Value, 8> &params, Operation *op,
-                      SparseTensorEncodingAttr &enc, uint32_t action,
+                      SparseTensorEncodingAttr &enc, Action action,
                       ValueRange szs, Value ptr = Value()) {
   Location loc = op->getLoc();
   ArrayRef<SparseTensorEncodingAttr::DimLevelType> dlt = enc.getDimLevelType();
   unsigned sz = dlt.size();
   // Sparsity annotations.
   SmallVector<Value, 4> attrs;
   for (unsigned i = 0; i < sz; i++)
-    attrs.push_back(constantI8(rewriter, loc, getDimLevelTypeEncoding(dlt[i])));
+    attrs.push_back(constantDimLevelTypeEncoding(rewriter, loc, dlt[i]));
   params.push_back(genBuffer(rewriter, loc, attrs));
   // Dimension sizes array of the enveloping tensor. Useful for either
   // verification of external data, or for construction of internal data.
   SmallVector<Value, 4> sizes;
   for (Value s : szs)
     sizes.push_back(s);
   params.push_back(genBuffer(rewriter, loc, sizes));
   // Dimension order permutation array. This is the "identity" permutation by
   // default, or otherwise the "reverse" permutation of a given ordering, so
   // that indices can be mapped quickly to the right position.
   SmallVector<Value, 4> rev(sz);
   if (AffineMap p = enc.getDimOrdering()) {
     for (unsigned i = 0; i < sz; i++)
       rev[p.getDimPosition(i)] = constantIndex(rewriter, loc, i);
   } else {
     for (unsigned i = 0; i < sz; i++)
       rev[i] = constantIndex(rewriter, loc, i);
   }
   params.push_back(genBuffer(rewriter, loc, rev));
   // Secondary and primary types encoding.
   ShapedType resType = op->getResult(0).getType().cast<ShapedType>();
-  uint32_t secPtr = getOverheadTypeEncoding(enc.getPointerBitWidth());
-  uint32_t secInd = getOverheadTypeEncoding(enc.getIndexBitWidth());
-  uint32_t primary = getPrimaryTypeEncoding(resType.getElementType());
-  assert(primary);
-  params.push_back(constantI32(rewriter, loc, secPtr));
-  params.push_back(constantI32(rewriter, loc, secInd));
-  params.push_back(constantI32(rewriter, loc, primary));
+  params.push_back(
+      constantOverheadTypeEncoding(rewriter, loc, enc.getPointerBitWidth()));
+  params.push_back(
+      constantOverheadTypeEncoding(rewriter, loc, enc.getIndexBitWidth()));
+  params.push_back(
+      constantPrimaryTypeEncoding(rewriter, loc, resType.getElementType()));
   // User action and pointer.
   Type pTp = LLVM::LLVMPointerType::get(rewriter.getI8Type());
   if (!ptr)
     ptr = rewriter.create<LLVM::NullOp>(loc, pTp);
-  params.push_back(constantI32(rewriter, loc, action));
+  params.push_back(constantI32(rewriter, loc, (uint32_t)action));
   params.push_back(ptr);
 }
 
 /// Generates the comparison `v != 0` where `v` is of numeric type `t`.
 /// For floating types, we use the "unordered" comparator (i.e., returns
 /// true if `v` is NaN).
 static Value genIsNonzero(ConversionPatternRewriter &rewriter, Location loc,
                           Value v) {
matchAndRewrite(NewOp op, OpAdaptor adaptor,
     if (!enc)
       return failure();
     // Generate the call to construct tensor from ptr. The sizes are
     // inferred from the result type of the new operator.
     SmallVector<Value, 4> sizes;
     SmallVector<Value, 8> params;
     sizesFromType(rewriter, sizes, op.getLoc(), resType.cast<ShapedType>());
     Value ptr = adaptor.getOperands()[0];
-    newParams(rewriter, params, op, enc, kFromFile, sizes, ptr);
+    newParams(rewriter, params, op, enc, Action::kFromFile, sizes, ptr);
     rewriter.replaceOp(op, genNewCall(rewriter, op, params));
     return success();
   }
 };
 
 /// Sparse conversion rule for the init operator.
 class SparseTensorInitConverter : public OpConversionPattern<InitOp> {
   using OpConversionPattern::OpConversionPattern;
   LogicalResult
   matchAndRewrite(InitOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     Type resType = op.getType();
     auto enc = getSparseTensorEncoding(resType);
     if (!enc)
       return failure();
     // Generate the call to construct empty tensor. The sizes are
     // explicitly defined by the arguments to the init operator.
     SmallVector<Value, 8> params;
-    newParams(rewriter, params, op, enc, kEmpty, adaptor.getOperands());
+    newParams(rewriter, params, op, enc, Action::kEmpty, adaptor.getOperands());
     rewriter.replaceOp(op, genNewCall(rewriter, op, params));
     return success();
   }
 };
 
 /// Sparse conversion rule for the convert operator.
 class SparseTensorConvertConverter : public OpConversionPattern<ConvertOp> {
   using OpConversionPattern::OpConversionPattern;
if (encDst && encSrc) {
       sizesFromPtr(rewriter, sizes, op, encSrc, srcType.cast<ShapedType>(),
                    src);
       // Set up encoding with right mix of src and dst so that the two
       // method calls can share most parameters, while still providing
       // the correct sparsity information to either of them.
       auto enc = SparseTensorEncodingAttr::get(
           op->getContext(), encDst.getDimLevelType(), encDst.getDimOrdering(),
           encSrc.getPointerBitWidth(), encSrc.getIndexBitWidth());
-      newParams(rewriter, params, op, enc, kToCOO, sizes, src);
+      newParams(rewriter, params, op, enc, Action::kToCOO, sizes, src);
       Value coo = genNewCall(rewriter, op, params);
-      params[3] = constantI32(
-          rewriter, loc, getOverheadTypeEncoding(encDst.getPointerBitWidth()));
-      params[4] = constantI32(
-          rewriter, loc, getOverheadTypeEncoding(encDst.getIndexBitWidth()));
-      params[6] = constantI32(rewriter, loc, kFromCOO);
+      params[3] = constantOverheadTypeEncoding(rewriter, loc,
+                                               encDst.getPointerBitWidth());
+      params[4] = constantOverheadTypeEncoding(rewriter, loc,
+                                               encDst.getIndexBitWidth());
+      params[6] = constantI32(rewriter, loc, (uint32_t)Action::kFromCOO);
       params[7] = coo;
       rewriter.replaceOp(op, genNewCall(rewriter, op, params));
       return success();
     }
     if (!encDst && encSrc) {
       // This is sparse => dense conversion, which is handled as follows:
       //   dst = new Tensor(0);
       //   iter = src->toCOO();
       //   iter->startIterator();
       //   while (elem = iter->getNext()) {
       //     dst[elem.indices] = elem.value;
       //   }
       RankedTensorType dstTensorTp = resType.cast<RankedTensorType>();
       RankedTensorType srcTensorTp = srcType.cast<RankedTensorType>();
       unsigned rank = dstTensorTp.getRank();
       Type elemTp = dstTensorTp.getElementType();
       // Fabricate a no-permutation encoding for newParams().
       // The pointer/index types must be those of `src`.
-      // The dimLevelTypes aren't actually used by kToIter.
+      // The dimLevelTypes aren't actually used by Action::kToIterator.
       encDst = SparseTensorEncodingAttr::get(
           op->getContext(),
           SmallVector<SparseTensorEncodingAttr::DimLevelType>(
               rank, SparseTensorEncodingAttr::DimLevelType::Dense),
           AffineMap(), encSrc.getPointerBitWidth(), encSrc.getIndexBitWidth());
       SmallVector<Value, 4> sizes;
       SmallVector<Value, 8> params;
       sizesFromPtr(rewriter, sizes, op, encSrc, srcTensorTp, src);
-      newParams(rewriter, params, op, encDst, kToIter, sizes, src);
+      newParams(rewriter, params, op, encDst, Action::kToIterator, sizes, src);
       Value iter = genNewCall(rewriter, op, params);
       Value ind = genAlloca(rewriter, loc, rank, rewriter.getIndexType());
       Value elemPtr = genAllocaScalar(rewriter, loc, elemTp);
       Value dst = allocDenseTensor(rewriter, loc, dstTensorTp, sizes);
       SmallVector<Value> noArgs;
       SmallVector<Type> noTypes;
       auto whileOp = rewriter.create<scf::WhileOp>(loc, noTypes, noArgs);
       Block *before = rewriter.createBlock(&whileOp.before(), {}, noTypes);
matchAndRewrite(ConvertOp op, OpAdaptor adaptor,
     // Also note that the code below only generates the "new" ops and
     // the loop-nest per se; whereas the entire body of the innermost
     // loop is generated by genAddElt().
     ShapedType stp = resType.cast<ShapedType>();
     unsigned rank = stp.getRank();
     SmallVector<Value, 4> sizes;
     SmallVector<Value, 8> params;
     sizesFromSrc(rewriter, sizes, loc, src);
-    newParams(rewriter, params, op, encDst, kEmptyCOO, sizes);
+    newParams(rewriter, params, op, encDst, Action::kEmptyCOO, sizes);
     Value ptr = genNewCall(rewriter, op, params);
     Value ind = genAlloca(rewriter, loc, rank, rewriter.getIndexType());
     Value perm = params[2];
     SmallVector<Value> lo;
     SmallVector<Value> hi;
     SmallVector<Value> st;
     Value zero = constantIndex(rewriter, loc, 0);
     Value one = constantIndex(rewriter, loc, 1);
scf::buildLoopNest(
             val = genIndexAndValueForSparse(rewriter, loc, indices, values, ind,
                                             ivs, rank);
           else
             val = genIndexAndValueForDense(rewriter, loc, src, ind, ivs);
           genAddEltCall(rewriter, op, eltType, ptr, val, ind, perm);
           return {};
         });
     // Final call to construct sparse tensor storage.
-    params[6] = constantI32(rewriter, loc, kFromCOO);
+    params[6] = constantI32(rewriter, loc, (uint32_t)Action::kFromCOO);
     params[7] = ptr;
     rewriter.replaceOp(op, genNewCall(rewriter, op, params));
     return success();
   }
 };
 
 /// Sparse conversion rule for the release operator.
 class SparseTensorReleaseConverter : public OpConversionPattern<ReleaseOp> {

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 <cctype>
 #include <cinttypes>
private:
   bool iteratorLocked;
   unsigned iteratorPos;
 };
 
 /// Abstract base class of sparse tensor storage. Note that we use
 /// function overloading to implement "partial" method specialization.
 class SparseTensorStorageBase {
 public:
-  enum DimLevelType : uint8_t { kDense = 0, kCompressed = 1, kSingleton = 2 };
-
   virtual uint64_t getDimSize(uint64_t) = 0;
 
   // Overhead storage.
   virtual void getPointers(std::vector<uint64_t> **, uint64_t) { fatal("p64"); }
   virtual void getPointers(std::vector<uint32_t> **, uint64_t) { fatal("p32"); }
   virtual void getPointers(std::vector<uint16_t> **, uint64_t) { fatal("p16"); }
   virtual void getPointers(std::vector<uint8_t> **, uint64_t) { fatal("p8"); }
   virtual void getIndices(std::vector<uint64_t> **, uint64_t) { fatal("i64"); }

 /// 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 with the given dimensions,
   /// permutation, and per-dimension dense/sparse annotations, using
   /// the coordinate scheme tensor for the initial contents if provided.
   SparseTensorStorage(const std::vector<uint64_t> &szs, const uint64_t *perm,
-                      const uint8_t *sparsity, SparseTensorCOO<V> *tensor)
+                      const DimLevelType *sparsity, SparseTensorCOO<V> *tensor)
       : sizes(szs), rev(getRank()), pointers(getRank()), indices(getRank()) {
     uint64_t rank = getRank();
     // Store "reverse" permutation.
     for (uint64_t r = 0; r < rank; r++)
       rev[perm[r]] = r;
     // Provide hints on capacity of pointers and indices.
     // TODO: needs fine-tuning based on sparsity
     for (uint64_t r = 0, s = 1; r < rank; r++) {
       s *= sizes[r];
-      if (sparsity[r] == kCompressed) {
+      if (sparsity[r] == DimLevelType::kCompressed) {
         pointers[r].reserve(s + 1);
         indices[r].reserve(s);
         s = 1;
       } else {
-        assert(sparsity[r] == kDense && "singleton not yet supported");
+        assert(sparsity[r] == DimLevelType::kDense &&
+               "singleton not yet supported");
       }
     }
     // Prepare sparse pointer structures for all dimensions.
     for (uint64_t r = 0; r < rank; r++)
-      if (sparsity[r] == kCompressed)
+      if (sparsity[r] == DimLevelType::kCompressed)
         pointers[r].push_back(0);
     // Then assign contents from coordinate scheme tensor if provided.
     if (tensor) {
       uint64_t nnz = tensor->getElements().size();
       values.reserve(nnz);
       fromCOO(tensor, sparsity, 0, nnz, 0);
     }
   }
public:
 
   /// Factory method. Constructs a sparse tensor storage scheme with the given
   /// dimensions, permutation, and per-dimension dense/sparse annotations,
   /// using the coordinate scheme tensor for the initial contents if provided.
   /// In the latter case, the coordinate scheme must respect the same
   /// permutation as is desired for the new sparse tensor storage.
   static SparseTensorStorage<P, I, V> *
   newSparseTensor(uint64_t rank, const uint64_t *sizes, const uint64_t *perm,
-                  const uint8_t *sparsity, SparseTensorCOO<V> *tensor) {
+                  const DimLevelType *sparsity, SparseTensorCOO<V> *tensor) {
     SparseTensorStorage<P, I, V> *n = nullptr;
     if (tensor) {
       assert(tensor->getRank() == rank);
       for (uint64_t r = 0; r < rank; r++)
         assert(sizes[r] == 0 || tensor->getSizes()[perm[r]] == sizes[r]);
       tensor->sort(); // sort lexicographically
       n = new SparseTensorStorage<P, I, V>(tensor->getSizes(), perm, sparsity,
                                            tensor);
       delete tensor;
     } else {
       std::vector<uint64_t> permsz(rank);
       for (uint64_t r = 0; r < rank; r++)
         permsz[perm[r]] = sizes[r];
       n = new SparseTensorStorage<P, I, V>(permsz, perm, sparsity, tensor);
     }
     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-dimension dense/sparse annotations.
-  void fromCOO(SparseTensorCOO<V> *tensor, const uint8_t *sparsity, uint64_t lo,
-               uint64_t hi, uint64_t d) {
+  void fromCOO(SparseTensorCOO<V> *tensor, const DimLevelType *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;
     }
     assert(d < getRank());
     // Visit all elements in this interval.
     uint64_t full = 0;
     while (lo < hi) {
       assert(lo < elements.size() && hi <= elements.size());
       // Find segment in interval with same index elements in this dimension.
       uint64_t idx = elements[lo].indices[d];
       uint64_t seg = lo + 1;
       while (seg < hi && elements[seg].indices[d] == idx)
         seg++;
       // Handle segment in interval for sparse or dense dimension.
-      if (sparsity[d] == kCompressed) {
+      if (sparsity[d] == DimLevelType::kCompressed) {
         indices[d].push_back(idx);
       } else {
         // For dense storage we must fill in all the zero values between
         // the previous element (when last we ran this for-loop) and the
         // current element.
         for (; full < idx; full++)
           fromCOO(tensor, sparsity, 0, 0, d + 1); // pass empty
         full++;
       }
       fromCOO(tensor, sparsity, lo, seg, d + 1);
       // And move on to next segment in interval.
       lo = seg;
     }
     // Finalize the sparse pointer structure at this dimension.
-    if (sparsity[d] == kCompressed) {
+    if (sparsity[d] == DimLevelType::kCompressed) {
       pointers[d].push_back(indices[d].size());
     } else {
       // For dense storage we must fill in all the zero values after
       // the last element.
       for (uint64_t sz = sizes[d]; full < sz; full++)
         fromCOO(tensor, sparsity, 0, 0, d + 1); // pass empty
     }
   }

 // 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.
 //
 //===----------------------------------------------------------------------===//
 
-enum OverheadTypeEnum : uint32_t { kU64 = 1, kU32 = 2, kU16 = 3, kU8 = 4 };
-
-enum PrimaryTypeEnum : uint32_t {
-  kF64 = 1,
-  kF32 = 2,
-  kI64 = 3,
-  kI32 = 4,
-  kI16 = 5,
-  kI8 = 6
-};
-
-enum Action : uint32_t {
-  kEmpty = 0,
-  kFromFile = 1,
-  kFromCOO = 2,
-  kEmptyCOO = 3,
-  kToCOO = 4,
-  kToIter = 5
-};
-
 #define CASE(p, i, v, P, I, V)                                                 \
   if (ptrTp == (p) && indTp == (i) && valTp == (v)) {                          \
     SparseTensorCOO<V> *tensor = nullptr;                                      \
-    if (action <= kFromCOO) {                                                  \
-      if (action == kFromFile) {                                               \
+    if (action <= Action::kFromCOO) {                                          \
+      if (action == Action::kFromFile) {                                       \
         char *filename = static_cast<char *>(ptr);                             \
         tensor = openSparseTensorCOO<V>(filename, rank, sizes, perm);          \
-      } else if (action == kFromCOO) {                                         \
+      } else if (action == Action::kFromCOO) {                                 \
         tensor = static_cast<SparseTensorCOO<V> *>(ptr);                       \
       } else {                                                                 \
-        assert(action == kEmpty);                                              \
+        assert(action == Action::kEmpty);                                      \
       }                                                                        \
       return SparseTensorStorage<P, I, V>::newSparseTensor(rank, sizes, perm,  \
                                                            sparsity, tensor);  \
-    } else if (action == kEmptyCOO) {                                          \
+    } else if (action == Action::kEmptyCOO) {                                  \
       return SparseTensorCOO<V>::newSparseTensorCOO(rank, sizes, perm);        \
     } else {                                                                   \
       tensor = static_cast<SparseTensorStorage<P, I, V> *>(ptr)->toCOO(perm);  \
-      if (action == kToIter) {                                                 \
+      if (action == Action::kToIterator) {                                     \
         tensor->startIterator();                                               \
       } else {                                                                 \
-        assert(action == kToCOO);                                              \
+        assert(action == Action::kToCOO);                                      \
       }                                                                        \
       return tensor;                                                           \
     }                                                                          \
   }
 
+#define CASE_SECSAME(p, v, P, V) CASE(p, p, v, P, P, V)
+
 #define IMPL_SPARSEVALUES(NAME, TYPE, LIB)                                     \
   void _mlir_ciface_##NAME(StridedMemRefType<TYPE, 1> *ref, void *tensor) {    \
     assert(ref);                                                               \
     assert(tensor);                                                            \
     std::vector<TYPE> *v;                                                      \
     static_cast<SparseTensorStorageBase *>(tensor)->LIB(&v);                   \
     ref->basePtr = ref->data = v->data();                                      \
     ref->offset = 0;                                                           \
for (uint64_t r = 0; r < isize; r++)                                       \
       indx[r] = elem->indices[r];                                              \
     *value = elem->value;                                                      \
     return true;                                                               \
   }
 
 /// Constructs a new sparse tensor. This is the "swiss army knife"
 /// method for materializing sparse tensors into the computation.
 ///
-/// action:
+/// 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
-/// kToIter = returns iterator from storage in ptr (call getNext() to use)
+/// kToIterator = returns iterator from storage in ptr (call getNext() to use)
 void *
-_mlir_ciface_newSparseTensor(StridedMemRefType<uint8_t, 1> *aref, // NOLINT
+_mlir_ciface_newSparseTensor(StridedMemRefType<DimLevelType, 1> *aref, // NOLINT
                              StridedMemRefType<index_t, 1> *sref,
                              StridedMemRefType<index_t, 1> *pref,
-                             uint32_t ptrTp, uint32_t indTp, uint32_t valTp,
-                             uint32_t action, void *ptr) {
+                             OverheadType ptrTp, OverheadType indTp,
+                             PrimaryType valTp, Action action, void *ptr) {
   assert(aref && sref && pref);
   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->offset;
+  const DimLevelType *sparsity = aref->data + aref->offset;
   const index_t *sizes = sref->data + sref->offset;
   const index_t *perm = pref->data + pref->offset;
   uint64_t rank = 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);
-  CASE(kU32, kU16, kF64, uint32_t, uint16_t, double);
-  CASE(kU32, kU8, kF64, uint32_t, uint8_t, double);
-  CASE(kU16, kU64, kF64, uint16_t, uint64_t, double);
-  CASE(kU16, kU32, kF64, uint16_t, uint32_t, double);
-  CASE(kU16, kU16, kF64, uint16_t, uint16_t, double);
-  CASE(kU16, kU8, kF64, uint16_t, uint8_t, double);
-  CASE(kU8, kU64, kF64, uint8_t, uint64_t, double);
-  CASE(kU8, kU32, kF64, uint8_t, uint32_t, double);
-  CASE(kU8, kU16, kF64, uint8_t, uint16_t, double);
-  CASE(kU8, kU8, kF64, uint8_t, uint8_t, double);
+  CASE(OverheadType::kU64, OverheadType::kU64, PrimaryType::kF64, uint64_t,
+       uint64_t, double);
+  CASE(OverheadType::kU64, OverheadType::kU32, PrimaryType::kF64, uint64_t,
+       uint32_t, double);
+  CASE(OverheadType::kU64, OverheadType::kU16, PrimaryType::kF64, uint64_t,
+       uint16_t, double);
+  CASE(OverheadType::kU64, OverheadType::kU8, PrimaryType::kF64, uint64_t,
+       uint8_t, double);
+  CASE(OverheadType::kU32, OverheadType::kU64, PrimaryType::kF64, uint32_t,
+       uint64_t, double);
+  CASE(OverheadType::kU32, OverheadType::kU32, PrimaryType::kF64, uint32_t,
+       uint32_t, double);
+  CASE(OverheadType::kU32, OverheadType::kU16, PrimaryType::kF64, uint32_t,
+       uint16_t, double);
+  CASE(OverheadType::kU32, OverheadType::kU8, PrimaryType::kF64, uint32_t,
+       uint8_t, double);
+  CASE(OverheadType::kU16, OverheadType::kU64, PrimaryType::kF64, uint16_t,
+       uint64_t, double);
+  CASE(OverheadType::kU16, OverheadType::kU32, PrimaryType::kF64, uint16_t,
+       uint32_t, double);
+  CASE(OverheadType::kU16, OverheadType::kU16, PrimaryType::kF64, uint16_t,
+       uint16_t, double);
+  CASE(OverheadType::kU16, OverheadType::kU8, PrimaryType::kF64, uint16_t,
+       uint8_t, double);
+  CASE(OverheadType::kU8, OverheadType::kU64, PrimaryType::kF64, uint8_t,
+       uint64_t, double);
+  CASE(OverheadType::kU8, OverheadType::kU32, PrimaryType::kF64, uint8_t,
+       uint32_t, double);
+  CASE(OverheadType::kU8, OverheadType::kU16, PrimaryType::kF64, uint8_t,
+       uint16_t, double);
+  CASE(OverheadType::kU8, OverheadType::kU8, PrimaryType::kF64, uint8_t,
+       uint8_t, double);
 
   // Float matrices with all combinations of overhead storage.
-  CASE(kU64, kU64, kF32, uint64_t, uint64_t, float);
-  CASE(kU64, kU32, kF32, uint64_t, uint32_t, float);
-  CASE(kU64, kU16, kF32, uint64_t, uint16_t, float);
-  CASE(kU64, kU8, kF32, uint64_t, uint8_t, float);
-  CASE(kU32, kU64, kF32, uint32_t, uint64_t, float);
-  CASE(kU32, kU32, kF32, uint32_t, uint32_t, float);
-  CASE(kU32, kU16, kF32, uint32_t, uint16_t, float);
-  CASE(kU32, kU8, kF32, uint32_t, uint8_t, float);
-  CASE(kU16, kU64, kF32, uint16_t, uint64_t, float);
-  CASE(kU16, kU32, kF32, uint16_t, uint32_t, float);
-  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);
-  CASE(kU16, kU16, kI8, uint16_t, uint16_t, int8_t);
-  CASE(kU8, kU8, kI32, uint8_t, uint8_t, int32_t);
-  CASE(kU8, kU8, kI16, uint8_t, uint8_t, int16_t);
-  CASE(kU8, kU8, kI8, uint8_t, uint8_t, int8_t);
+  CASE(OverheadType::kU64, OverheadType::kU64, PrimaryType::kF32, uint64_t,
+       uint64_t, float);
+  CASE(OverheadType::kU64, OverheadType::kU32, PrimaryType::kF32, uint64_t,
+       uint32_t, float);
+  CASE(OverheadType::kU64, OverheadType::kU16, PrimaryType::kF32, uint64_t,
+       uint16_t, float);
+  CASE(OverheadType::kU64, OverheadType::kU8, PrimaryType::kF32, uint64_t,
+       uint8_t, float);
+  CASE(OverheadType::kU32, OverheadType::kU64, PrimaryType::kF32, uint32_t,
+       uint64_t, float);
+  CASE(OverheadType::kU32, OverheadType::kU32, PrimaryType::kF32, uint32_t,
+       uint32_t, float);
+  CASE(OverheadType::kU32, OverheadType::kU16, PrimaryType::kF32, uint32_t,
+       uint16_t, float);
+  CASE(OverheadType::kU32, OverheadType::kU8, PrimaryType::kF32, uint32_t,
+       uint8_t, float);
+  CASE(OverheadType::kU16, OverheadType::kU64, PrimaryType::kF32, uint16_t,
+       uint64_t, float);
+  CASE(OverheadType::kU16, OverheadType::kU32, PrimaryType::kF32, uint16_t,
+       uint32_t, float);
+  CASE(OverheadType::kU16, OverheadType::kU16, PrimaryType::kF32, uint16_t,
+       uint16_t, float);
+  CASE(OverheadType::kU16, OverheadType::kU8, PrimaryType::kF32, uint16_t,
+       uint8_t, float);
+  CASE(OverheadType::kU8, OverheadType::kU64, PrimaryType::kF32, uint8_t,
+       uint64_t, float);
+  CASE(OverheadType::kU8, OverheadType::kU32, PrimaryType::kF32, uint8_t,
+       uint32_t, float);
+  CASE(OverheadType::kU8, OverheadType::kU16, PrimaryType::kF32, uint8_t,
+       uint16_t, float);
+  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::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::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::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);
 
   // Unsupported case (add above if needed).
   fputs("unsupported combination of types\n", stderr);
   exit(1);
 }
 
 /// Methods that provide direct access to pointers.
 IMPL_GETOVERHEAD(sparsePointers, index_t, getPointers)

 ///      values  = [1.0, 5.0, 3.0]
 ///      indices = [ 0, 0,  1, 1,  1, 2]
 //
 // TODO: for now f64 tensors only, no dim ordering, all dimensions compressed
 //
 void *convertToMLIRSparseTensor(uint64_t rank, uint64_t nse, uint64_t *shape,
                                 double *values, uint64_t *indices) {
   // Setup all-dims compressed and default ordering.
-  std::vector<uint8_t> sparse(rank, SparseTensorStorageBase::kCompressed);
+  std::vector<DimLevelType> sparse(rank, DimLevelType::kCompressed);
   std::vector<uint64_t> perm(rank);
   std::iota(perm.begin(), perm.end(), 0);
   // Convert external format to internal COO.
   SparseTensorCOO<double> *tensor = SparseTensorCOO<double>::newSparseTensorCOO(
       rank, shape, perm.data(), nse);
   std::vector<uint64_t> idx(rank);
   for (uint64_t i = 0, base = 0; i < nse; i++) {
     for (uint64_t r = 0; r < rank; r++)

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

deps = [
         ":StandardOps",
         "//llvm:Support",
     ],
 )
 
 cc_library(
     name = "SparseTensorTransforms",
     srcs = glob(["lib/Dialect/SparseTensor/Transforms/*.cpp"]),
-    hdrs = ["include/mlir/Dialect/SparseTensor/Transforms/Passes.h"],
+    hdrs = [
+        "include/mlir/Dialect/SparseTensor/Transforms/Passes.h",
+        "include/mlir/ExecutionEngine/SparseTensorUtils.h",
+    ],
     includes = ["include"],
     deps = [
         ":Affine",
         ":ArithmeticDialect",
         ":IR",
         ":LLVMDialect",
         ":LinalgOps",
         ":LinalgTransforms",

 )
 
 cc_library(
     name = "mlir_c_runner_utils",
     srcs = [
         "lib/ExecutionEngine/CRunnerUtils.cpp",
         "lib/ExecutionEngine/SparseTensorUtils.cpp",
     ],
-    hdrs = ["include/mlir/ExecutionEngine/CRunnerUtils.h"],
+    hdrs = [
+        "include/mlir/ExecutionEngine/CRunnerUtils.h",
+        "include/mlir/ExecutionEngine/SparseTensorUtils.h",
+    ],
     includes = ["include"],
 )
 
 cc_library(
     name = "mlir_async_runtime_api",
     hdrs = ["include/mlir/ExecutionEngine/AsyncRuntime.h"],
     includes = ["include"],
 )