[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

12

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.

17

move this line inside the macro ifdef

21

empty line before this line

46

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

54

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–693

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

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

21

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

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

662–693

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 ↗

(On Diff #384848)

SparseTensorUtils.h now ;-)

14 ↗

(On Diff #384848)

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/SparseUtils.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
+// While that's not very hygienic, it helps avoid hard to track down
+// bugs about mismatching magic numbers.
+//
+//===----------------------------------------------------------------------===//
+
+#include <cinttypes>
+
+#ifndef MLIR_EXECUTIONENGINE_SPARSEUTILS_H_
+#define MLIR_EXECUTIONENGINE_SPARSEUTILS_H_
+extern "C" {
+// TODO(wrengr): can these really be changed into "enum class" instead, while
+// still retaining the extern-C?
+
+// TODO(wrengr): should rename these so they look sufficiently different
+// from the enumerators of PrimaryType; also in case we ever wanted to
+// add unsigned types to PrimaryType.
+//
+/// 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
+};
+
+/// A copy of `SparseTensorEncodingAttr::DimLevelType` for breaking dependency
+/// cycles.
+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/SparseUtils.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;
+  // DO_NOT_SUBMIT(wrengr): we could try to collapse these first two cases by
+  // switching on FloatType::getWidth(); howeve, that unfortunately conflates
+  // Float16Type with BFloat16Type.  More generally we could theoretically
+  // switch on FloatType::getFloatSemantics(), but that's a struct so that won't
+  // work out very well.
+  if (tp.isF64()) {
+    primary = PrimaryType::kF64;
+  } else if (tp.isF32()) {
+    primary = PrimaryType::kF32;
+  } else if (auto intTy = tp.dyn_cast<IntegerType>()) {
+    // DO_NOT_SUBMIT(wrengr): this is consistent with the old if-chain, however
+    // it does ignore the signedness semantics which we really shouldn't do.
+    switch (intTy.getWidth()) {
+    case 64:
+      primary = PrimaryType::kI64;
+      break;
+    case 32:
+      primary = PrimaryType::kI32;
+      break;
+    case 16:
+      primary = PrimaryType::kI16;
+      break;
+    case 8:
+      primary = PrimaryType::kI8;
+      break;
+    default:
+      llvm_unreachable("Unsupported integer element type");
+    }
+  } 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) {
+  // TODO(wrengr): figure out how to avoid the repetition, yet retain the
+  // llvm_unreachable check. Then again, I doubt C++ does any sort of
+  // completeness analysis on switch statements, so there wouldn't be much lost
+  // by using a default case.
+  switch (dlt) {
+  case SparseTensorEncodingAttr::DimLevelType::Dense:
+    return constantI8(rewriter, loc, (uint8_t)DimLevelType::kDense);
+  case SparseTensorEncodingAttr::DimLevelType::Compressed:
+    return constantI8(rewriter, loc, (uint8_t)DimLevelType::kCompressed);
+  case SparseTensorEncodingAttr::DimLevelType::Singleton:
+    return constantI8(rewriter, loc, (uint8_t)DimLevelType::kSingleton);
+  }
+  llvm_unreachable("Unknown SparseTensorEncodingAttr::DimLevelType");
+}
+
 /// 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) {
       if (encDst == encSrc) {
         rewriter.replaceOp(op, adaptor.getOperands()); // hidden nop cast
         return success();
       }
       SmallVector<Value, 4> sizes;
       SmallVector<Value, 8> params;
       sizesFromPtr(rewriter, sizes, op, encSrc, srcType.cast<ShapedType>(),
                    src);
-      newParams(rewriter, params, op, encDst, kToCOO, sizes, src);
+      newParams(rewriter, params, op, encDst, Action::kToCOO, sizes, src);
       Value coo = genNewCall(rewriter, op, params);
-      params[6] = constantI32(rewriter, loc, kFromCOO);
+      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;
       //   }
       //
       // ConvertOp::verify() ensures that these casts will succeed,
       // and that dstTensorTp subsumes srcTensorTp.
       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/SparseUtils.cpp

 //===- SparseUtils.cpp - Sparse 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/SparseUtils.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 IMPL1(NAME, TYPE, LIB)                                                 \
   void _mlir_ciface_##NAME(StridedMemRefType<TYPE, 1> *ref, void *tensor) {    \
     assert(ref);                                                               \
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:
-/// 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 IMPL_GETNEXT to use)
+/// Action:
+/// Empty = returns empty storage to fill later
+/// FromFile = returns storage, where ptr contains filename to read
+/// FromCOO = returns storage, where ptr contains coordinate scheme to assign
+/// EmptyCOO = returns empty coordinate scheme to fill and use with FromCOO
+/// ToCOO = returns coordinate scheme from storage in ptr to use with FromCOO
+/// ToIterator = returns iterator from storage in ptr (call IMPL_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);
+  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 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::kI64, uint64_t,
+       uint64_t, int64_t);
+  CASE(OverheadType::kU64, OverheadType::kU64, PrimaryType::kI32, uint64_t,
+       uint64_t, int32_t);
+  CASE(OverheadType::kU64, OverheadType::kU64, PrimaryType::kI16, uint64_t,
+       uint64_t, int16_t);
+  CASE(OverheadType::kU64, OverheadType::kU64, PrimaryType::kI8, uint64_t,
+       uint64_t, int8_t);
+  CASE(OverheadType::kU32, OverheadType::kU32, PrimaryType::kI32, uint32_t,
+       uint32_t, int32_t);
+  CASE(OverheadType::kU32, OverheadType::kU32, PrimaryType::kI16, uint32_t,
+       uint32_t, int16_t);
+  CASE(OverheadType::kU32, OverheadType::kU32, PrimaryType::kI8, uint32_t,
+       uint32_t, int8_t);
+  CASE(OverheadType::kU16, OverheadType::kU16, PrimaryType::kI32, uint16_t,
+       uint16_t, int32_t);
+  CASE(OverheadType::kU16, OverheadType::kU16, PrimaryType::kI16, uint16_t,
+       uint16_t, int16_t);
+  CASE(OverheadType::kU16, OverheadType::kU16, PrimaryType::kI8, uint16_t,
+       uint16_t, int8_t);
+  CASE(OverheadType::kU8, OverheadType::kU8, PrimaryType::kI32, uint8_t,
+       uint8_t, int32_t);
+  CASE(OverheadType::kU8, OverheadType::kU8, PrimaryType::kI16, uint8_t,
+       uint8_t, int16_t);
+  CASE(OverheadType::kU8, OverheadType::kU8, PrimaryType::kI8, uint8_t, 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, indices, and values.
 IMPL2(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/SparseUtils.h",
+    ],
     includes = ["include"],
     deps = [
         ":Affine",
         ":ArithmeticDialect",
         ":IR",
         ":LLVMDialect",
         ":LinalgOps",
         ":LinalgTransforms",

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