[mlir][sparse] Use DLT in the mangled function names for insertion.

[aartbik]

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

501

how about ordered/unordered? the code may not differ right now, but perhaps in the future? seems good to get a unique method per unique situation?

[wrengr]

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

499–501

Now would be a good time to factor this part out into a standalone function of type DLT -> StringRef for returning the mangling itself (e.g., so we can use the same mangling elsewhere, as needed). To minimize string copying, you should write it like how toMLIRString is done (with a switch and returning the whole string at once, rather than gluing it together from parts).

501

+1, if only for the sake of getting it out of the way at the same time.

[bixia]

Modify toMLIRString to add toString and use it here.

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

499–501

Add toString and use it here.

[aartbik]

mlir/include/mlir/Dialect/SparseTensor/IR/Enums.h

197

This now has runtime overhead. Since we only use it for debugging at one place, let's fix the only other place

./IR/SparseTensorDialect.cpp:

change
printer << toMLIRString(getDimLevelType()[i]);
into
printer << "\"" << toMLIRString(getDimLevelType()[i]) << "\""

and remove this one (and call your onetoMLIRString again

[bixia]

Remove toString and modified toMLIRString.

[aartbik]

mlir/include/mlir/Dialect/SparseTensor/IR/Enums.h

194

be careful with removing this! even though all cases are covered, some tools want an "llvm_unreachable" here, and Wren was going to write one that does not need to pull in dependences into this pure include

but when you remove it, perhaps some builds will break....

[bixia]

Add the return statement back.

[bixia]

Rebase.

mlir/include/mlir/Dialect/SparseTensor/IR/Enums.h

enum class LevelFormat : uint8_t {
   Singleton = 16, // 0b100_00
 };
 
 /// Returns string representation of the given dimension level type.
 inline std::string toMLIRString(DimLevelType dlt) {
   switch (dlt) {
   // TODO: should probably raise an error instead of printing it...
   case DimLevelType::Undef:
-    return "\"undef\"";
+    return "undef";
   case DimLevelType::Dense:
-    return "\"dense\"";
+    return "dense";
   case DimLevelType::Compressed:
-    return "\"compressed\"";
+    return "compressed";
   case DimLevelType::CompressedNu:
-    return "\"compressed-nu\"";
+    return "compressed-nu";
   case DimLevelType::CompressedNo:
-    return "\"compressed-no\"";
+    return "compressed-no";
   case DimLevelType::CompressedNuNo:
-    return "\"compressed-nu-no\"";
+    return "compressed-nu-no";
   case DimLevelType::Singleton:
-    return "\"singleton\"";
+    return "singleton";
   case DimLevelType::SingletonNu:
-    return "\"singleton-nu\"";
+    return "singleton-nu";
   case DimLevelType::SingletonNo:
-    return "\"singleton-no\"";
+    return "singleton-no";
   case DimLevelType::SingletonNuNo:
-    return "\"singleton-nu-no\"";
+    return "singleton-nu-no";
   }
-  return "";
 }
 
 /// Check that the `DimLevelType` contains a valid (possibly undefined) value.
 constexpr bool isValidDLT(DimLevelType dlt) {
   const uint8_t formatBits = static_cast<uint8_t>(dlt) >> 2;
   const uint8_t propertyBits = static_cast<uint8_t>(dlt) & 3;
   // If undefined or dense, then must be unique and ordered.
   // Otherwise, the format must be one of the known ones.
   return (formatBits <= 1) ? (propertyBits == 0)
                            : (formatBits == 2 || formatBits == 4);
 }
 

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

Attribute SparseTensorEncodingAttr::parse(AsmParser &parser, Type type) {
   return parser.getChecked<SparseTensorEncodingAttr>(
       parser.getContext(), dlt, dimOrd, higherOrd, ptr, ind);
 }
 
 void SparseTensorEncodingAttr::print(AsmPrinter &printer) const {
   // Print the struct-like storage in dictionary fashion.
   printer << "<{ dimLevelType = [ ";
   for (unsigned i = 0, e = getDimLevelType().size(); i < e; i++) {
-    printer << toMLIRString(getDimLevelType()[i]);
+    printer << "\"" << toMLIRString(getDimLevelType()[i]) << "\"";
     if (i != e - 1)
       printer << ", ";
   }
   printer << " ]";
   // Print remaining members only for non-default values.
   if (!hasIdDimOrdering())
     printer << ", dimOrdering = affine_map<" << getDimOrdering() << ">";
   if (getHigherOrdering())

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

 //===----------------------------------------------------------------------===//
 
 #include "CodegenUtils.h"
 
 #include "mlir/Dialect/Bufferization/IR/Bufferization.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/Linalg/Utils/Utils.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
+#include "mlir/Dialect/SparseTensor/IR/Enums.h"
 #include "mlir/Dialect/SparseTensor/IR/SparseTensor.h"
 #include "mlir/Dialect/SparseTensor/Transforms/Passes.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/Transforms/DialectConversion.h"
 
 using namespace mlir;
 using namespace mlir::sparse_tensor;
 

 /// function doesn't exist yet, call `createFunc` to generate the function.
 static void genInsertionCallHelper(OpBuilder &builder,
                                    MutSparseTensorDescriptor desc,
                                    SmallVectorImpl<Value> &indices, Value value,
                                    func::FuncOp insertPoint,
                                    StringRef namePrefix,
                                    FuncGeneratorType createFunc) {
   // The mangled name of the function has this format:
-  //   <namePrefix>_[C|S|D]_<shape>_<ordering>_<eltType>
+  //   <namePrefix>_<DLT>_<shape>_<ordering>_<eltType>
   //     _<indexBitWidth>_<pointerBitWidth>
   RankedTensorType rtp = desc.getTensorType();
   SmallString<32> nameBuffer;
   llvm::raw_svector_ostream nameOstream(nameBuffer);
   nameOstream << namePrefix;
   unsigned rank = rtp.getShape().size();
   assert(rank == indices.size());
   for (unsigned d = 0; d < rank; d++) {
-    if (isCompressedDim(rtp, d)) {
-      nameOstream << "C_";
-    } else if (isSingletonDim(rtp, d)) {
-      nameOstream << "S_";
-    } else {
-      nameOstream << "D_";
-    }
+    nameOstream << toMLIRString(getDimLevelType(rtp, d)) << "_";
   }
   // Static dim sizes are used in the generated code while dynamic sizes are
   // loaded from the dimSizes buffer. This is the reason for adding the shape
   // to the function name.
   for (auto d : rtp.getShape())
     nameOstream << d << "_";
   SparseTensorEncodingAttr enc = getSparseTensorEncoding(rtp);
   // Permutation information is also used in generating insertion.
   if (enc.getDimOrdering() && !enc.getDimOrdering().isIdentity())

mlir/test/Dialect/SparseTensor/codegen.mlir

-// RUN: mlir-opt %s --sparse-tensor-codegen  --canonicalize --cse | FileCheck %s
+// RUN: mlir-opt %s --sparse-tensor-codegen  --canonicalize -cse | FileCheck %s
 
 #SV = #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>
 
 #SparseVector = #sparse_tensor.encoding<{
   dimLevelType = [ "compressed" ],
   indexBitWidth = 64,
   pointerBitWidth = 32
 }>
#DCSR = #sparse_tensor.encoding<{
   pointerBitWidth = 32
 }>
 
 #Dense3D = #sparse_tensor.encoding<{
   dimLevelType = [ "dense", "dense", "dense" ],
   dimOrdering = affine_map<(i, j, k) -> (k, i, j)>
 }>
 
