[mlir][sparse] Add sparse_tensor.specifier and sparse_tensor.set_specifier.

[bixia]

Summary

This allows post-sparsification-rewriters to update sparse tensor storage
specifiers.

[Peiming]

mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorOps.td

267

do you mean #COO instead of #CSR here?

291

Same here.

[aartbik]

I prefer to come up with a solution (for the problem this solves) that does not expose this "field", which only makes sense along the direct IR codegen path....

mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorOps.td

let description = [{
     ```mlir
     %1 = sparse_tensor.values %0 : tensor<64x64xf64, #CSR> to memref<?xf64>
     ```
   }];
   let assemblyFormat = "$tensor attr-dict `:` type($tensor) `to` type($result)";
   let hasVerifier = 1;
 }
 
+def SparseTensor_ToSpecifierOp : SparseTensor_Op<"specifier", [Pure]>,
+    Arguments<(ins AnySparseTensor:$tensor)>,
+    Results<(outs SparseTensorStorageSpecifier:$result)> {
+  let summary = "";
+  let description = [{
+    Get the storage specifier for the tensor.
+
+    Example of getting the storage specifier for a COO tensor:
+
+    ```mlir
+    %0 = sparse_tensor.specifier %0 : tensor<64x64xf64, #CSR> to
+       !sparse_tensor.storage_specifier<#COO>
+    ```
+  }];
+  let assemblyFormat = "$tensor attr-dict `:` "
+                       "type($tensor) `to` qualified(type($result))";
+  let builders = [
+    // Builds an op by inferring the result type.
+    OpBuilder<(ins "Value":$tensor)>
+  ];
+}
+
+def SparseTensor_SetSpecifierOp : SparseTensor_Op<"set_specifier",
+[AllTypesMatch<["tensor", "result"]>]>,
+    Arguments<(ins AnySparseTensor:$tensor, SparseTensorStorageSpecifier:$specifier)>,
+    Results<(outs AnySparseTensor:$result)> {
+  let summary = "";
+  let description = [{
+    Set the storage specifier for a tensor and return the new representation of
+    the tensor.
+
+    Example of setting the storage specifier for a CSR tensor:
+
+    ```mlir
+    %0 = sparse_tensor.set_specifier %0, %1: tensor<64x64xf64, #CSR>,
+       !sparse_tensor.storage_specifier<#COO> to tensor<64x64xf64, #CSR>
+    ```
+  }];
+  let assemblyFormat = "$tensor `,` $specifier attr-dict `:` "
+                       "type($tensor) `,` qualified(type($specifier)) `to` type($result)";
+}
+
 def SparseTensor_StorageSpecifierInitOp : SparseTensor_Op<"storage_specifier.init", [Pure]>,
     Results<(outs SparseTensorStorageSpecifier:$result)> {
   let summary = "";
   let description = [{
     Returns an initial storage specifier value. A storage specifier value holds
     the sizes for tensor dimensions, pointer arrays, index arrays, and the value array.
 
     Example:

mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp

 LogicalResult ToValuesOp::verify() {
   auto ttp = getRankedTensorType(getTensor());
   auto mtp = getMemRefType(getResult());
   if (ttp.getElementType() != mtp.getElementType())
     return emitError("unexpected mismatch in element types");
   return success();
 }
 
+void ToSpecifierOp::build(OpBuilder &builder, OperationState &result,
+                          Value tensor) {
+  build(builder, result,
+        StorageSpecifierType::get(getSparseTensorEncoding(tensor.getType())),
+        tensor);
+}
+
 LogicalResult GetStorageSpecifierOp::verify() {
   RETURN_FAILURE_IF_FAILED(verifySparsifierGetterSetter(
       getSpecifierKind(), getDim(), getSpecifier(), getOperation()))
   // Checks the result type
   if (getSpecifier().getType().getFieldType(getSpecifierKind(), getDim()) !=
       getResult().getType()) {
     return emitError(
         "type mismatch between requested specifier field and result value");

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

matchAndRewrite(ToValuesOp op, OpAdaptor adaptor,
     // The cast_op is inserted by type converter to intermix 1:N type
     // conversion.
     auto desc = getDescriptorFromTensorTuple(adaptor.getTensor());
     rewriter.replaceOp(op, desc.getValMemRef());
     return success();
   }
 };
 
+/// Sparse codegen rule for storage specifier accesses.
+class SparseToSpecifierConverter : public OpConversionPattern<ToSpecifierOp> {
+public:
+  using OpAdaptor = typename ToSpecifierOp::Adaptor;
+  using OpConversionPattern<ToSpecifierOp>::OpConversionPattern;
+  LogicalResult
+  matchAndRewrite(ToSpecifierOp op, OpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override {
+    auto desc = getDescriptorFromTensorTuple(adaptor.getTensor());
+    rewriter.replaceOp(op, desc.getSpecifier());
+    return success();
+  }
+};
+
+/// Sparse codegen rule for setting storage specifier.
+class SparseSetSpecifierConverter : public OpConversionPattern<SetSpecifierOp> {
+public:
+  using OpAdaptor = typename SetSpecifierOp::Adaptor;
+  using OpConversionPattern<SetSpecifierOp>::OpConversionPattern;
+  LogicalResult
+  matchAndRewrite(SetSpecifierOp op, OpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override {
+    SmallVector<Value> fields;
+    auto desc = getMutDescriptorFromTensorTuple(adaptor.getTensor(), fields);
+    desc.setSpecifier(op.getSpecifier());
+    Location loc = op.getLoc();
+    rewriter.replaceOp(op, genTuple(rewriter, loc, op.getType(), fields));
+    return success();
+  }
+};
+
 /// Sparse codegen rule for the convert operator.
 class SparseConvertConverter : public OpConversionPattern<ConvertOp> {
 public:
   using OpConversionPattern::OpConversionPattern;
   LogicalResult
   matchAndRewrite(ConvertOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     SparseTensorEncodingAttr encDst = getSparseTensorEncoding(op.getType());
void mlir::populateSparseTensorCodegenPatterns(
     bool enableBufferInitialization) {
   patterns.add<SparsePackOpConverter, SparseReturnConverter,
                SparseCallConverter, SparseDimOpConverter, SparseCastConverter,
                SparseTensorDeallocConverter, SparseExtractSliceCoverter,
                SparseTensorLoadConverter, SparseExpandConverter,
                SparseCompressConverter, SparseInsertConverter,
                SparseToPointersConverter, SparseToIndicesConverter,
                SparseToIndicesBufferConverter, SparseToValuesConverter,
+               SparseToSpecifierConverter, SparseSetSpecifierConverter,
                SparseConvertConverter, SparseNumberOfEntriesConverter>(
       typeConverter, patterns.getContext());
   patterns.add<SparseTensorAllocConverter>(typeConverter, patterns.getContext(),
                                            enableBufferInitialization);
 }

mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorStorageLayout.h

public:
   }
 
   unsigned getNumFields() const { return fields.size(); }
 
   ///
   /// Getters: get the value for required field.
   ///
 
+  Value getSpecifier() const { return fields.back(); }
+
   Value getSpecifierField(OpBuilder &builder, Location loc,
                           StorageSpecifierKind kind,
                           std::optional<unsigned> dim) const {
     SparseTensorSpecifier md(fields.back());
     return md.getSpecifierField(builder, loc, kind, dim);
   }
 
   Value getDimSize(OpBuilder &builder, Location loc, unsigned dim) const {
void setMemRefField(unsigned fidx, Value v) {
     fields[fidx] = v;
   }
 
   void setField(unsigned fidx, Value v) {
     assert(fidx < fields.size());
     fields[fidx] = v;
   }
 
+  void setSpecifier(Value spec) { fields[fields.size() - 1] = spec; }
+
   void setSpecifierField(OpBuilder &builder, Location loc,
                          StorageSpecifierKind kind, std::optional<unsigned> dim,
                          Value v) {
     SparseTensorSpecifier md(fields.back());
     md.setSpecifierField(builder, loc, v, kind, dim);
     fields.back() = md;
   }
 

mlir/test/Dialect/SparseTensor/codegen.mlir

 //  CHECK-SAME: %[[A4:.*4]]: memref<?xf64>,
 //  CHECK-SAME: %[[A5:.*5]]: !sparse_tensor.storage_specifier
 //       CHECK: return %[[A4]] : memref<?xf64>
 func.func @sparse_values_coo(%arg0: tensor<?x?x?xf64, #ccoo>) -> memref<?xf64> {
   %0 = sparse_tensor.values %arg0 : tensor<?x?x?xf64, #ccoo> to memref<?xf64>
   return %0 : memref<?xf64>
 }
 
+// CHECK-LABEL: func.func @sparse_specifier(
+//  CHECK-SAME: %[[A0:.*0]]: memref<?xindex>,
+//  CHECK-SAME: %[[A1:.*1]]: memref<?xindex>,
+//  CHECK-SAME: %[[A2:.*2]]: memref<?xindex>,
+//  CHECK-SAME: %[[A3:.*3]]: memref<?xindex>,
+//  CHECK-SAME: %[[A4:.*4]]: memref<?xf64>,
+//  CHECK-SAME: %[[A5:.*5]]: !sparse_tensor.storage_specifier
+//       CHECK: return %[[A5]] :  !sparse_tensor.storage_specifier
+func.func @sparse_specifier(%arg0: tensor<?x?x?xf64, #ccoo>) ->  !sparse_tensor.storage_specifier<#ccoo> {
+  %0 = sparse_tensor.specifier %arg0 : tensor<?x?x?xf64, #ccoo> to  !sparse_tensor.storage_specifier<#ccoo>
+  return %0 :  !sparse_tensor.storage_specifier<#ccoo>
+}
+
+// CHECK-LABEL: func.func @sparse_set_specifier(
+//  CHECK-SAME: %[[A0:.*0]]: memref<?xindex>,
+//  CHECK-SAME: %[[A1:.*1]]: memref<?xindex>,
+//  CHECK-SAME: %[[A2:.*2]]: memref<?xindex>,
+//  CHECK-SAME: %[[A3:.*3]]: memref<?xindex>,
+//  CHECK-SAME: %[[A4:.*4]]: memref<?xf64>,
+//  CHECK-SAME: %[[A5:.*5]]: !sparse_tensor.storage_specifier
+//  CHECK-SAME: %[[A6:.*6]]: !sparse_tensor.storage_specifier
+//       CHECK: return %[[A0]], %[[A1]], %[[A2]], %[[A3]], %[[A4]], %[[A6]]
+func.func @sparse_set_specifier(%arg0: tensor<?x?x?xf64, #ccoo>, %arg1: !sparse_tensor.storage_specifier<#ccoo>) -> tensor<?x?x?xf64, #ccoo> {
+  %0 = sparse_tensor.set_specifier %arg0, %arg1 : tensor<?x?x?xf64, #ccoo>, !sparse_tensor.storage_specifier<#ccoo> to tensor<?x?x?xf64, #ccoo>
+  return %0 : tensor<?x?x?xf64, #ccoo>
+}
 
 // CHECK-LABEL:   func.func @sparse_indices_coo(
 //  CHECK-SAME: %[[A0:.*0]]: memref<?xindex>,
 //  CHECK-SAME: %[[A1:.*1]]: memref<?xindex>,
 //  CHECK-SAME: %[[A2:.*2]]: memref<?xindex>,
 //  CHECK-SAME: %[[A3:.*3]]: memref<?xindex>,
 //  CHECK-SAME: %[[A4:.*4]]: memref<?xf64>,
 //  CHECK-SAME: %[[A5:.*5]]: !sparse_tensor.storage_specifier

mlir/test/Dialect/SparseTensor/roundtrip.mlir

func.func @sparse_values(%arg0: tensor<128xf64, #SparseVector>) -> memref<?xf64> {
   %0 = sparse_tensor.values %arg0 : tensor<128xf64, #SparseVector> to memref<?xf64>
   return %0 : memref<?xf64>
 }
 
 // -----
 
 #SparseVector = #sparse_tensor.encoding<{dimLevelType = ["compressed"]}>
 
+// CHECK-LABEL: func @sparse_get_specifier(
+//  CHECK-SAME: %[[A:.*]]: tensor<128xf64, #{{.*}}>)
+//       CHECK: %[[T:.*]] = sparse_tensor.specifier %[[A]] : tensor<128xf64, #{{.*}}> to !sparse_tensor.storage_specifier<#{{.*}}>
+//       CHECK: return %[[T]] : !sparse_tensor.storage_specifier<#{{.*}}>
+func.func @sparse_get_specifier(%arg0: tensor<128xf64, #SparseVector>)
+  -> !sparse_tensor.storage_specifier<#SparseVector> {
+  %0 = sparse_tensor.specifier %arg0 : tensor<128xf64, #SparseVector> to
+    !sparse_tensor.storage_specifier<#SparseVector>
+  return %0 : !sparse_tensor.storage_specifier<#SparseVector>
+}
+
+// -----
+
+#SparseVector = #sparse_tensor.encoding<{dimLevelType = ["compressed"]}>
+
+// CHECK-LABEL: func @sparse_set_specifier(
+//  CHECK-SAME: %[[A:.*]]: tensor<128xf64, #{{.*}}>,
+//  CHECK-SAME: %[[B:.*]]: !sparse_tensor.storage_specifier<#{{.*}}>)
+//       CHECK: %[[T:.*]] = sparse_tensor.set_specifier %[[A]], %[[B]] : tensor<128xf64, #{{.*}}>, !sparse_tensor.storage_specifier<#{{.*}}> to tensor<128xf64, #{{.*}}>
+//       CHECK: return %[[T]] : tensor<128xf64, #{{.*}}>
+func.func @sparse_set_specifier(%arg0: tensor<128xf64, #SparseVector>,
+  %arg1:!sparse_tensor.storage_specifier<#SparseVector>)
+  -> tensor<128xf64, #SparseVector> {
+  %0 = sparse_tensor.set_specifier %arg0, %arg1 : tensor<128xf64, #SparseVector>,
+    !sparse_tensor.storage_specifier<#SparseVector> to tensor<128xf64, #SparseVector>
+  return %0 : tensor<128xf64, #SparseVector>
+}
+
+
+// -----
+
+#SparseVector = #sparse_tensor.encoding<{dimLevelType = ["compressed"]}>
+
 // CHECK-LABEL: func @sparse_metadata_init(
 //       CHECK: %[[T:.*]] = sparse_tensor.storage_specifier.init : !sparse_tensor.storage_specifier<#{{.*}}>
 //       CHECK: return %[[T]] : !sparse_tensor.storage_specifier<#{{.*}}>
 func.func @sparse_metadata_init() -> !sparse_tensor.storage_specifier<#SparseVector> {
   %0 = sparse_tensor.storage_specifier.init : !sparse_tensor.storage_specifier<#SparseVector>
   return %0 : !sparse_tensor.storage_specifier<#SparseVector>
 }