[mlir][sparse] Avoid values buffer reallocation for annotated all dense tensors.

[bixia]

Summary

Previously, we rely on the InsertOp to gradually increase the size of the
storage for all sparse tensors. We now allocate the full size values buffer
for annotated all dense tensors when we first allocate the tensor. This avoids
the cost of gradually increasing the buffer and allows accessing the values
buffer as if it were a dense tensor.

[Peiming]

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

244

LG, but you can probably do the same optimization for idxMemRef at the last level, right?

[Peiming]

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

244

NVM, you are doing this for all dense tensors.

Probably we should extend the alloc tensor to take some heuristic, there are many time when the NNZ can be computed.

[bixia]

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

244

Right, there is already a TODO in the code for improving the heuristic.

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

 ///
 /// TODO: for efficiency, we will need heuristis to make educated guesses
 ///       on the required capacities (see heuristic variable).
 ///
 static void createAllocFields(OpBuilder &builder, Location loc, Type type,
                               ValueRange dynSizes, bool enableInit,
                               SmallVectorImpl<Value> &fields) {
   RankedTensorType rtp = type.cast<RankedTensorType>();
-  Value heuristic = constantIndex(builder, loc, 16);
+  // Build original sizes.
+  SmallVector<Value> sizes;
+  auto shape = rtp.getShape();
+  unsigned rank = shape.size();
+  for (unsigned r = 0, o = 0; r < rank; r++) {
+    if (ShapedType::isDynamic(shape[r]))
+      sizes.push_back(dynSizes[o++]);
+    else
+      sizes.push_back(constantIndex(builder, loc, shape[r]));
+  }
 
+  Value heuristic = constantIndex(builder, loc, 16);
+  Value valHeuristic = heuristic;
+  SparseTensorEncodingAttr enc = getSparseTensorEncoding(rtp);
+  if (enc.isAllDense()) {
+    Value linear = sizes[0];
+    for (unsigned r = 1; r < rank; r++) {
+      linear = builder.create<arith::MulIOp>(loc, linear, sizes[r]);
+    }
+    valHeuristic = linear;
+  }
   foreachFieldAndTypeInSparseTensor(
       rtp,
-      [&builder, &fields, rtp, loc, heuristic,
+      [&builder, &fields, rtp, loc, heuristic, valHeuristic,
        enableInit](Type fType, unsigned fIdx, SparseTensorFieldKind fKind,
                    unsigned /*dim*/, DimLevelType /*dlt*/) -> bool {
         assert(fields.size() == fIdx);
         Value field;
         switch (fKind) {
         case SparseTensorFieldKind::StorageSpec:
           field = SparseTensorSpecifier::getInitValue(builder, loc, rtp);
           break;
         case SparseTensorFieldKind::PtrMemRef:
         case SparseTensorFieldKind::IdxMemRef:
         case SparseTensorFieldKind::ValMemRef:
           field = createAllocation(builder, loc, fType.cast<MemRefType>(),
-                                   heuristic, enableInit);
+                                   fKind == SparseTensorFieldKind::ValMemRef
+                                       ? valHeuristic
+                                       : heuristic,
+                                   enableInit);
           break;
         }
         assert(field);
         fields.push_back(field);
         // Returns true to continue the iteration.
         return true;
       });
 
   MutSparseTensorDescriptor desc(rtp, fields);
 
-  // Build original sizes.
-  SmallVector<Value> sizes;
-  auto shape = rtp.getShape();
-  unsigned rank = shape.size();
-  for (unsigned r = 0, o = 0; r < rank; r++) {
-    if (ShapedType::isDynamic(shape[r]))
-      sizes.push_back(dynSizes[o++]);
-    else
-      sizes.push_back(constantIndex(builder, loc, shape[r]));
-  }
   // Initialize the storage scheme to an empty tensor. Initialized memSizes
   // to all zeros, sets the dimSizes to known values and gives all pointer
   // fields an initial zero entry, so that it is easier to maintain the
   // "linear + 1" length property.
   Value ptrZero =
       constantZero(builder, loc, getSparseTensorEncoding(rtp).getPointerType());
   for (unsigned r = 0; r < rank; r++) {
     unsigned ro = toOrigDim(rtp, r);

mlir/test/Dialect/SparseTensor/codegen.mlir

 }
 
 // CHECK-LABEL:   func.func @sparse_alloc_3d() -> (memref<?xf64>, !sparse_tensor.storage_specifier
 //       CHECK:     %[[A0:.*]] = arith.constant 6000 : index
 //       CHECK:     %[[A1:.*]] = arith.constant 20 : i64
 //       CHECK:     %[[A2:.*]] = arith.constant 10 : i64
 //       CHECK:     %[[A3:.*]] = arith.constant 30 : i64
 //       CHECK:     %[[A4:.*]] = arith.constant 0.000000e+00 : f64
-//       CHECK:     %[[A5:.*]] = memref.alloc() : memref<16xf64>
-//       CHECK:     %[[A6:.*]] = memref.cast %[[A5]] : memref<16xf64> to memref<?xf64>
+//       CHECK:     %[[A5:.*]] = memref.alloc() : memref<6000xf64>
+//       CHECK:     %[[A6:.*]] = memref.cast %[[A5]] : memref<6000xf64> to memref<?xf64>
 //       CHECK:     %[[A7:.*]] = sparse_tensor.storage_specifier.init : !sparse_tensor.storage_specifier
 //       CHECK:     %[[A8:.*]] = sparse_tensor.storage_specifier.set %[[A7]]  dim_sz at 0 with %[[A3]] : i64, !sparse_tensor.storage_specifier
 //       CHECK:     %[[A9:.*]] = sparse_tensor.storage_specifier.set %[[A8]]  dim_sz at 1 with %[[A2]] : i64, !sparse_tensor.storage_specifier
 //       CHECK:     %[[A10:.*]] = sparse_tensor.storage_specifier.set %[[A9]]  dim_sz at 2 with %[[A1]] : i64, !sparse_tensor.storage_specifier
 //       CHECK:     %[[A11:.*]] = sparse_tensor.storage_specifier.get %[[A10]]  val_mem_sz : !sparse_tensor.storage_specifier
 //       CHECK:     %[[A12:.*]] = arith.index_cast %[[A11]] : i64 to index
 //       CHECK:     %[[A13:.*]], %[[A14:.*]] = sparse_tensor.push_back %[[A12]], %[[A6]], %[[A4]], %[[A0]] : index, memref<?xf64>, f64, index
 //       CHECK:     %[[A15:.*]] = arith.index_cast %[[A14]] : index to i64