+#Coo = #sparse_tensor.encoding<{
+  dimLevelType = [ "compressed-nu", "singleton" ]
+}>
+
 // CHECK-LABEL: func @sparse_nop(
 //  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: return %[[A0]], %[[A1]], %[[A2]], %[[A3]], %[[A4]] :
 //  CHECK-SAME:   memref<1xindex>, memref<3xindex>, memref<?xi32>, memref<?xi64>, memref<?xf64>

 //       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>
   return %added : memref<?xindex>
 }
 
-// CHECK-LABEL: func.func private @_insert_C_100_f64_0_0(
+// CHECK-LABEL: func.func private @_insert_compressed_100_f64_0_0(
 //  CHECK-SAME: %[[A0:.*0]]: memref<1xindex>,
 //  CHECK-SAME: %[[A1:.*1]]: memref<3xindex>,
 //  CHECK-SAME: %[[A2:.*2]]: memref<?xindex>,
 //  CHECK-SAME: %[[A3:.*3]]: memref<?xindex>,
 //  CHECK-SAME: %[[A4:.*4]]: memref<?xf64>,
 //  CHECK-SAME: %[[A5:.*5]]: index,
 //  CHECK-SAME: %[[A6:.*6]]: f64)
 //       CHECK:   %[[PV:.*]] = sparse_tensor.push_back %[[A1]], %[[A4]], %[[A6]] {idx = 2 : index} : memref<3xindex>, memref<?xf64>, f64

 //   CHECK-DAG: %[[F0:.*]] = arith.constant 0.000000e+00 : f64
 //   CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
 //   CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index
 //       CHECK: sparse_tensor.sort %[[A8]], %[[A7]] : memref<?xindex>
 //       CHECK: %[[R:.*]]:5 = scf.for %[[I:.*]] = %[[C0]] to %[[A8]] step %[[C1]]
 //  CHECK-SAME: iter_args(%[[P0:.*]] = %[[A0]], %[[P1:.*]] = %[[A1]], %[[P2:.*]] = %[[A2]], %[[P3:.*]] = %[[A3]], %[[P4:.*]] = %[[A4]]) -> (memref<1xindex>, memref<3xindex>, memref<?xindex>, memref<?xindex>, memref<?xf64>) {
 //       CHECK:   %[[INDEX:.*]] = memref.load %[[A7]][%[[I]]] : memref<?xindex>
 //       CHECK:   %[[VAL:.*]] = memref.load %[[A5]][%[[INDEX]]] : memref<?xf64>
