[mlir][sparse] lower number of entries op to actual code

[aartbik]

Summary

works both along runtime path and pure codegen path

[wrengr]

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

760

Is there any better way to abstract over this? (to avoid magic numbers / ensure it stays in sync with the struct definition)

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

952–956

Should factor out a genSparseValuesCall function, to deduplicate this code and the version in SparseTensorToValuesConverter::matchAndRewrite

[aartbik]

rebased against main

[aartbik]

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

760

Yeah, that is a good suggestion, since the -2 has bitten me before. Added a general util for MemSizes field

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

952–956

That is really just two lines of code and will break the nice symmetry with the other "getters".
Is it worth it?

[aartbik]

added dimSizes index util (to avoid magic subtractions all over the place)

[aartbik]

dimSizes -> memSizes of course

[wrengr]

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

952–956

The biggest reason to share it is to avoid the "magic constants" problem re the 15 and "sparseValues". Most of the other genFooCall functions are similarly short, but exist to address the same maintenance concerns. Perhaps there won't be any other places that use it in the future, but whenever I see the same magic constants more than once that's a warning flag for me, especially when they're strings :)

I don't quite follow what you mean about symmetry with other getters, but this would give symmetry with the other genFooCall functions.

[aartbik]

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

952–956

Ok (by symmetric I meant that all SparseTensorToXXXConverters looked very similar and this would break that, but not a big deal).

The <15> is a good point [is there a templating way of getting rid of the need for that, btw, it seems we could strlen somehow at compile-time?!

[aartbik]

refactored out getValues call

[wrengr]

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

952–956

by symmetric I meant that all SparseTensorToXXXConverters

Gotcha. Yeah, when looking over the code a while back for the Enums refactoring, I saw a few places where the setup for createFuncCall wasn't abstracted into a genFooCall wrapper; but since they were the single place the given function was called, I didn't bother writing up a CL to add the wrappers (because definitely not worth it if there's only one callsite :)

is there a templating way of getting rid of the need for that, btw, it seems we could strlen somehow at compile-time?!

It's doable, but I'm not sure it's worth the code it'd take to do it. Last time I thought about it, the simplest way I could come up with was to store the base-name string literal as a global constexpr, and then have some macro for threading that constexpr into both the size and the value; i.e., something like:

static constexpr const char *kFooBaseFunctionName = "foo";
static constexpr unsigned kMaxLenPrimaryTypeFunctionSuffix = 3;
...
# define PRIMARYTYPE_FUNCTION_NAME(VAR, BASE, ETP) \
    SmallString<std::strlen(BASE) + kMaxLenPrimaryTypeFunctionSuffix> VAR{BASE, primaryTypeFunctionSuffix(ETP)};
...
static Value genFooCall(...) {
    PRIMARYTYPE_FUNCTION_NAME(name, kFooBaseFunctionName, elemTp)
    ...
}

Of course we'd also need to define our own version of strlen since std::strlen isn't constexpr. (But that's easy enough to do, since this is compile-time and the strings are short, therefore performance isn't much of a concern.)

We could probably convert the macro into a proper template by using tricks similar to https://en.cppreference.com/w/cpp/types/integral_constant; i.e., so we could define the moral equivalent of:

template<string BaseName>
class PrimaryTypeFunctionName {
  SmallString<strlen(BaseName) + kMaxLenPrimaryTypeFunctionSuffix> name;
public:
  PrimaryTypeFunctionName(Type elemTp) : name(BaseName, primaryTypeFunctionSuffix(elemTp)) {}
  StringRef getName() const { return name; }
}

Then we could just use it à la PrimaryTypeFunctionName<"foo"> name(elemTp);. Though I'm not sure how close to this ideal we could really get. Afterall, the whole point of using SmallString was to allocate these small buffers directly on the stack; so we wouldn't want to break that.

There might be something else in the STL (or LLVM/MLIR utils) which can help simplify this metaprogramming, but I didn't see anything last time I thought about all this.

[aartbik]

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

952–956

It's doable, but I'm not sure it's worth the code it'd take to do it.

Neat! Thanks for figuring that out. I agree with you that this is probably not worth the extra machinery, given that we only have a handful of such methods.

[Peiming]

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

934

Neat, I was worrying whether there is an implicit copy here, but seems runtime library just copy the pointer ;-)

[aartbik]

rebased with main

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

static scf::ForOp createFor(OpBuilder &builder, Location loc, Value count,
   Value one = constantOne(builder, loc, indexType);
   scf::ForOp forOp = builder.create<scf::ForOp>(loc, zero, count, one, fields);
   for (unsigned i = 0, e = fields.size(); i < e; i++)
     fields[i] = forOp.getRegionIterArg(i);
   builder.setInsertionPointToStart(forOp.getBody());
   return forOp;
 }
 
+/// Translates field index to memSizes index.
+static unsigned getMemSizesIndex(unsigned field) {
+  assert(2 <= field);
+  return field - 2;
+}
+
 /// Creates a pushback op for given field and updates the fields array
 /// accordingly.
 static void createPushback(OpBuilder &builder, Location loc,
                            SmallVectorImpl<Value> &fields, unsigned field,
                            Value value) {
   assert(2 <= field && field < fields.size());
   Type etp = fields[field].getType().cast<ShapedType>().getElementType();
   if (value.getType() != etp)
     value = builder.create<arith::IndexCastOp>(loc, etp, value);
-  fields[field] =
-      builder.create<PushBackOp>(loc, fields[field].getType(), fields[1],
-                                 fields[field], value, APInt(64, field - 2));
+  fields[field] = builder.create<PushBackOp>(
+      loc, fields[field].getType(), fields[1], fields[field], value,
+      APInt(64, getMemSizesIndex(field)));
 }
 
 /// Generates insertion code.
 //
 // TODO: generalize this for any rank and format currently it is just sparse
 //       vectors as a proof of concept that we have everything in place!
 //
 static void genInsert(OpBuilder &builder, Location loc, RankedTensorType rtp,
if (encDst != encSrc) {
       // This should be handled by rewriting before codegen.
       return failure();
     }
     rewriter.replaceOp(op, adaptor.getSource());
     return success();
   }
 };
 
