[mlir][sparse] add memSizes array to sparse storage format

[aartbik]

Summary

Rationale:
For every dynamic memref (memref<?xtype>), the stored size really
indicates the capacity and the entry in the memSizes indicates
the actual size. This allows us to use memref's as "vectors".

[aartbik]

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

104

I simply moved this up so we can precompute #fields easier

546

I moved base class closer to actual usage. It was split.

626–627

I changed this to order in which they appear in file

[aartbik]

typo

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

static Optional<Value> sizeFromTensorAtDim(OpBuilder &rewriter, Location loc,
   auto tuple =
       llvm::cast<UnrealizedConversionCastOp>(adaptedValue.getDefiningOp());
   return rewriter
       .create<memref::LoadOp>(loc, tuple.getInputs().front(),
                               constantIndex(rewriter, loc, toStored(enc, dim)))
       .getResult();
 }
 
+/// Returns field index of sparse tensor type for pointers/indices, when set.
+static unsigned getFieldIndex(Type type, unsigned ptrDim, unsigned idxDim) {
+  auto enc = getSparseTensorEncoding(type);
+  assert(enc);
+  RankedTensorType rType = type.cast<RankedTensorType>();
+  unsigned field = 2; // start past sizes
+  unsigned ptr = 0;
+  unsigned idx = 0;
+  for (unsigned r = 0, rank = rType.getShape().size(); r < rank; r++) {
+    switch (enc.getDimLevelType()[r]) {
+    case SparseTensorEncodingAttr::DimLevelType::Dense:
+      break; // no fields
+    case SparseTensorEncodingAttr::DimLevelType::Compressed:
+    case SparseTensorEncodingAttr::DimLevelType::CompressedNu:
+    case SparseTensorEncodingAttr::DimLevelType::CompressedNo:
+    case SparseTensorEncodingAttr::DimLevelType::CompressedNuNo:
+      if (ptr++ == ptrDim)
+        return field;
+      field++;
+      if (idx++ == idxDim)
+        return field;
+      field++;
+      break;
+    case SparseTensorEncodingAttr::DimLevelType::Singleton:
+    case SparseTensorEncodingAttr::DimLevelType::SingletonNu:
+    case SparseTensorEncodingAttr::DimLevelType::SingletonNo:
+    case SparseTensorEncodingAttr::DimLevelType::SingletonNuNo:
+      if (idx++ == idxDim)
+        return field;
+      field++;
+      break;
+    }
+  }
+  return field + 1; // return values field index
+}
+
 /// Maps a sparse tensor type to the appropriate compounded buffers.
 static Optional<LogicalResult>
 convertSparseTensorType(Type type, SmallVectorImpl<Type> &fields) {
   auto enc = getSparseTensorEncoding(type);
   if (!enc)
     return llvm::None;
   // Construct the basic types.
   auto context = type.getContext();
   unsigned idxWidth = enc.getIndexBitWidth();
   unsigned ptrWidth = enc.getPointerBitWidth();
   RankedTensorType rType = type.cast<RankedTensorType>();
   Type indexType = IndexType::get(context);
   Type idxType = idxWidth ? IntegerType::get(context, idxWidth) : indexType;
   Type ptrType = ptrWidth ? IntegerType::get(context, ptrWidth) : indexType;
   Type eltType = rType.getElementType();
   //
   // Sparse tensor storage for rank-dimensional tensor is organized as a
-  // single compound type with the following fields:
+  // single compound type with the following fields. Note that every
+  // memref with ? size actually behaves as a "vector", i.e. the stored
+  // size is the capacity and the used size resides in the memSizes array.
   //
   // struct {
   //   memref<rank x index> dimSizes     ; size in each dimension
+  //   memref<n x index> memSizes        ; sizes of ptrs/inds/values
   //   ; per-dimension d:
   //   ;  if dense:
   //        <nothing>
   //   ;  if compresed:
   //        memref<? x ptr>  pointers-d  ; pointers for sparse dim d
   //        memref<? x idx>  indices-d   ; indices for sparse dim d
   //   ;  if singleton:
   //        memref<? x idx>  indices-d   ; indices for singleton dim d
   //   memref<? x eltType> values        ; values
   // };
   //
   unsigned rank = rType.getShape().size();
   // The dimSizes array.
   fields.push_back(MemRefType::get({rank}, indexType));
+  // The memSizes array.
+  unsigned lastField = getFieldIndex(type, -1, -1);
+  fields.push_back(MemRefType::get({lastField - 2}, indexType));
   // Per-dimension storage.
   for (unsigned r = 0; r < rank; r++) {
     // Dimension level types apply in order to the reordered dimension.
     // As a result, the compound type can be constructed directly in the given
     // order. Clients of this type know what field is what from the sparse
     // tensor type.
     switch (enc.getDimLevelType()[r]) {
     case SparseTensorEncodingAttr::DimLevelType::Dense:
for (unsigned r = 0; r < rank; r++) {
     case SparseTensorEncodingAttr::DimLevelType::SingletonNo:
     case SparseTensorEncodingAttr::DimLevelType::SingletonNuNo:
       fields.push_back(MemRefType::get({ShapedType::kDynamicSize}, idxType));
       break;
     }
   }
   // The values array.
   fields.push_back(MemRefType::get({ShapedType::kDynamicSize}, eltType));
+  assert(fields.size() == lastField);
   return success();
 }
 
-// Returns field index of sparse tensor type for pointers/indices, when set.
-static unsigned getFieldIndex(Type type, unsigned ptrDim, unsigned idxDim) {
-  auto enc = getSparseTensorEncoding(type);
-  assert(enc);
-  RankedTensorType rType = type.cast<RankedTensorType>();
-  unsigned field = 1; // start at DimSizes;
-  unsigned ptr = 0;
-  unsigned idx = 0;
-  for (unsigned r = 0, rank = rType.getShape().size(); r < rank; r++) {
-    switch (enc.getDimLevelType()[r]) {
-    case SparseTensorEncodingAttr::DimLevelType::Dense:
-      break; // no fields
-    case SparseTensorEncodingAttr::DimLevelType::Compressed:
-    case SparseTensorEncodingAttr::DimLevelType::CompressedNu:
-    case SparseTensorEncodingAttr::DimLevelType::CompressedNo:
-    case SparseTensorEncodingAttr::DimLevelType::CompressedNuNo:
-      if (ptr++ == ptrDim)
-        return field;
-      field++;
-      if (idx++ == idxDim)
-        return field;
-      field++;
-      break;
-    case SparseTensorEncodingAttr::DimLevelType::Singleton:
-    case SparseTensorEncodingAttr::DimLevelType::SingletonNu:
-    case SparseTensorEncodingAttr::DimLevelType::SingletonNo:
-    case SparseTensorEncodingAttr::DimLevelType::SingletonNuNo:
-      if (idx++ == idxDim)
-        return field;
-      field++;
-      break;
-    }
-  }
-  llvm_unreachable("failed to find ptr/idx field index");
-  return -1;
-}
-
 /// Create allocation operation.
 static Value createAllocation(OpBuilder &builder, Location loc, Type type,
                               Value sz) {
   auto memType = MemRefType::get({ShapedType::kDynamicSize}, type);
   return builder.create<memref::AllocOp>(loc, memType, sz);
 }
 
-/// Creates allocation for each field in sparse tensor type.
+/// Creates allocation for each field in sparse tensor type. Note that
+/// for all dynamic memrefs, the memory size is really the capacity of
+/// the "vector", while the actual size resides in the sizes array.
 ///
 /// TODO: for efficiency, we will need heuristis to make educated guesses
-///       on the required final sizes; also, we will need an improved
-///       memory allocation scheme with capacity and reallocation
+///       on the required capacities
 ///
 static void createAllocFields(OpBuilder &builder, Location loc, Type type,
                               ValueRange dynSizes,
                               SmallVectorImpl<Value> &fields) {
   auto enc = getSparseTensorEncoding(type);
   assert(enc);
   // Construct the basic types.
   unsigned idxWidth = enc.getIndexBitWidth();
if (ShapedType::isDynamic(shape[r]))
       sizes.push_back(dynSizes[o++]);
     else
       sizes.push_back(constantIndex(builder, loc, shape[r]));
   }
   // The dimSizes array.
   Value dimSizes =
       builder.create<memref::AllocOp>(loc, MemRefType::get({rank}, indexType));
   fields.push_back(dimSizes);
+  // The sizes array.
+  unsigned lastField = getFieldIndex(type, -1, -1);
+  Value memSizes = builder.create<memref::AllocOp>(
+      loc, MemRefType::get({lastField - 2}, indexType));
+  fields.push_back(memSizes);
   // Per-dimension storage.
   for (unsigned r = 0; r < rank; r++) {
     // Get the original dimension (ro) for the current stored dimension.
     unsigned ro = toOrig(enc, r);
     builder.create<memref::StoreOp>(loc, sizes[ro], dimSizes,
                                     constantIndex(builder, loc, r));
     linear = builder.create<arith::MulIOp>(loc, linear, sizes[ro]);
     // Allocate fiels.
case SparseTensorEncodingAttr::DimLevelType::SingletonNuNo:
       allDense = false;
       break;
     }
   }
   // The values array. For all-dense, the full length is required.
   // In all other case, we resort to the heuristical initial value.
   Value valuesSz = allDense ? linear : heuristic;
   fields.push_back(createAllocation(builder, loc, eltType, valuesSz));