-//       CHECK:   %[[C:.*]]:5 = func.call @_insert_C_100_f64_0_0(%[[P0]], %[[P1]], %[[P2]], %[[P3]], %[[P4]], %[[INDEX]], %[[VAL]])
+//       CHECK:   %[[C:.*]]:5 = func.call @_insert_compressed_100_f64_0_0(%[[P0]], %[[P1]], %[[P2]], %[[P3]], %[[P4]], %[[INDEX]], %[[VAL]])
 //       CHECK:   memref.store %[[F0]], %[[A5]][%[[INDEX]]] : memref<?xf64>
 //       CHECK:   memref.store %[[B0]], %[[A6]][%[[INDEX]]] : memref<?xi1>
 //       CHECK:   scf.yield %[[C]]#0, %[[C]]#1, %[[C]]#2, %[[C]]#3, %[[C]]#4 : memref<1xindex>, memref<3xindex>, memref<?xindex>, memref<?xindex>, memref<?xf64>
 //       CHECK: }
 //       CHECK: memref.dealloc %[[A5]] : memref<?xf64>
 //       CHECK: memref.dealloc %[[A6]] : memref<?xi1>
 //       CHECK: memref.dealloc %[[A7]] : memref<?xindex>
 //       CHECK: return %[[R]]#0, %[[R]]#1, %[[R]]#2, %[[R]]#3, %[[R]]#4
 //  CHECK-SAME:   memref<1xindex>, memref<3xindex>, memref<?xindex>, memref<?xindex>, memref<?xf64>
 func.func @sparse_compression_1d(%tensor: tensor<100xf64, #SV>,
                                  %values: memref<?xf64>,
                                  %filled: memref<?xi1>,
                                  %added: memref<?xindex>,
                                  %count: index) -> tensor<100xf64, #SV> {
   %0 = sparse_tensor.compress %values, %filled, %added, %count into %tensor[]
     : memref<?xf64>, memref<?xi1>, memref<?xindex>, tensor<100xf64, #SV>
   %1 = sparse_tensor.load %0 hasInserts : tensor<100xf64, #SV>
   return %1 : tensor<100xf64, #SV>
 }
 