+/// Sparse codegen rule for number of entries operator.
+class SparseNumberOfEntriesConverter
+    : public OpConversionPattern<NumberOfEntriesOp> {
+public:
+  using OpConversionPattern::OpConversionPattern;
+  LogicalResult
+  matchAndRewrite(NumberOfEntriesOp op, OpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override {
+    // Query memSizes for the actually stored values size.
+    auto tuple = getTuple(adaptor.getTensor());
+    auto fields = tuple.getInputs();
+    unsigned lastField = fields.size() - 1;
+    Value field =
+        constantIndex(rewriter, op.getLoc(), getMemSizesIndex(lastField));
+    rewriter.replaceOpWithNewOp<memref::LoadOp>(op, fields[1], field);
+    return success();
+  }
+};
+
 } // namespace
 
 //===----------------------------------------------------------------------===//
 // Sparse tensor type conversion into an actual buffer.
 //===----------------------------------------------------------------------===//
 
 mlir::SparseTensorTypeToBufferConverter::SparseTensorTypeToBufferConverter() {
   addConversion([](Type type) { return type; });

 void mlir::populateSparseTensorCodegenPatterns(TypeConverter &typeConverter,
                                                RewritePatternSet &patterns) {
   patterns.add<SparseReturnConverter, SparseCallConverter, SparseDimOpConverter,
                SparseCastConverter, SparseTensorAllocConverter,
                SparseTensorDeallocConverter, SparseTensorLoadConverter,
                SparseExpandConverter, SparseCompressConverter,
                SparseInsertConverter, SparseToPointersConverter,
                SparseToIndicesConverter, SparseToValuesConverter,
-               SparseConvertConverter>(typeConverter, patterns.getContext());
+               SparseConvertConverter, SparseNumberOfEntriesConverter>(
+      typeConverter, patterns.getContext());
 }

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

static void newParams(OpBuilder &builder, SmallVector<Value, 8> &params,
   // User action.
   params.push_back(constantAction(builder, loc, action));
   // Payload pointer.
   if (!ptr)
     ptr = builder.create<LLVM::NullOp>(loc, getOpaquePointerType(builder));
   params.push_back(ptr);
 }
 
+/// Generates a call to obtain the values array.
+static Value genValuesCall(OpBuilder &builder, Location loc, ShapedType tp,
+                           ValueRange ptr) {
+  SmallString<15> name{"sparseValues",
+                       primaryTypeFunctionSuffix(tp.getElementType())};
+  return createFuncCall(builder, loc, name, tp, ptr, EmitCInterface::On)
+      .getResult(0);
+}
+
 /// Generates a call to release/delete a `SparseTensorCOO`.
 static void genDelCOOCall(OpBuilder &builder, Location loc, Type elemTp,
                           Value coo) {
   SmallString<21> name{"delSparseTensorCOO", primaryTypeFunctionSuffix(elemTp)};
   createFuncCall(builder, loc, name, {}, coo, EmitCInterface::Off);
 }
 
 /// Generates a call to release/delete a `SparseTensorIterator`.

 
 /// Sparse conversion rule for value accesses.
 class SparseTensorToValuesConverter : public OpConversionPattern<ToValuesOp> {
 public:
   using OpConversionPattern::OpConversionPattern;
   LogicalResult
   matchAndRewrite(ToValuesOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
-    Type resType = op.getType();
-    Type eltType = resType.cast<ShapedType>().getElementType();
-    SmallString<15> name{"sparseValues", primaryTypeFunctionSuffix(eltType)};
-    replaceOpWithFuncCall(rewriter, op, name, resType, adaptor.getOperands(),
-                          EmitCInterface::On);
+    auto resType = op.getType().cast<ShapedType>();
+    rewriter.replaceOp(op, genValuesCall(rewriter, op.getLoc(), resType,
+                                         adaptor.getOperands()));
+    return success();
+  }
+};
+
+/// Sparse conversion rule for number of entries operator.
+class SparseNumberOfEntriesConverter
+    : public OpConversionPattern<NumberOfEntriesOp> {
+public:
+  using OpConversionPattern::OpConversionPattern;
+  LogicalResult
+  matchAndRewrite(NumberOfEntriesOp op, OpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override {
+    Location loc = op.getLoc();
+    // Query values array size for the actually stored values size.
+    Type eltType = op.getTensor().getType().cast<ShapedType>().getElementType();
+    auto resTp = MemRefType::get({ShapedType::kDynamicSize}, eltType);
+    Value values = genValuesCall(rewriter, loc, resTp, adaptor.getOperands());
+    rewriter.replaceOpWithNewOp<memref::DimOp>(op, values,
+                                               constantIndex(rewriter, loc, 0));
     return success();
   }
 };
 
 /// Sparse conversion rule for tensor rematerialization.
 class SparseTensorLoadConverter : public OpConversionPattern<LoadOp> {
 public:
   using OpConversionPattern::OpConversionPattern;
   LogicalResult
   matchAndRewrite(LoadOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     if (op.getHasInserts()) {
       // Finalize any pending insertions.
       StringRef name = "endInsert";
       createFuncCall(rewriter, op->getLoc(), name, {}, adaptor.getOperands(),
                      EmitCInterface::Off);
     }
     rewriter.replaceOp(op, adaptor.getOperands());
     return success();
   }
 };
 
 /// Sparse conversion rule for the insertion operator.
 class SparseTensorInsertConverter : public OpConversionPattern<InsertOp> {
 public:
   using OpConversionPattern::OpConversionPattern;
   LogicalResult
   matchAndRewrite(InsertOp op, OpAdaptor adaptor,
void mlir::populateSparseTensorConversionPatterns(
     const SparseTensorConversionOptions &options) {
   patterns.add<SparseReturnConverter, SparseTensorToDimSizeConverter,
                SparseCastConverter, SparseTensorNewConverter,
                SparseReshapeConverter<tensor::ExpandShapeOp>,
                SparseReshapeConverter<tensor::CollapseShapeOp>,
                SparseTensorConcatConverter, SparseTensorAllocConverter,
                SparseTensorDeallocConverter, SparseTensorToPointersConverter,
                SparseTensorToIndicesConverter, SparseTensorToValuesConverter,
-               SparseTensorLoadConverter, SparseTensorInsertConverter,
-               SparseTensorExpandConverter, SparseTensorCompressConverter,
-               SparseTensorOutConverter>(typeConverter, patterns.getContext());
+               SparseNumberOfEntriesConverter, SparseTensorLoadConverter,
+               SparseTensorInsertConverter, SparseTensorExpandConverter,
+               SparseTensorCompressConverter, SparseTensorOutConverter>(
+      typeConverter, patterns.getContext());
 
   patterns.add<SparseTensorConvertConverter>(typeConverter,
                                              patterns.getContext(), options);
 }

mlir/test/Dialect/SparseTensor/codegen.mlir

 //  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_noe(
+//  CHECK-SAME: %[[A0:.*0]]: memref<1xindex>,
+//  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: %[[C2:.*]] = arith.constant 2 : index
+//       CHECK: %[[NOE:.*]] = memref.load %[[A1]][%[[C2]]] : memref<3xindex>
+//       CHECK: return %[[NOE]] : index
+func.func @sparse_noe(%arg0: tensor<128xf64, #SparseVector>) -> index {
+  %0 = sparse_tensor.number_of_entries %arg0 : tensor<128xf64, #SparseVector>
+  return %0 : index
+}
+
 // CHECK-LABEL: func @sparse_dealloc_csr(
 //  CHECK-SAME: %[[A0:.*0]]: memref<2xindex>,
 //  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<3xindex>

mlir/test/Dialect/SparseTensor/conversion.mlir

 //  CHECK-SAME: %[[A:.*]]: !llvm.ptr<i8>)
 //       CHECK: %[[T:.*]] = call @sparseValuesI8(%[[A]]) : (!llvm.ptr<i8>) -> memref<?xi8>
 //       CHECK: return %[[T]] : memref<?xi8>
 func.func @sparse_valuesi8(%arg0: tensor<128xi8, #SparseVector>) -> memref<?xi8> {
   %0 = sparse_tensor.values %arg0: tensor<128xi8, #SparseVector> to memref<?xi8>
   return %0 : memref<?xi8>
 }
 
+// CHECK-LABEL: func @sparse_noe(
+//  CHECK-SAME: %[[A:.*]]: !llvm.ptr<i8>)
+//   CHECK-DAG: %[[C:.*]] = arith.constant 0 : index
+//   CHECK-DAG: %[[T:.*]] = call @sparseValuesF64(%[[A]]) : (!llvm.ptr<i8>) -> memref<?xf64>
+//       CHECK: %[[NOE:.*]] = memref.dim %[[T]], %[[C]] : memref<?xf64>
+//       CHECK: return %[[NOE]] : index
+func.func @sparse_noe(%arg0: tensor<128xf64, #SparseVector>) -> index {
+  %0 = sparse_tensor.number_of_entries %arg0 : tensor<128xf64, #SparseVector>
+  return %0 : index
+}
+
 // CHECK-LABEL: func @sparse_reconstruct(
 //  CHECK-SAME: %[[A:.*]]: !llvm.ptr<i8>
 //       CHECK: return %[[A]] : !llvm.ptr<i8>
 func.func @sparse_reconstruct(%arg0: tensor<128xf32, #SparseVector>) -> tensor<128xf32, #SparseVector> {
   %0 = sparse_tensor.load %arg0 : tensor<128xf32, #SparseVector>
   return %0 : tensor<128xf32, #SparseVector>
 }
 

mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_dot.mlir

func.func @entry() {
     %z = arith.constant 0.0 : f32
     %x = tensor.insert %z into %t[] : tensor<f32>
     %0 = call @sparse_dot(%s1, %s2, %x) : (tensor<1024xf32, #SparseVector>,
                                            tensor<1024xf32, #SparseVector>,
                                            tensor<f32>) -> tensor<f32>
     %1 = tensor.extract %0[] : tensor<f32>
     vector.print %1 : f32
 
+    // Print number of entries in the sparse vectors.
+    //
+    // CHECK: 5
+    // CHECK: 3
+    //
+    %noe1 = sparse_tensor.number_of_entries %s1 : tensor<1024xf32, #SparseVector>
+    %noe2 = sparse_tensor.number_of_entries %s2 : tensor<1024xf32, #SparseVector>
+    vector.print %noe1 : index
+    vector.print %noe2 : index
+
     // Release the resources.
     bufferization.dealloc_tensor %s1 : tensor<1024xf32, #SparseVector>
     bufferization.dealloc_tensor %s2 : tensor<1024xf32, #SparseVector>
 
     return
   }
 }