[mlir][sparse] adding OverheadType::kIndex

[wrengr]

Summary

New version of D115010 to decouple it from D115008.

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]

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/SparseTensorConversion.cpp

static Value constantAction(ConversionPatternRewriter &rewriter, Location loc,
   return constantI32(rewriter, loc, static_cast<uint32_t>(action));
 }
 
 /// Generates a constant of the internal type encoding for overhead storage.
 static Value constantOverheadTypeEncoding(ConversionPatternRewriter &rewriter,
                                           Location loc, unsigned width) {
   OverheadType sec;
   switch (width) {
-  default:
+  case 64:
     sec = OverheadType::kU64;
     break;
   case 32:
     sec = OverheadType::kU32;
     break;
   case 16:
     sec = OverheadType::kU16;
     break;
   case 8:
     sec = OverheadType::kU8;
     break;
+  case 0:
+    sec = OverheadType::kIndex;
+    break;
+  default:
+    llvm_unreachable("Unsupported overhead bitwidth");
   }
   return constantI32(rewriter, loc, static_cast<uint32_t>(sec));
 }
 
 /// Generates a constant of the internal type encoding for pointer
 /// overhead storage.
 static Value constantPointerTypeEncoding(ConversionPatternRewriter &rewriter,
                                          Location loc,

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