-// CHECK-LABEL: func.func private @_insert_D_C_8_8_f64_64_32(
+// CHECK-LABEL: func.func private @_insert_dense_compressed_8_8_f64_64_32(
 //  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-SAME: %[[A5:.*5]]: index,
 //  CHECK-SAME: %[[A6:.*6]]: index,
 //  CHECK-SAME: %[[A7:.*7]]: f64)

 //   CHECK-DAG: %[[F0:.*]] = arith.constant 0.000000e+00 : f64
 //   CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
 //   CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index
 //       CHECK: sparse_tensor.sort %[[A8]], %[[A7]] : memref<?xindex>
 //       CHECK: %[[R:.*]]:5 = scf.for %[[I:.*]] = %[[C0]] to %[[A8]] step %[[C1]]
 //  CHECK-SAME: iter_args(%[[P0:.*]] = %[[A0]], %[[P1:.*]] = %[[A1]], %[[P2:.*]] = %[[A2]], %[[P3:.*]] = %[[A3]], %[[P4:.*]] = %[[A4]]) -> (memref<2xindex>, memref<3xindex>, memref<?xi32>, memref<?xi64>, memref<?xf64>) {
 //       CHECK:   %[[INDEX:.*]] = memref.load %[[A7]][%[[I]]] : memref<?xindex>
 //       CHECK:   %[[VAL:.*]] = memref.load %[[A5]][%[[INDEX]]] : memref<?xf64>
-//       CHECK:   %[[C:.*]]:5 = func.call @_insert_D_C_8_8_f64_64_32(%[[P0]], %[[P1]], %[[P2]], %[[P3]], %[[P4]], %[[A9]], %[[INDEX]], %[[VAL]])
+//       CHECK:   %[[C:.*]]:5 = func.call @_insert_dense_compressed_8_8_f64_64_32(%[[P0]], %[[P1]], %[[P2]], %[[P3]], %[[P4]], %[[A9]], %[[INDEX]], %[[VAL]])
 //       CHECK:   memref.store %[[F0]], %[[A5]][%[[INDEX]]] : memref<?xf64>
 //       CHECK:   memref.store %[[B0]], %[[A6]][%[[INDEX]]] : memref<?xi1>
 //       CHECK:   scf.yield %[[C]]#0, %[[C]]#1, %[[C]]#2, %[[C]]#3, %[[C]]#4 : memref<2xindex>, memref<3xindex>, memref<?xi32>, memref<?xi64>, memref<?xf64>
 //       CHECK: }
 //       CHECK: memref.dealloc %[[A5]] : memref<?xf64>
 //       CHECK: memref.dealloc %[[A6]] : memref<?xi1>
 //       CHECK: memref.dealloc %[[A7]] : memref<?xindex>
 //       CHECK: return %[[R]]#0, %[[R]]#1, %[[R]]#2, %[[R]]#3, %[[R]]#4
 //  CHECK-SAME:   memref<2xindex>, memref<3xindex>, memref<?xi32>, memref<?xi64>, memref<?xf64>
 func.func @sparse_compression(%tensor: tensor<8x8xf64, #CSR>,
                               %values: memref<?xf64>,
                               %filled: memref<?xi1>,
                               %added: memref<?xindex>,
                               %count: index,
                               %i: index) -> tensor<8x8xf64, #CSR> {
   %0 = sparse_tensor.compress %values, %filled, %added, %count into %tensor[%i]
     : memref<?xf64>, memref<?xi1>, memref<?xindex>, tensor<8x8xf64, #CSR>
   %1 = sparse_tensor.load %0 hasInserts : tensor<8x8xf64, #CSR>
   return %1 : tensor<8x8xf64, #CSR>
 }
 
-// CHECK-LABEL: func.func private @_insert_D_C_8_8_f64_0_0(
+// CHECK-LABEL: func.func private @"_insert_dense_compressed-no_8_8_f64_0_0"(
 //  CHECK-SAME: %[[A0:.*0]]: memref<2xindex>,
 //  CHECK-SAME: %[[A1:.*1]]: memref<3xindex>,
 //  CHECK-SAME: %[[A2:.*2]]: memref<?xindex>,
 //  CHECK-SAME: %[[A3:.*3]]: memref<?xindex>,
 //  CHECK-SAME: %[[A4:.*4]]: memref<?xf64>,
 //  CHECK-SAME: %[[A5:.*5]]: index,
 //  CHECK-SAME: %[[A6:.*6]]: index,
 //  CHECK-SAME: %[[A7:.*7]]: f64)

 //   CHECK-DAG: %[[F0:.*]] = arith.constant 0.000000e+00 : f64
 //   CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
 //   CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index
 //   CHECK-NOT: sparse_tensor.sort
 //       CHECK: %[[R:.*]]:5 = scf.for %[[I:.*]] = %[[C0]] to %[[A8]] step %[[C1]]
 //  CHECK-SAME: iter_args(%[[P0:.*]] = %[[A0]], %[[P1:.*]] = %[[A1]], %[[P2:.*]] = %[[A2]], %[[P3:.*]] = %[[A3]], %[[P4:.*]] = %[[A4]]) -> (memref<2xindex>, memref<3xindex>, memref<?xindex>, memref<?xindex>, memref<?xf64>) {
 //       CHECK:   %[[INDEX:.*]] = memref.load %[[A7]][%[[I]]] : memref<?xindex>
 //       CHECK:   %[[VAL:.*]] = memref.load %[[A5]][%[[INDEX]]] : memref<?xf64>