+  // Set memSizes.
+  if (allDense)
+    builder.create<memref::StoreOp>(
+        loc, valuesSz, memSizes,
+        constantIndex(builder, loc, 0)); // TODO: avoid memSizes in this case?
+  else
+    builder.create<linalg::FillOp>(
+        loc, ValueRange{constantZero(builder, loc, indexType)},
+        ValueRange{memSizes});
+  assert(fields.size() == lastField);
 }
 
 //===----------------------------------------------------------------------===//
 // Codegen rules.
 //===----------------------------------------------------------------------===//
 
 /// Sparse tensor storage conversion rule for returns.
 class SparseReturnConverter : public OpConversionPattern<func::ReturnOp> {
if (op.getHasInserts()) {
       // Finalize any pending insertions.
       // TODO: implement
     }
     rewriter.replaceOp(op, adaptor.getOperands());
     return success();
   }
 };
 
-/// Base class for getter-like operations, e.g., to_indices, to_pointers.
-template <typename SourceOp, typename Base>
-class SparseGetterOpConverter : public OpConversionPattern<SourceOp> {
-public:
-  using OpAdaptor = typename SourceOp::Adaptor;
-  using OpConversionPattern<SourceOp>::OpConversionPattern;
-  LogicalResult
-  matchAndRewrite(SourceOp op, OpAdaptor adaptor,
-                  ConversionPatternRewriter &rewriter) const override {
-    // Replace the requested pointer access with corresponding field.
-    // The cast_op is inserted by type converter to intermix 1:N type
-    // conversion.
-    auto tuple = llvm::cast<UnrealizedConversionCastOp>(
-        adaptor.getTensor().getDefiningOp());
-    unsigned idx = Base::getIndexForOp(tuple, op);
-    auto fields = tuple.getInputs();
-    assert(idx < fields.size());
-    rewriter.replaceOp(op, fields[idx]);
-    return success();
-  }
-};
-
 /// Sparse codegen rule for the expand op.
 class SparseExpandConverter : public OpConversionPattern<ExpandOp> {
 public:
   using OpConversionPattern::OpConversionPattern;
   LogicalResult
   matchAndRewrite(ExpandOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     Location loc = op->getLoc();
rewriter.create<linalg::FillOp>(
         ValueRange{filled});
     // Replace expansion op with these buffers and initial index.
     assert(op.getNumResults() == 4);
     rewriter.replaceOp(op, {values, filled, indices, zero});
     return success();
   }
 };
 
