[mlir][sparse] adding OverheadType::kIndex

[wrengr]

Summary

Depends On D115008

This change opens the way for D115012, and removes some corner cases in CodegenUtils.cpp. The SparseTensorAttrDefs.td already specifies that we allow 0 bitwidth for the two overhead types and that it is interpreted to mean the architecture's native width.

[wrengr]

rebasing for changes to comments in the dependency D115008

[aartbik]

mlir/include/mlir/ExecutionEngine/SparseTensorUtils.h

20–21

"elemental type" does not seem right, given that we aslo say that for L24

mlir/lib/Dialect/SparseTensor/Utils/CodegenUtils.cpp

29 ↗

(On Diff #391514)

this default makes less sense now

mlir/lib/ExecutionEngine/SparseTensorUtils.cpp

850–851

This seems to be a compiler assert that should be outside any method

[wrengr]

Fixing comments

mlir/include/mlir/ExecutionEngine/SparseTensorUtils.h

20–21

yikes, how'd that get there. Fixed.

[aartbik]

Good to go after you address my last nits

mlir/test/Dialect/SparseTensor/conversion_sparse2dense.mlir

31–32

Let's remove this clarification. In the future, people may use automated tools to regenerate (some of) these CHECKs, and very few people actually read inside the CHECK blocks anyway

70–71

same

[wrengr]

fixing nits

[wrengr]

rebasing to updated D115008

[wrengr]

rebasing, to hopefully fix the windows build failure

[wrengr]

rebasing for changed D115008

[wrengr]

rebase

mlir/include/mlir/ExecutionEngine/SparseTensorUtils.h

 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_EXECUTIONENGINE_SPARSETENSORUTILS_H_
 #define MLIR_EXECUTIONENGINE_SPARSETENSORUTILS_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 };
+/// This type is used in the public API at all places where MLIR expects
+/// values with the built-in type "index". For now, we simply assume that
+/// type is 64-bit, but targets with different "index" bit widths should link
+/// with an alternatively built runtime support library.
+// TODO: support such targets?
+using index_t = uint64_t;
+
+/// Encoding of overhead types (both pointer overhead and indices
+/// overhead), for "overloading" @newSparseTensor.
+enum class OverheadType : uint32_t {
+  kIndex = 0,
+  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,

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

 using namespace mlir::sparse_tensor;
 
 //===----------------------------------------------------------------------===//
 // ExecutionEngine/SparseTensorUtils helper functions.
 //===----------------------------------------------------------------------===//
 
 OverheadType mlir::sparse_tensor::overheadTypeEncoding(unsigned width) {
   switch (width) {
-  default:
+  case 64:
     return OverheadType::kU64;
   case 32:
     return OverheadType::kU32;
   case 16:
     return OverheadType::kU16;
   case 8:
     return OverheadType::kU8;
+  case 0:
+    return OverheadType::kIndex;
   }
+  llvm_unreachable("Unsupported overhead bitwidth");
 }
 
 Type mlir::sparse_tensor::getOverheadType(Builder &builder, OverheadType ot) {
   switch (ot) {
+  case OverheadType::kIndex:
+    return builder.getIndexType();
   case OverheadType::kU64:
     return builder.getIntegerType(64);
   case OverheadType::kU32:
     return builder.getIntegerType(32);
   case OverheadType::kU16:
     return builder.getIntegerType(16);
   case OverheadType::kU8:
     return builder.getIntegerType(8);
   }
   llvm_unreachable("Unknown OverheadType");
 }
 
 Type mlir::sparse_tensor::getPointerOverheadType(
     Builder &builder, const SparseTensorEncodingAttr &enc) {
-  // NOTE(wrengr): This workaround will be fixed in D115010.
-  unsigned width = enc.getPointerBitWidth();
-  if (width == 0)
-    return builder.getIndexType();
-  return getOverheadType(builder, overheadTypeEncoding(width));
+  return getOverheadType(builder,
+                         overheadTypeEncoding(enc.getPointerBitWidth()));
 }
 
 Type mlir::sparse_tensor::getIndexOverheadType(
     Builder &builder, const SparseTensorEncodingAttr &enc) {
-  // NOTE(wrengr): This workaround will be fixed in D115010.
-  unsigned width = enc.getIndexBitWidth();
-  if (width == 0)
-    return builder.getIndexType();
-  return getOverheadType(builder, overheadTypeEncoding(width));
+  return getOverheadType(builder, overheadTypeEncoding(enc.getIndexBitWidth()));
 }
 
 PrimaryType mlir::sparse_tensor::primaryTypeEncoding(Type elemTp) {
   if (elemTp.isF64())
     return PrimaryType::kF64;
   if (elemTp.isF32())
     return PrimaryType::kF32;
   if (elemTp.isInteger(64))

mlir/lib/ExecutionEngine/SparseTensorUtils.cpp

static SparseTensorCOO<V> *openSparseTensorCOO(char *filename, uint64_t rank,
   fclose(file);
   return tensor;
 }
 
 } // namespace
 
 extern "C" {
 
-/// This type is used in the public API at all places where MLIR expects
-/// values with the built-in type "index". For now, we simply assume that
-/// type is 64-bit, but targets with different "index" bit widths should link
-/// with an alternatively built runtime support library.
-// TODO: support such targets?
-using index_t = uint64_t;
-
 //===----------------------------------------------------------------------===//
 //
 // Public API with methods that operate on MLIR buffers (memrefs) to interact
 // 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.
void _mlir_ciface_##NAME(                                                    \
     index_t *cursor = cref->data + cref->offset;                               \
     V *values = vref->data + vref->offset;                                     \
     bool *filled = fref->data + fref->offset;                                  \
     index_t *added = aref->data + aref->offset;                                \
     static_cast<SparseTensorStorageBase *>(tensor)->expInsert(                 \
         cursor, values, filled, added, count);                                 \
   }
 
+// Assume index_t is in fact uint64_t, so that _mlir_ciface_newSparseTensor
+// can safely rewrite kIndex to kU64.  We make this assertion to guarantee
+// that this file cannot get out of sync with its header.
+static_assert(std::is_same<index_t, uint64_t>::value,
+              "Expected index_t == uint64_t");
+
 /// 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
_mlir_ciface_newSparseTensor(StridedMemRefType<DimLevelType, 1> *aref, // NOLINT
   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 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];
 
+  // Rewrite kIndex to kU64, to avoid introducing a bunch of new cases.
+  // This is safe because of the static_assert above.
+  if (ptrTp == OverheadType::kIndex)
+    ptrTp = OverheadType::kU64;
+  if (indTp == OverheadType::kIndex)
+    indTp = OverheadType::kU64;
+
   // Double matrices with all combinations of overhead storage.
   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,

mlir/test/Dialect/SparseTensor/conversion_sparse2dense.mlir

 //   CHECK-DAG: %[[Attr0:.*]] = arith.constant 0 : i8
 //   CHECK-DAG: memref.store %[[Attr0]], %[[AttrsS]][%[[I0]]] : memref<1xi8>
 //   CHECK-DAG: %[[SizesS:.*]] = memref.alloca() : memref<1xindex>
 //   CHECK-DAG: %[[SizesD:.*]] = memref.cast %[[SizesS]] : memref<1xindex> to memref<?xindex>
 //   CHECK-DAG: memref.store %[[I13]], %[[SizesS]][%[[I0]]] : memref<1xindex>
 //   CHECK-DAG: %[[PermS:.*]] = memref.alloca() : memref<1xindex>
 //   CHECK-DAG: %[[PermD:.*]] = memref.cast %[[PermS]] : memref<1xindex> to memref<?xindex>
 //   CHECK-DAG: memref.store %[[I0]], %[[PermS]][%[[I0]]] : memref<1xindex>
-//   CHECK-DAG: %[[SecTp:.*]] = arith.constant 1 : i32
+//   CHECK-DAG: %[[zeroI32:.*]] = arith.constant 0 : i32
 //   CHECK-DAG: %[[ElemTp:.*]] = arith.constant 4 : i32
 //   CHECK-DAG: %[[ActionToIter:.*]] = arith.constant 5 : i32
-//   CHECK-DAG: %[[Iter:.*]] = call @newSparseTensor(%[[AttrsD]], %[[SizesD]], %[[PermD]], %[[SecTp]], %[[SecTp]], %[[ElemTp]], %[[ActionToIter]], %[[Arg]]) : (memref<?xi8>, memref<?xindex>, memref<?xindex>, i32, i32, i32, i32, !llvm.ptr<i8>) -> !llvm.ptr<i8>
+//   CHECK-DAG: %[[Iter:.*]] = call @newSparseTensor(%[[AttrsD]], %[[SizesD]], %[[PermD]], %[[zeroI32]], %[[zeroI32]], %[[ElemTp]], %[[ActionToIter]], %[[Arg]]) : (memref<?xi8>, memref<?xindex>, memref<?xindex>, i32, i32, i32, i32, !llvm.ptr<i8>) -> !llvm.ptr<i8>
 //   CHECK-DAG: %[[IndS:.*]] = memref.alloca() : memref<1xindex>
 //   CHECK-DAG: %[[IndD:.*]] = memref.cast %[[IndS]] : memref<1xindex> to memref<?xindex>
 //   CHECK-DAG: %[[ElemBuffer:.*]] = memref.alloca() : memref<i32>
 //   CHECK-DAG: %[[M:.*]] = memref.alloc() : memref<13xi32>
-//   CHECK-DAG: %[[E0:.*]] = arith.constant 0 : i32
-//   CHECK-DAG: linalg.fill(%[[E0]], %[[M]]) : i32, memref<13xi32>
+//   CHECK-DAG: linalg.fill(%[[zeroI32]], %[[M]]) : i32, memref<13xi32>
 //       CHECK: scf.while : () -> () {
 //       CHECK:   %[[Cond:.*]] = call @getNextI32(%[[Iter]], %[[IndD]], %[[ElemBuffer]]) : (!llvm.ptr<i8>, memref<?xindex>, memref<i32>) -> i1
 //       CHECK:   scf.condition(%[[Cond]])
 //       CHECK: } do {
 //       CHECK:   %[[Iv0:.*]] = memref.load %[[IndS]][%[[I0]]] : memref<1xindex>
 //       CHECK:   %[[ElemVal:.*]] = memref.load %[[ElemBuffer]][] : memref<i32>
 //       CHECK:   memref.store %[[ElemVal]], %[[M]][%[[Iv0]]] : memref<13xi32>
 //       CHECK:   scf.yield

 //   CHECK-DAG: memref.store %[[Attr0]], %[[AttrsS]][%[[I0]]] : memref<1xi8>
 //   CHECK-DAG: %[[SizesS:.*]] = memref.alloca() : memref<1xindex>
 //   CHECK-DAG: %[[SizesD:.*]] = memref.cast %[[SizesS]] : memref<1xindex> to memref<?xindex>
 //   CHECK-DAG: %[[SizeI0:.*]] = call @sparseDimSize(%[[Arg]], %[[I0]]) : (!llvm.ptr<i8>, index) -> index
 //   CHECK-DAG: memref.store %[[SizeI0]], %[[SizesS]][%[[I0]]] : memref<1xindex>
 //   CHECK-DAG: %[[PermS:.*]] = memref.alloca() : memref<1xindex>
 //   CHECK-DAG: %[[PermD:.*]] = memref.cast %[[PermS]] : memref<1xindex> to memref<?xindex>
 //   CHECK-DAG: memref.store %[[I0]], %[[PermS]][%[[I0]]] : memref<1xindex>
-//   CHECK-DAG: %[[SecTp:.*]] = arith.constant 1 : i32
+//   CHECK-DAG: %[[zeroI32:.*]] = arith.constant 0 : i32
 //   CHECK-DAG: %[[ElemTp:.*]] = arith.constant 4 : i32
 //   CHECK-DAG: %[[ActionToIter:.*]] = arith.constant 5 : i32
-//   CHECK-DAG: %[[Iter:.*]] = call @newSparseTensor(%[[AttrsD]], %[[SizesD]], %[[PermD]], %[[SecTp]], %[[SecTp]], %[[ElemTp]], %[[ActionToIter]], %[[Arg]]) : (memref<?xi8>, memref<?xindex>, memref<?xindex>, i32, i32, i32, i32, !llvm.ptr<i8>) -> !llvm.ptr<i8>
+//   CHECK-DAG: %[[Iter:.*]] = call @newSparseTensor(%[[AttrsD]], %[[SizesD]], %[[PermD]], %[[zeroI32]], %[[zeroI32]], %[[ElemTp]], %[[ActionToIter]], %[[Arg]]) : (memref<?xi8>, memref<?xindex>, memref<?xindex>, i32, i32, i32, i32, !llvm.ptr<i8>) -> !llvm.ptr<i8>
 //   CHECK-DAG: %[[IndS:.*]] = memref.alloca() : memref<1xindex>
 //   CHECK-DAG: %[[IndD:.*]] = memref.cast %[[IndS]] : memref<1xindex> to memref<?xindex>
 //   CHECK-DAG: %[[ElemBuffer:.*]] = memref.alloca() : memref<i32>
 //   CHECK-DAG: %[[M:.*]] = memref.alloc(%[[SizeI0]]) : memref<?xi32>
-//   CHECK-DAG: %[[E0:.*]] = arith.constant 0 : i32
-//   CHECK-DAG: linalg.fill(%[[E0]], %[[M]]) : i32, memref<?xi32>
+//   CHECK-DAG: linalg.fill(%[[zeroI32]], %[[M]]) : i32, memref<?xi32>
 //       CHECK: scf.while : () -> () {
 //       CHECK:   %[[Cond:.*]] = call @getNextI32(%[[Iter]], %[[IndD]], %[[ElemBuffer]]) : (!llvm.ptr<i8>, memref<?xindex>, memref<i32>) -> i1
 //       CHECK:   scf.condition(%[[Cond]])
 //       CHECK: } do {
 //       CHECK:   %[[Iv0:.*]] = memref.load %[[IndS]][%[[I0]]] : memref<1xindex>
 //       CHECK:   %[[ElemVal:.*]] = memref.load %[[ElemBuffer]][] : memref<i32>
 //       CHECK:   memref.store %[[ElemVal]], %[[M]][%[[Iv0]]] : memref<?xi32>
 //       CHECK:   scf.yield