-//       CHECK:   %[[C:.*]]:5 = func.call @_insert_D_C_8_8_f64_0_0(%[[P0]], %[[P1]], %[[P2]], %[[P3]], %[[P4]], %[[A9]], %[[INDEX]], %[[VAL]])
+//       CHECK:   %[[C:.*]]:5 = func.call @"_insert_dense_compressed-no_8_8_f64_0_0"(%[[P0]], %[[P1]], %[[P2]], %[[P3]], %[[P4]], %[[A9]], %[[INDEX]], %[[VAL]])
 //       CHECK:   memref.store %[[F0]], %[[A5]][%[[INDEX]]] : memref<?xf64>
 //       CHECK:   memref.store %[[B0]], %[[A6]][%[[INDEX]]] : memref<?xi1>
 //       CHECK:   scf.yield %[[C]]#0, %[[C]]#1, %[[C]]#2, %[[C]]#3, %[[C]]#4 : memref<2xindex>, memref<3xindex>, memref<?xindex>, memref<?xindex>, memref<?xf64>
 //       CHECK: }
 //       CHECK: memref.dealloc %[[A5]] : memref<?xf64>
 //       CHECK: memref.dealloc %[[A6]] : memref<?xi1>
 //       CHECK: memref.dealloc %[[A7]] : memref<?xindex>
 //       CHECK: return %[[R]]#0, %[[R]]#1, %[[R]]#2, %[[R]]#3, %[[R]]#4
 //  CHECK-SAME:   memref<2xindex>, memref<3xindex>, memref<?xindex>, memref<?xindex>, memref<?xf64>
 func.func @sparse_compression_unordered(%tensor: tensor<8x8xf64, #UCSR>,
                                         %values: memref<?xf64>,
                                         %filled: memref<?xi1>,
                                         %added: memref<?xindex>,
                                         %count: index,
                                         %i: index) -> tensor<8x8xf64, #UCSR> {
   %0 = sparse_tensor.compress %values, %filled, %added, %count into %tensor[%i]
     : memref<?xf64>, memref<?xi1>, memref<?xindex>, tensor<8x8xf64, #UCSR>
   %1 = sparse_tensor.load %0 hasInserts : tensor<8x8xf64, #UCSR>
   return %1 : tensor<8x8xf64, #UCSR>
 }
 
-// CHECK-LABEL: func.func private @_insert_C_128_f64_0_0(
+// CHECK-LABEL: func.func private @_insert_compressed_128_f64_0_0(
 //  CHECK-SAME: %[[A0:.*0]]: memref<1xindex>,
 //  CHECK-SAME: %[[A1:.*1]]: memref<3xindex>,
 //  CHECK-SAME: %[[A2:.*2]]: memref<?xindex>,
 //  CHECK-SAME: %[[A3:.*3]]: memref<?xindex>,
 //  CHECK-SAME: %[[A4:.*4]]: memref<?xf64>,
 //  CHECK-SAME: %[[A5:.*5]]: index,
 //  CHECK-SAME: %[[A6:.*6]]: f64)
 //       CHECK: %[[P:.*]] = sparse_tensor.push_back %[[A1]], %[[A4]], %[[A6]]  {idx = 2 : index} : memref<3xindex>, memref<?xf64>, f64
 //       CHECK: return %[[A0]], %[[A1]], %[[A2]], %{{.*}}, %[[P]] :
 //       CHECK: func @sparse_insert(
 //  CHECK-SAME: %[[A0:.*0]]: memref<1xindex>,
 //  CHECK-SAME: %[[A1:.*1]]: memref<3xindex>,
 //  CHECK-SAME: %[[A2:.*2]]: memref<?xindex>,
 //  CHECK-SAME: %[[A3:.*3]]: memref<?xindex>,
 //  CHECK-SAME: %[[A4:.*4]]: memref<?xf64>,
 //  CHECK-SAME: %[[A5:.*5]]: index,
 //  CHECK-SAME: %[[A6:.*6]]: f64)