+/// Base class for getter-like operations, e.g., to_indices, to_pointers.
+template <typename SourceOp, typename Base>
+class SparseGetterOpConverter : public OpConversionPattern<SourceOp> {
+public:
+  using OpAdaptor = typename SourceOp::Adaptor;
+  using OpConversionPattern<SourceOp>::OpConversionPattern;
+  LogicalResult
+  matchAndRewrite(SourceOp op, OpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override {
+    // Replace the requested pointer access with corresponding field.
+    // The cast_op is inserted by type converter to intermix 1:N type
+    // conversion.
+    auto tuple = llvm::cast<UnrealizedConversionCastOp>(
+        adaptor.getTensor().getDefiningOp());
+    unsigned idx = Base::getIndexForOp(tuple, op);
+    auto fields = tuple.getInputs();
+    assert(idx < fields.size());
+    rewriter.replaceOp(op, fields[idx]);
+    return success();
+  }
+};
+
 /// Sparse codegen rule for pointer accesses.
 class SparseToPointersConverter
     : public SparseGetterOpConverter<ToPointersOp, SparseToPointersConverter> {
 public:
   using SparseGetterOpConverter::SparseGetterOpConverter;
   // Callback for SparseGetterOpConverter.
   static unsigned getIndexForOp(UnrealizedConversionCastOp /*tuple*/,
                                 ToPointersOp op) {

 //===----------------------------------------------------------------------===//
 // Public method for populating conversion rules.
 //===----------------------------------------------------------------------===//
 
 /// Populates the given patterns list with conversion rules required for
 /// the sparsification of linear algebra operations.
 void mlir::populateSparseTensorCodegenPatterns(TypeConverter &typeConverter,
                                                RewritePatternSet &patterns) {
   patterns.add<SparseReturnConverter, SparseCallConverter, SparseDimOpConverter,
-               SparseCastConverter, SparseExpandConverter,
-               SparseTensorAllocConverter, SparseTensorDeallocConverter,
-               SparseToPointersConverter, SparseToIndicesConverter,
-               SparseToValuesConverter, SparseTensorLoadConverter>(
+               SparseCastConverter, SparseTensorAllocConverter,
+               SparseTensorDeallocConverter, SparseTensorLoadConverter,
+               SparseExpandConverter, SparseToPointersConverter,
+               SparseToIndicesConverter, SparseToValuesConverter>(
       typeConverter, patterns.getContext());
 }

mlir/test/Dialect/SparseTensor/codegen.mlir

 
 #Dense3D = #sparse_tensor.encoding<{
   dimLevelType = [ "dense", "dense", "dense" ],
   dimOrdering = affine_map<(i, j, k) -> (k, i, j)>
 }>
 
 // CHECK-LABEL: func @sparse_nop(
 //  CHECK-SAME: %[[A0:.*0]]: memref<1xindex>,
-//  CHECK-SAME: %[[A1:.*1]]: memref<?xi32>,
-//  CHECK-SAME: %[[A2:.*2]]: memref<?xi64>,
-//  CHECK-SAME: %[[A3:.*3]]: memref<?xf64>)
-//       CHECK: return %[[A0]], %[[A1]], %[[A2]], %[[A3]] : memref<1xindex>, memref<?xi32>, memref<?xi64>, memref<?xf64>
+//  CHECK-SAME: %[[A1:.*1]]: memref<3xindex>,
+//  CHECK-SAME: %[[A2:.*2]]: memref<?xi32>,
+//  CHECK-SAME: %[[A3:.*3]]: memref<?xi64>,
+//  CHECK-SAME: %[[A4:.*4]]: memref<?xf64>)
+//       CHECK: return %[[A0]], %[[A1]], %[[A2]], %[[A3]], %[[A4]] : memref<1xindex>, memref<3xindex>, memref<?xi32>, memref<?xi64>, memref<?xf64>
 func.func @sparse_nop(%arg0: tensor<?xf64, #SparseVector>) -> tensor<?xf64, #SparseVector> {
   return %arg0 : tensor<?xf64, #SparseVector>
 }
 
 // CHECK-LABEL: func @sparse_nop_multi_ret(
 //  CHECK-SAME: %[[A0:.*0]]: memref<1xindex>,
-//  CHECK-SAME: %[[A1:.*1]]: memref<?xi32>,
-//  CHECK-SAME: %[[A2:.*2]]: memref<?xi64>,
-//  CHECK-SAME: %[[A3:.*3]]: memref<?xf64>,
-//  CHECK-SAME: %[[A4:.*4]]: memref<1xindex>,
-//  CHECK-SAME: %[[A5:.*5]]: memref<?xi32>,
-//  CHECK-SAME: %[[A6:.*6]]: memref<?xi64>,
-//  CHECK-SAME: %[[A7:.*7]]: memref<?xf64>) ->
-//       CHECK: return %[[A0]], %[[A1]], %[[A2]], %[[A3]], %[[A4]], %[[A5]], %[[A6]], %[[A7]]
+//  CHECK-SAME: %[[A1:.*1]]: memref<3xindex>,
+//  CHECK-SAME: %[[A2:.*2]]: memref<?xi32>,
+//  CHECK-SAME: %[[A3:.*3]]: memref<?xi64>,
+//  CHECK-SAME: %[[A4:.*4]]: memref<?xf64>,
+//  CHECK-SAME: %[[A5:.*5]]: memref<1xindex>,
+//  CHECK-SAME: %[[A6:.*6]]: memref<3xindex>,
+//  CHECK-SAME: %[[A7:.*7]]: memref<?xi32>,
+//  CHECK-SAME: %[[A8:.*8]]: memref<?xi64>,
+//  CHECK-SAME: %[[A9:.*9]]: memref<?xf64>) ->
+//       CHECK: return %[[A0]], %[[A1]], %[[A2]], %[[A3]], %[[A4]], %[[A5]], %[[A6]], %[[A7]], %[[A8]], %[[A9]]
 func.func @sparse_nop_multi_ret(%arg0: tensor<?xf64, #SparseVector>,
                                 %arg1: tensor<?xf64, #SparseVector>) ->
                                 (tensor<?xf64, #SparseVector>, tensor<?xf64, #SparseVector>) {
   return %arg0, %arg1 : tensor<?xf64, #SparseVector>, tensor<?xf64, #SparseVector>
 }
 
 // CHECK-LABEL: func @sparse_nop_call(
 //  CHECK-SAME: %[[A0:.*0]]: memref<1xindex>,
-//  CHECK-SAME: %[[A1:.*1]]: memref<?xi32>,
-//  CHECK-SAME: %[[A2:.*2]]: memref<?xi64>,
-//  CHECK-SAME: %[[A3:.*3]]: memref<?xf64>,
-//  CHECK-SAME: %[[A4:.*4]]: memref<1xindex>,
-//  CHECK-SAME: %[[A5:.*5]]: memref<?xi32>,
-//  CHECK-SAME: %[[A6:.*6]]: memref<?xi64>,
-//  CHECK-SAME: %[[A7:.*7]]: memref<?xf64>) 
-//       CHECK: %[[T0:.*]]:8 = call @sparse_nop_multi_ret(%[[A0]], %[[A1]], %[[A2]], %[[A3]], %[[A4]], %[[A5]], %[[A6]], %[[A7]]) 
-//       CHECK: return %[[T0]]#0, %[[T0]]#1, %[[T0]]#2, %[[T0]]#3, %[[T0]]#4, %[[T0]]#5, %[[T0]]#6, %[[T0]]#7 
+//  CHECK-SAME: %[[A1:.*1]]: memref<3xindex>,
+//  CHECK-SAME: %[[A2:.*2]]: memref<?xi32>,
+//  CHECK-SAME: %[[A3:.*3]]: memref<?xi64>,
+//  CHECK-SAME: %[[A4:.*4]]: memref<?xf64>,
+//  CHECK-SAME: %[[A5:.*5]]: memref<1xindex>,
+//  CHECK-SAME: %[[A6:.*6]]: memref<3xindex>,
+//  CHECK-SAME: %[[A7:.*7]]: memref<?xi32>,
+//  CHECK-SAME: %[[A8:.*8]]: memref<?xi64>,
+//  CHECK-SAME: %[[A9:.*9]]: memref<?xf64>)
+//       CHECK: %[[T0:.*]]:10 = call @sparse_nop_multi_ret(%[[A0]], %[[A1]], %[[A2]], %[[A3]], %[[A4]], %[[A5]], %[[A6]], %[[A7]], %[[A8]], %[[A9]])
+//       CHECK: return %[[T0]]#0, %[[T0]]#1, %[[T0]]#2, %[[T0]]#3, %[[T0]]#4, %[[T0]]#5, %[[T0]]#6, %[[T0]]#7, %[[T0]]#8, %[[T0]]#9
 func.func @sparse_nop_call(%arg0: tensor<?xf64, #SparseVector>,
                            %arg1: tensor<?xf64, #SparseVector>) ->
                            (tensor<?xf64, #SparseVector>, tensor<?xf64, #SparseVector>) {
   %1, %2 = call @sparse_nop_multi_ret(%arg0, %arg1) :
                            (tensor<?xf64, #SparseVector>, tensor<?xf64, #SparseVector>) ->
                            (tensor<?xf64, #SparseVector>, tensor<?xf64, #SparseVector>)
   return %1, %2: tensor<?xf64, #SparseVector>, tensor<?xf64, #SparseVector>
 }
 
-//
 // CHECK-LABEL: func @sparse_nop_cast(
 //  CHECK-SAME: %[[A0:.*0]]: memref<1xindex>,
-//  CHECK-SAME: %[[A1:.*1]]: memref<?xi32>,
-//  CHECK-SAME: %[[A2:.*2]]: memref<?xi64>,
-//  CHECK-SAME: %[[A3:.*3]]: memref<?xf32>)
-//       CHECK: return %[[A0]], %[[A1]], %[[A2]], %[[A3]] : memref<1xindex>, memref<?xi32>, memref<?xi64>, memref<?xf32>
+//  CHECK-SAME: %[[A1:.*1]]: memref<3xindex>,
+//  CHECK-SAME: %[[A2:.*2]]: memref<?xi32>,
+//  CHECK-SAME: %[[A3:.*3]]: memref<?xi64>,
+//  CHECK-SAME: %[[A4:.*4]]: memref<?xf32>)
+//       CHECK: return %[[A0]], %[[A1]], %[[A2]], %[[A3]], %[[A4]] : memref<1xindex>, memref<3xindex>, memref<?xi32>, memref<?xi64>, memref<?xf32>
 func.func @sparse_nop_cast(%arg0: tensor<64xf32, #SparseVector>) -> tensor<?xf32, #SparseVector> {
   %0 = tensor.cast %arg0 : tensor<64xf32, #SparseVector> to tensor<?xf32, #SparseVector>
   return %0 : tensor<?xf32, #SparseVector>
 }
 
-//
 // CHECK-LABEL: func @sparse_nop_cast_3d(
 //  CHECK-SAME: %[[A0:.*0]]: memref<3xindex>,
-//  CHECK-SAME: %[[A1:.*1]]: memref<?xf32>)
-//       CHECK: return %[[A0]], %[[A1]] : memref<3xindex>, memref<?xf32>
+//  CHECK-SAME: %[[A1:.*1]]: memref<1xindex>,
+//  CHECK-SAME: %[[A2:.*2]]: memref<?xf32>)
+//       CHECK: return %[[A0]], %[[A1]], %[[A2]] : memref<3xindex>, memref<1xindex>, memref<?xf32>
 func.func @sparse_nop_cast_3d(%arg0: tensor<10x20x30xf32, #Dense3D>) -> tensor<?x?x?xf32, #Dense3D> {
   %0 = tensor.cast %arg0 : tensor<10x20x30xf32, #Dense3D> to tensor<?x?x?xf32, #Dense3D>
   return %0 : tensor<?x?x?xf32, #Dense3D>
 }
 
-//
 // CHECK-LABEL: func @sparse_dense_2d(
 //  CHECK-SAME: %[[A0:.*0]]: memref<2xindex>,
-//  CHECK-SAME: %[[A1:.*1]]: memref<?xf64>) {
+//  CHECK-SAME: %[[A1:.*1]]: memref<1xindex>,
+//  CHECK-SAME: %[[A2:.*2]]: memref<?xf64>)
 //       CHECK: return
 func.func @sparse_dense_2d(%arg0: tensor<?x?xf64, #Dense2D>) {
   return
 }
 
-//
 // CHECK-LABEL: func @sparse_row(
 //  CHECK-SAME: %[[A0:.*0]]: memref<2xindex>,
-//  CHECK-SAME: %[[A1:.*1]]: memref<?xi32>,
-//  CHECK-SAME: %[[A2:.*2]]: memref<?xi64>,
-//  CHECK-SAME: %[[A3:.*3]]: memref<?xf64>) {
+//  CHECK-SAME: %[[A1:.*1]]: memref<3xindex>,
+//  CHECK-SAME: %[[A2:.*2]]: memref<?xi32>,
+//  CHECK-SAME: %[[A3:.*3]]: memref<?xi64>,
+//  CHECK-SAME: %[[A4:.*4]]: memref<?xf64>)
 //       CHECK: return
 func.func @sparse_row(%arg0: tensor<?x?xf64, #Row>) {
   return
 }
 
-//
 // CHECK-LABEL: func @sparse_csr(
 //  CHECK-SAME: %[[A0:.*0]]: memref<2xindex>,
-//  CHECK-SAME: %[[A1:.*1]]: memref<?xi32>,
-//  CHECK-SAME: %[[A2:.*2]]: memref<?xi64>,
-//  CHECK-SAME: %[[A3:.*3]]: memref<?xf64>) {
+//  CHECK-SAME: %[[A1:.*1]]: memref<3xindex>,
+//  CHECK-SAME: %[[A2:.*2]]: memref<?xi32>,
+//  CHECK-SAME: %[[A3:.*3]]: memref<?xi64>,
+//  CHECK-SAME: %[[A4:.*4]]: memref<?xf64>)
 //       CHECK: return
 func.func @sparse_csr(%arg0: tensor<?x?xf64, #CSR>) {
   return
 }
 
-//
 // CHECK-LABEL: func @sparse_dcsr(
 //  CHECK-SAME: %[[A0:.*0]]: memref<2xindex>,
-//  CHECK-SAME: %[[A1:.*1]]: memref<?xi32>,
-//  CHECK-SAME: %[[A2:.*2]]: memref<?xi64>,
-//  CHECK-SAME: %[[A3:.*3]]: memref<?xi32>,
-//  CHECK-SAME: %[[A4:.*4]]: memref<?xi64>,
-//  CHECK-SAME: %[[A5:.*5]]: memref<?xf64>) {
+//  CHECK-SAME: %[[A1:.*1]]: memref<5xindex>,
+//  CHECK-SAME: %[[A2:.*2]]: memref<?xi32>,
+//  CHECK-SAME: %[[A3:.*3]]: memref<?xi64>,
+//  CHECK-SAME: %[[A4:.*4]]: memref<?xi32>,
+//  CHECK-SAME: %[[A5:.*5]]: memref<?xi64>,
+//  CHECK-SAME: %[[A6:.*6]]: memref<?xf64>)
 //       CHECK: return
 func.func @sparse_dcsr(%arg0: tensor<?x?xf64, #DCSR>) {
   return
 }
 
 //
 // Querying for dimension 1 in the tensor type can immediately
 // fold using the original static dimension sizes.
 //
-//
 // CHECK-LABEL: func @sparse_dense_3d(
 //  CHECK-SAME: %[[A0:.*0]]: memref<3xindex>,
-//  CHECK-SAME: %[[A1:.*1]]: memref<?xf64>)
+//  CHECK-SAME: %[[A1:.*1]]: memref<1xindex>,
+//  CHECK-SAME: %[[A2:.*2]]: memref<?xf64>)
 //       CHECK: %[[C:.*]] = arith.constant 20 : index
 //       CHECK: return %[[C]] : index
 func.func @sparse_dense_3d(%arg0: tensor<10x20x30xf64, #Dense3D>) -> index {
   %c = arith.constant 1 : index
   %0 = tensor.dim %arg0, %c : tensor<10x20x30xf64, #Dense3D>
   return %0 : index
 }
 
 //
 // Querying for dimension 1 in the tensor type needs to be permuted
 // into querying for dimension 2 in the stored sparse tensor scheme,
 // since the latter honors the dimOrdering.
 //
-//
 // CHECK-LABEL: func @sparse_dense_3d_dyn(
 //  CHECK-SAME: %[[A0:.*0]]: memref<3xindex>,
-//  CHECK-SAME: %[[A1:.*1]]: memref<?xf64>)
+//  CHECK-SAME: %[[A1:.*1]]: memref<1xindex>,
+//  CHECK-SAME: %[[A2:.*2]]: memref<?xf64>)
 //       CHECK: %[[C:.*]] = arith.constant 2 : index
 //       CHECK: %[[L:.*]] = memref.load %[[A0]][%[[C]]] : memref<3xindex>
 //       CHECK: return %[[L]] : index
 func.func @sparse_dense_3d_dyn(%arg0: tensor<?x?x?xf64, #Dense3D>) -> index {
   %c = arith.constant 1 : index
   %0 = tensor.dim %arg0, %c : tensor<?x?x?xf64, #Dense3D>
   return %0 : index
 }
 
-//
 // CHECK-LABEL: func @sparse_pointers_dcsr(
 //  CHECK-SAME: %[[A0:.*0]]: memref<2xindex>,
-//  CHECK-SAME: %[[A1:.*1]]: memref<?xi32>,
-//  CHECK-SAME: %[[A2:.*2]]: memref<?xi64>,
-//  CHECK-SAME: %[[A3:.*3]]: memref<?xi32>,
-//  CHECK-SAME: %[[A4:.*4]]: memref<?xi64>,
-//  CHECK-SAME: %[[A5:.*5]]: memref<?xf64>)
-//       CHECK: return %[[A3]] : memref<?xi32>
+//  CHECK-SAME: %[[A1:.*1]]: memref<5xindex>,
+//  CHECK-SAME: %[[A2:.*2]]: memref<?xi32>,
+//  CHECK-SAME: %[[A3:.*3]]: memref<?xi64>,
+//  CHECK-SAME: %[[A4:.*4]]: memref<?xi32>,
+//  CHECK-SAME: %[[A5:.*5]]: memref<?xi64>,
+//  CHECK-SAME: %[[A6:.*6]]: memref<?xf64>)
+//       CHECK: return %[[A4]] : memref<?xi32>
 func.func @sparse_pointers_dcsr(%arg0: tensor<?x?xf64, #DCSR>) -> memref<?xi32> {
   %0 = sparse_tensor.pointers %arg0 { dimension = 1 : index } : tensor<?x?xf64, #DCSR> to memref<?xi32>
   return %0 : memref<?xi32>
 }
 
-//
 // CHECK-LABEL: func @sparse_indices_dcsr(
 //  CHECK-SAME: %[[A0:.*0]]: memref<2xindex>,
-//  CHECK-SAME: %[[A1:.*1]]: memref<?xi32>,
-//  CHECK-SAME: %[[A2:.*2]]: memref<?xi64>,
-//  CHECK-SAME: %[[A3:.*3]]: memref<?xi32>,
-//  CHECK-SAME: %[[A4:.*4]]: memref<?xi64>,
-//  CHECK-SAME: %[[A5:.*5]]: memref<?xf64>)
-//       CHECK: return %[[A4]] : memref<?xi64>
+//  CHECK-SAME: %[[A1:.*1]]: memref<5xindex>,
+//  CHECK-SAME: %[[A2:.*2]]: memref<?xi32>,
+//  CHECK-SAME: %[[A3:.*3]]: memref<?xi64>,
+//  CHECK-SAME: %[[A4:.*4]]: memref<?xi32>,
+//  CHECK-SAME: %[[A5:.*5]]: memref<?xi64>,
+//  CHECK-SAME: %[[A6:.*6]]: memref<?xf64>)
+//       CHECK: return %[[A5]] : memref<?xi64>
 func.func @sparse_indices_dcsr(%arg0: tensor<?x?xf64, #DCSR>) -> memref<?xi64> {
   %0 = sparse_tensor.indices %arg0 { dimension = 1 : index } : tensor<?x?xf64, #DCSR> to memref<?xi64>
   return %0 : memref<?xi64>
 }
 
-//
 // CHECK-LABEL: func @sparse_values_dcsr(
 //  CHECK-SAME: %[[A0:.*0]]: memref<2xindex>,
-//  CHECK-SAME: %[[A1:.*1]]: memref<?xi32>,
-//  CHECK-SAME: %[[A2:.*2]]: memref<?xi64>,
-//  CHECK-SAME: %[[A3:.*3]]: memref<?xi32>,
-//  CHECK-SAME: %[[A4:.*4]]: memref<?xi64>,
-//  CHECK-SAME: %[[A5:.*5]]: memref<?xf64>)
-//       CHECK: return %[[A5]] : memref<?xf64>
+//  CHECK-SAME: %[[A1:.*1]]: memref<5xindex>,
+//  CHECK-SAME: %[[A2:.*2]]: memref<?xi32>,
+//  CHECK-SAME: %[[A3:.*3]]: memref<?xi64>,
+//  CHECK-SAME: %[[A4:.*4]]: memref<?xi32>,
+//  CHECK-SAME: %[[A5:.*5]]: memref<?xi64>,
+//  CHECK-SAME: %[[A6:.*6]]: memref<?xf64>)
+//       CHECK: return %[[A6]] : memref<?xf64>
 func.func @sparse_values_dcsr(%arg0: tensor<?x?xf64, #DCSR>) -> memref<?xf64> {
   %0 = sparse_tensor.values %arg0 : tensor<?x?xf64, #DCSR> to memref<?xf64>
   return %0 : memref<?xf64>
 }
 
-//
 // CHECK-LABEL: func @sparse_dealloc_csr(
 //  CHECK-SAME: %[[A0:.*0]]: memref<2xindex>,
-//  CHECK-SAME: %[[A1:.*1]]: memref<?xi32>,
-//  CHECK-SAME: %[[A2:.*2]]: memref<?xi64>,
-//  CHECK-SAME: %[[A3:.*3]]: memref<?xf64>) {
+//  CHECK-SAME: %[[A1:.*1]]: memref<3xindex>,
+//  CHECK-SAME: %[[A2:.*2]]: memref<?xi32>,
+//  CHECK-SAME: %[[A3:.*3]]: memref<?xi64>,
+//  CHECK-SAME: %[[A4:.*4]]: memref<?xf64>)
 //       CHECK: memref.dealloc %[[A0]] : memref<2xindex>
-//       CHECK: memref.dealloc %[[A1]] : memref<?xi32>
-//       CHECK: memref.dealloc %[[A2]] : memref<?xi64>
-//       CHECK: memref.dealloc %[[A3]] : memref<?xf64>
+//       CHECK: memref.dealloc %[[A1]] : memref<3xindex>
+//       CHECK: memref.dealloc %[[A2]] : memref<?xi32>
+//       CHECK: memref.dealloc %[[A3]] : memref<?xi64>
+//       CHECK: memref.dealloc %[[A4]] : memref<?xf64>
 //       CHECK: return
 func.func @sparse_dealloc_csr(%arg0: tensor<?x?xf64, #CSR>) {
   bufferization.dealloc_tensor %arg0 : tensor<?x?xf64, #CSR>
   return
 }
 
-//        CHECK-LABEL: func @sparse_alloc_csc(
-//         CHECK-SAME: %[[A:.*]]: index) ->
-//         CHECK-SAME: memref<2xindex>, memref<?xindex>, memref<?xindex>, memref<?xf64>
-//          CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
-//          CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index
-//          CHECK-DAG: %[[C10:.*]] = arith.constant 10 : index
-//              CHECK: %[[T0:.*]] = memref.alloc() : memref<2xindex>
-//              CHECK: memref.store %[[A]], %[[T0]][%[[C0]]] : memref<2xindex>
-//              CHECK: memref.store %[[C10]], %[[T0]][%[[C1]]] : memref<2xindex>
-//              CHECK: %[[T1:.*]] = memref.alloc() : memref<1xindex>
-//              CHECK: %[[T2:.*]] = memref.cast %[[T1]] : memref<1xindex> to memref<?xindex>
-//              CHECK: %[[T3:.*]] = memref.alloc() : memref<1xindex>
-//              CHECK: %[[T4:.*]] = memref.cast %[[T3]] : memref<1xindex> to memref<?xindex>
-//              CHECK: %[[T5:.*]] = memref.alloc() : memref<1xf64>
-//              CHECK: %[[T6:.*]] = memref.cast %[[T5]] : memref<1xf64> to memref<?xf64>
-//              CHECK: return %[[T0]], %[[T2]], %[[T4]], %[[T6]]
+// CHECK-LABEL: func @sparse_alloc_csc(
+//  CHECK-SAME: %[[A:.*]]: index) ->
+//  CHECK-SAME: memref<2xindex>, memref<3xindex>, memref<?xindex>, memref<?xindex>, memref<?xf64>
+//   CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
+//   CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index
+//   CHECK-DAG: %[[C10:.*]] = arith.constant 10 : index
+//       CHECK: %[[T0:.*]] = memref.alloc() : memref<2xindex>
+//       CHECK: %[[T1:.*]] = memref.alloc() : memref<3xindex>
+//       CHECK: memref.store %[[A]], %[[T0]][%[[C0]]] : memref<2xindex>
+//       CHECK: memref.store %[[C10]], %[[T0]][%[[C1]]] : memref<2xindex>
+//       CHECK: %[[T2:.*]] = memref.alloc() : memref<1xindex>
+//       CHECK: %[[T3:.*]] = memref.cast %[[T2]] : memref<1xindex> to memref<?xindex>
+//       CHECK: %[[T4:.*]] = memref.alloc() : memref<1xindex>
+//       CHECK: %[[T5:.*]] = memref.cast %[[T4]] : memref<1xindex> to memref<?xindex>
+//       CHECK: %[[T6:.*]] = memref.alloc() : memref<1xf64>
+//       CHECK: %[[T7:.*]] = memref.cast %[[T6]] : memref<1xf64> to memref<?xf64>
+//       CHECK: linalg.fill ins(%[[C0]] : index) outs(%[[T1]] : memref<3xindex>)
+//       CHECK: return %[[T0]], %[[T1]], %[[T3]], %[[T5]], %[[T7]]
 func.func @sparse_alloc_csc(%arg0: index) -> tensor<10x?xf64, #CSC> {
   %0 = bufferization.alloc_tensor(%arg0) : tensor<10x?xf64, #CSC>
   %1 = sparse_tensor.load %0 : tensor<10x?xf64, #CSC>
   return %1 : tensor<10x?xf64, #CSC>
 }
 
-//        CHECK-LABEL: func @sparse_alloc_3d() ->
-//         CHECK-SAME: memref<3xindex>, memref<?xf64>
-//          CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
-//          CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index
-//          CHECK-DAG: %[[C2:.*]] = arith.constant 2 : index
-//          CHECK-DAG: %[[C10:.*]] = arith.constant 10 : index
-//          CHECK-DAG: %[[C20:.*]] = arith.constant 20 : index
-//          CHECK-DAG: %[[C30:.*]] = arith.constant 30 : index
-//              CHECK: %[[A0:.*]] = memref.alloc() : memref<3xindex>
-//              CHECK: memref.store %[[C30]], %[[A0]][%[[C0]]] : memref<3xindex>
-//              CHECK: memref.store %[[C10]], %[[A0]][%[[C1]]] : memref<3xindex>
-//              CHECK: memref.store %[[C20]], %[[A0]][%[[C2]]] : memref<3xindex>
-//              CHECK: %[[A:.*]] = memref.alloc() : memref<6000xf64>
-//              CHECK: %[[A1:.*]] = memref.cast %[[A]] : memref<6000xf64> to memref<?xf64>
-//              CHECK: return %[[A0]], %[[A1]] : memref<3xindex>, memref<?xf64>
+// CHECK-LABEL: func @sparse_alloc_3d() ->
+//  CHECK-SAME: memref<3xindex>, memref<1xindex>, memref<?xf64>
+//   CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
+//   CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index
+//   CHECK-DAG: %[[C2:.*]] = arith.constant 2 : index
+//   CHECK-DAG: %[[C10:.*]] = arith.constant 10 : index
+//   CHECK-DAG: %[[C20:.*]] = arith.constant 20 : index
+//   CHECK-DAG: %[[C30:.*]] = arith.constant 30 : index
+//   CHECK-DAG: %[[C6000:.*]] = arith.constant 6000 : index
+//       CHECK: %[[A0:.*]] = memref.alloc() : memref<3xindex>
+//       CHECK: %[[A1:.*]] = memref.alloc() : memref<1xindex>
+//       CHECK: memref.store %[[C30]], %[[A0]][%[[C0]]] : memref<3xindex>
+//       CHECK: memref.store %[[C10]], %[[A0]][%[[C1]]] : memref<3xindex>
+//       CHECK: memref.store %[[C20]], %[[A0]][%[[C2]]] : memref<3xindex>
+//       CHECK: %[[A:.*]] = memref.alloc() : memref<6000xf64>
+//       CHECK: %[[A2:.*]] = memref.cast %[[A]] : memref<6000xf64> to memref<?xf64>
+//       CHECK: memref.store %[[C6000]], %[[A1]][%[[C0]]] : memref<1xindex>
+//       CHECK: return %[[A0]], %[[A1]], %[[A2]] : memref<3xindex>, memref<1xindex>, memref<?xf64>
 func.func @sparse_alloc_3d() -> tensor<10x20x30xf64, #Dense3D> {
   %0 = bufferization.alloc_tensor() : tensor<10x20x30xf64, #Dense3D>
   %1 = sparse_tensor.load %0 : tensor<10x20x30xf64, #Dense3D>
   return %1 : tensor<10x20x30xf64, #Dense3D>
 }
 
-//   CHECK-LABEL: func.func @sparse_expansion1()
-//         CHECK: %[[A:.*]] = memref.alloc() : memref<8xf64>
-//         CHECK: %[[B:.*]] = memref.alloc() : memref<8xi1>
-//         CHECK: %[[C:.*]] = memref.alloc() : memref<8xindex>
-//         CHECK: %[[D:.*]] = memref.cast %[[C]] : memref<8xindex> to memref<?xindex>
-//     CHECK-DAG: linalg.fill ins(%{{.*}}  : f64) outs(%[[A]] : memref<8xf64>)
-//     CHECK-DAG: linalg.fill ins(%{{.*}}  : i1) outs(%[[B]] : memref<8xi1>)
-//         CHECK: return %[[D]] : memref<?xindex>
+// CHECK-LABEL: func.func @sparse_expansion1()
+//       CHECK: %[[A:.*]] = memref.alloc() : memref<8xf64>
+//       CHECK: %[[B:.*]] = memref.alloc() : memref<8xi1>
+//       CHECK: %[[C:.*]] = memref.alloc() : memref<8xindex>
+//       CHECK: %[[D:.*]] = memref.cast %[[C]] : memref<8xindex> to memref<?xindex>
+//   CHECK-DAG: linalg.fill ins(%{{.*}}  : f64) outs(%[[A]] : memref<8xf64>)
+//   CHECK-DAG: linalg.fill ins(%{{.*}}  : i1) outs(%[[B]] : memref<8xi1>)
+//       CHECK: return %[[D]] : memref<?xindex>
 func.func @sparse_expansion1() -> memref<?xindex> {
   %0 = bufferization.alloc_tensor() : tensor<4x8xf64, #CSR>
   %values, %filled, %added, %count = sparse_tensor.expand %0
     : tensor<4x8xf64, #CSR> to memref<?xf64>, memref<?xi1>, memref<?xindex>, index
   return %added : memref<?xindex>
 }
 
-//   CHECK-LABEL: func.func @sparse_expansion2()
-//         CHECK: %[[A:.*]] = memref.alloc() : memref<4xf64>
-//         CHECK: %[[B:.*]] = memref.alloc() : memref<4xi1>
-//         CHECK: %[[C:.*]] = memref.alloc() : memref<4xindex>
-//         CHECK: %[[D:.*]] = memref.cast %[[C]] : memref<4xindex> to memref<?xindex>
-//     CHECK-DAG: linalg.fill ins(%{{.*}}  : f64) outs(%[[A]] : memref<4xf64>)
-//     CHECK-DAG: linalg.fill ins(%{{.*}}  : i1) outs(%[[B]] : memref<4xi1>)
-//         CHECK: return %[[D]] : memref<?xindex>
+// CHECK-LABEL: func.func @sparse_expansion2()
+//       CHECK: %[[A:.*]] = memref.alloc() : memref<4xf64>
+//       CHECK: %[[B:.*]] = memref.alloc() : memref<4xi1>
+//       CHECK: %[[C:.*]] = memref.alloc() : memref<4xindex>
+//       CHECK: %[[D:.*]] = memref.cast %[[C]] : memref<4xindex> to memref<?xindex>
+//   CHECK-DAG: linalg.fill ins(%{{.*}}  : f64) outs(%[[A]] : memref<4xf64>)
+//   CHECK-DAG: linalg.fill ins(%{{.*}}  : i1) outs(%[[B]] : memref<4xi1>)
+//       CHECK: return %[[D]] : memref<?xindex>
 func.func @sparse_expansion2() -> memref<?xindex> {
   %0 = bufferization.alloc_tensor() : tensor<4x8xf64, #CSC>
   %values, %filled, %added, %count = sparse_tensor.expand %0
     : tensor<4x8xf64, #CSC> to memref<?xf64>, memref<?xi1>, memref<?xindex>, index
   return %added : memref<?xindex>
 }
 
-//   CHECK-LABEL: func.func @sparse_expansion3(
-//    CHECK-SAME: %[[D0:.*]]: index,
-//    CHECK-SAME: %{{.*}}: index) -> memref<?xindex> {
-//         CHECK: %[[C1:.*]] = arith.constant 1 : index
-//         CHECK: %[[S0:.*]] = memref.alloc() : memref<2xindex>
-//         CHECK: memref.store %[[D0]], %[[S0]]{{\[}}%[[C1]]] : memref<2xindex>
-//         CHECK: %[[D1:.*]] = memref.load %[[S0]]{{\[}}%[[C1]]] : memref<2xindex>
-//         CHECK: %[[V:.*]] = memref.alloc(%[[D1]]) : memref<?xf64>
-//         CHECK: %[[B:.*]] = memref.alloc(%[[D1]]) : memref<?xi1>
-//         CHECK: %[[D:.*]] = memref.alloc(%[[D1]]) : memref<?xindex>
-//         CHECK: linalg.fill ins(%{{.*}} : f64) outs(%[[V]] : memref<?xf64>)
-//         CHECK: linalg.fill ins(%{{.*}} : i1) outs(%[[B]] : memref<?xi1>)
-//         CHECK: return %[[D]] : memref<?xindex>
+// CHECK-LABEL: func.func @sparse_expansion3(
+//  CHECK-SAME: %[[D0:.*]]: index,
+//  CHECK-SAME: %{{.*}}: index) -> memref<?xindex> {
+//       CHECK: %[[C1:.*]] = arith.constant 1 : index
+//       CHECK: %[[S0:.*]] = memref.alloc() : memref<2xindex>
+//       CHECK: memref.store %[[D0]], %[[S0]]{{\[}}%[[C1]]] : memref<2xindex>
+//       CHECK: %[[D1:.*]] = memref.load %[[S0]]{{\[}}%[[C1]]] : memref<2xindex>
+//       CHECK: %[[V:.*]] = memref.alloc(%[[D1]]) : memref<?xf64>
+//       CHECK: %[[B:.*]] = memref.alloc(%[[D1]]) : memref<?xi1>
+//       CHECK: %[[D:.*]] = memref.alloc(%[[D1]]) : memref<?xindex>
+//       CHECK: linalg.fill ins(%{{.*}} : f64) outs(%[[V]] : memref<?xf64>)
+//       CHECK: linalg.fill ins(%{{.*}} : i1) outs(%[[B]] : memref<?xi1>)
+//       CHECK: return %[[D]] : memref<?xindex>
 func.func @sparse_expansion3(%arg0: index, %arg1: index) -> memref<?xindex> {
   %0 = bufferization.alloc_tensor(%arg0, %arg1) : tensor<?x?xf64, #CSC>
   %values, %filled, %added, %count = sparse_tensor.expand %0
     : tensor<?x?xf64, #CSC> to memref<?xf64>, memref<?xi1>, memref<?xindex>, index
   return %added : memref<?xindex>
 }