-//       CHECK: %[[R:.*]]:5 = call @_insert_C_128_f64_0_0(%[[A0]], %[[A1]], %[[A2]], %[[A3]], %[[A4]], %[[A5]], %[[A6]])
+//       CHECK: %[[R:.*]]:5 = call @_insert_compressed_128_f64_0_0(%[[A0]], %[[A1]], %[[A2]], %[[A3]], %[[A4]], %[[A5]], %[[A6]])
 //       CHECK: return %[[R]]#0, %[[R]]#1, %[[R]]#2, %[[R]]#3, %[[R]]#4
 //  CHECK-SAME:   memref<1xindex>, memref<3xindex>, memref<?xindex>, memref<?xindex>, memref<?xf64>
 func.func @sparse_insert(%arg0: tensor<128xf64, #SV>, %arg1: index, %arg2: f64) -> tensor<128xf64, #SV> {
   %0 = sparse_tensor.insert %arg2 into %arg0[%arg1] : tensor<128xf64, #SV>
   %1 = sparse_tensor.load %0 hasInserts : tensor<128xf64, #SV>
   return %1 : tensor<128xf64, #SV>
 }
 
-// CHECK-LABEL: func.func private @_insert_C_128_f64_64_32(
+// CHECK-LABEL: func.func private @_insert_compressed_128_f64_64_32(
 //  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-SAME: %[[A5:.*5]]: index,
 //  CHECK-SAME: %[[A6:.*6]]: f64)
 //       CHECK: %[[P:.*]] = sparse_tensor.push_back %[[A1]], %[[A4]], %[[A6]]  {idx = 2 : index} : memref<3xindex>, memref<?xf64>, f64
 //       CHECK: return %[[A0]], %[[A1]], %[[A2]], %{{.*}}, %[[P]] :
 //       CHECK: func @sparse_insert_typed(
 //  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-SAME: %[[A5:.*5]]: index,
 //  CHECK-SAME: %[[A6:.*6]]: f64)
-//       CHECK: %[[R:.*]]:5 = call @_insert_C_128_f64_64_32(%[[A0]], %[[A1]], %[[A2]], %[[A3]], %[[A4]], %[[A5]], %[[A6]])
+//       CHECK: %[[R:.*]]:5 = call @_insert_compressed_128_f64_64_32(%[[A0]], %[[A1]], %[[A2]], %[[A3]], %[[A4]], %[[A5]], %[[A6]])
 //       CHECK: return %[[R]]#0, %[[R]]#1, %[[R]]#2, %[[R]]#3, %[[R]]#4
 //  CHECK-SAME:   memref<1xindex>, memref<3xindex>, memref<?xi32>, memref<?xi64>, memref<?xf64>
 func.func @sparse_insert_typed(%arg0: tensor<128xf64, #SparseVector>, %arg1: index, %arg2: f64) -> tensor<128xf64, #SparseVector> {
   %0 = sparse_tensor.insert %arg2 into %arg0[%arg1] : tensor<128xf64, #SparseVector>
   %1 = sparse_tensor.load %0 hasInserts : tensor<128xf64, #SparseVector>
   return %1 : tensor<128xf64, #SparseVector>
 }
 
+// CHECK-LABEL: func.func private @"_insert_compressed-nu_singleton_5_6_f64_0_0"(
+//  CHECK-SAME: %[[A0:.*0]]: memref<2xindex>,
+//  CHECK-SAME: %[[A1:.*1]]: memref<4xindex>,
+//  CHECK-SAME: %[[A2:.*2]]: memref<?xindex>,
+//  CHECK-SAME: %[[A3:.*3]]: memref<?xindex>,
+//  CHECK-SAME: %[[A4:.*4]]: memref<?xindex>,
+//  CHECK-SAME: %[[A5:.*5]]: memref<?xf64>,
+//  CHECK-SAME: %[[A6:.*6]]: index,
+//  CHECK-SAME: %[[A7:.*7]]: index,
+//  CHECK-SAME: %[[A8:.*8]]: f64)
+//       CHECK: %[[P0:.*]] = sparse_tensor.push_back %[[A1]], %[[A3]], %[[A6]]
+//       CHECK: %[[P1:.*]] = sparse_tensor.push_back %[[A1]], %[[A4]], %[[A7]]
+//       CHECK: %[[P2:.*]] = sparse_tensor.push_back %[[A1]], %[[A5]], %[[A8]]
+//       CHECK: return %[[A0]], %[[A1]], %[[A2]], %[[P0]], %[[P1]], %[[P2]]
+//       CHECK: func.func @sparse_insert_coo(
+//  CHECK-SAME: %[[A0:.*0]]: memref<2xindex>,
+//  CHECK-SAME: %[[A1:.*1]]: memref<4xindex>,
+//  CHECK-SAME: %[[A2:.*2]]: memref<?xindex>,
+//  CHECK-SAME: %[[A3:.*3]]: memref<?xindex>,
+//  CHECK-SAME: %[[A4:.*4]]: memref<?xindex>,
+//  CHECK-SAME: %[[A5:.*5]]: memref<?xf64>,
+//  CHECK-SAME: %[[A6:.*6]]: index,
+//  CHECK-SAME: %[[A7:.*7]]: f64)
+//       CHECK: %[[R:.*]]:6 = call @"_insert_compressed-nu_singleton_5_6_f64_0_0"(%[[A0]], %[[A1]], %[[A2]], %[[A3]], %[[A4]], %[[A5]], %[[A6]], %[[A6]], %[[A7]])
+//       CHECK: return %[[R]]#0, %[[R]]#1, %[[R]]#2, %[[R]]#3, %[[R]]#4, %[[R]]#5
+func.func @sparse_insert_coo(%arg0: tensor<5x6xf64, #Coo>, %arg1: index, %arg2: f64) -> tensor<5x6xf64, #Coo> {
+  %0 = sparse_tensor.insert %arg2 into %arg0[%arg1, %arg1] : tensor<5x6xf64, #Coo>
+  %1 = sparse_tensor.load %0 hasInserts : tensor<5x6xf64, #Coo>
+  return %1 : tensor<5x6xf64, #Coo>
+}
+
 // CHECK-LABEL: func.func @sparse_nop_convert(
 //  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<?xf32>)
 //       CHECK: return %[[A0]], %[[A1]], %[[A2]], %[[A3]], %[[A4]] :
 //  CHECK-SAME: memref<1xindex>, memref<3xindex>, memref<?xi32>, memref<?xi64>, memref<?xf32>
 func.func @sparse_nop_convert(%arg0: tensor<32xf32, #SparseVector>) -> tensor<?xf32, #SparseVector> {
   %0 = sparse_tensor.convert %arg0 : tensor<32xf32, #SparseVector> to tensor<?xf32, #SparseVector>
   return %0 : tensor<?xf32, #SparseVector>
 }
 No newline at end of file

mlir/test/Dialect/SparseTensor/sparse_matmul_codegen.mlir

 // NOTE: Assertions have been autogenerated by utils/generate-test-checks.py
 
 // RUN: mlir-opt %s --linalg-generalize-named-ops \
 // RUN:  --sparsification --sparse-tensor-codegen \
 // RUN:  --canonicalize --cse | FileCheck %s
 
 #CSR = #sparse_tensor.encoding<{
   dimLevelType = [ "dense", "compressed" ],
   dimOrdering = affine_map<(i,j) -> (i,j)>
 }>
 
 //
 // Computes C = A x B with all matrices sparse (SpMSpM) in CSR.
 //
-// CHECK-LABEL:   func.func private @_insert_D_C_4_4_f64_0_0(
+// CHECK-LABEL:   func.func private @_insert_dense_compressed_4_4_f64_0_0(
 // CHECK-SAME:      %[[VAL_0:.*]]: memref<2xindex>,
 // CHECK-SAME:      %[[VAL_1:.*]]: memref<3xindex>,
 // CHECK-SAME:      %[[VAL_2:[^ ]+]]: memref<?xindex>,
 // CHECK-SAME:      %[[VAL_3:.*]]: memref<?xindex>,
 // CHECK-SAME:      %[[VAL_4:.*]]: memref<?xf64>,
 // CHECK-SAME:      %[[VAL_5:[^ ]+]]: index,
 // CHECK-SAME:      %[[VAL_6:.*]]: index,
 // CHECK-SAME:      %[[VAL_7:.*]]: f64) -> (memref<2xindex>, memref<3xindex>, memref<?xindex>, memref<?xindex>, memref<?xf64>) {

 // CHECK:                 }
 // CHECK:                 memref.store %[[VAL_55]], %[[VAL_26]]{{\[}}%[[VAL_51]]] : memref<4xf64>
 // CHECK:                 scf.yield %[[VAL_60:.*]] : index
 // CHECK:               } {"Emitted from" = "linalg.generic"}
 // CHECK:             sparse_tensor.sort %[[VAL_62:.*]], %[[VAL_29]] : memref<?xindex>
 // CHECK:             %[[VAL_63:.*]]:5 = scf.for %[[VAL_64:.*]] = %[[VAL_12]] to %[[VAL_62]] step %[[VAL_13]] iter_args(%[[VAL_65:.*]] = %[[VAL_32]], %[[VAL_66:.*]] = %[[VAL_33]], %[[VAL_67:.*]] = %[[VAL_34]], %[[VAL_68:.*]] = %[[VAL_35]], %[[VAL_69:.*]] = %[[VAL_36]]) -> (memref<2xindex>, memref<3xindex>, memref<?xindex>, memref<?xindex>, memref<?xf64>) {
 // CHECK:               %[[VAL_70:.*]] = memref.load %[[VAL_28]]{{\[}}%[[VAL_64]]] : memref<4xindex>
 // CHECK:               %[[VAL_71:.*]] = memref.load %[[VAL_26]]{{\[}}%[[VAL_70]]] : memref<4xf64>
-// CHECK:               %[[VAL_72:.*]]:5 = func.call @_insert_D_C_4_4_f64_0_0(%[[VAL_65]], %[[VAL_66]], %[[VAL_67]], %[[VAL_68]], %[[VAL_69]], %[[VAL_31]], %[[VAL_70]], %[[VAL_71]]) : (memref<2xindex>, memref<3xindex>, memref<?xindex>, memref<?xindex>, memref<?xf64>, index, index, f64) -> (memref<2xindex>, memref<3xindex>, memref<?xindex>, memref<?xindex>, memref<?xf64>)
+// CHECK:               %[[VAL_72:.*]]:5 = func.call @_insert_dense_compressed_4_4_f64_0_0(%[[VAL_65]], %[[VAL_66]], %[[VAL_67]], %[[VAL_68]], %[[VAL_69]], %[[VAL_31]], %[[VAL_70]], %[[VAL_71]]) : (memref<2xindex>, memref<3xindex>, memref<?xindex>, memref<?xindex>, memref<?xf64>, index, index, f64) -> (memref<2xindex>, memref<3xindex>, memref<?xindex>, memref<?xindex>, memref<?xf64>)
 // CHECK:               memref.store %[[VAL_11]], %[[VAL_26]]{{\[}}%[[VAL_70]]] : memref<4xf64>
 // CHECK:               memref.store %[[VAL_14]], %[[VAL_27]]{{\[}}%[[VAL_70]]] : memref<4xi1>
 // CHECK:               scf.yield %[[VAL_72]]#0, %[[VAL_72]]#1, %[[VAL_72]]#2, %[[VAL_72]]#3, %[[VAL_72]]#4 : memref<2xindex>, memref<3xindex>, memref<?xindex>, memref<?xindex>, memref<?xf64>
 // CHECK:             }
 // CHECK:             scf.yield %[[VAL_73:.*]]#0, %[[VAL_73]]#1, %[[VAL_73]]#2, %[[VAL_73]]#3, %[[VAL_73]]#4 : memref<2xindex>, memref<3xindex>, memref<?xindex>, memref<?xindex>, memref<?xf64>
 // CHECK:           } {"Emitted from" = "linalg.generic"}
 // CHECK:           memref.dealloc %[[VAL_26]] : memref<4xf64>
 // CHECK:           memref.dealloc %[[VAL_27]] : memref<4xi1>