[mlir][sparse] partially implement codegen for sparse_tensor.compress

[aartbik]

removed TABs

[aartbik]

better var naming

[aartbik]

fix

[aartbik]

rebased with main

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

builder.create<memref::StoreOp>(
         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);
 }
 
+/// Creates a straightforward counting for-loop.
+static scf::ForOp createFor(OpBuilder &builder, Location loc, Value count) {
+  Type indexType = builder.getIndexType();
+  Value zero = constantZero(builder, loc, indexType);
+  Value one = constantOne(builder, loc, indexType);
+  scf::ForOp forOp = builder.create<scf::ForOp>(loc, zero, count, one);
+  builder.setInsertionPointToStart(forOp.getBody());
+  return forOp;
+}
+
 //===----------------------------------------------------------------------===//
 // Codegen rules.
 //===----------------------------------------------------------------------===//
 
 /// Sparse tensor storage conversion rule for returns.
 class SparseReturnConverter : public OpConversionPattern<func::ReturnOp> {
 public:
   using OpConversionPattern::OpConversionPattern;
matchAndRewrite(ExpandOp op, OpAdaptor adaptor,
     auto sz = sizeFromTensorAtDim(rewriter, loc, srcType, adaptor.getTensor(),
                                   innerDim);
     assert(sz); // This for sure is a sparse tensor
     // Generate a memref for `sz` elements of type `t`.
     auto genAlloc = [&](Type t) {
       auto memTp = MemRefType::get({ShapedType::kDynamicSize}, t);
       return rewriter.create<memref::AllocOp>(loc, memTp, ValueRange{*sz});
     };
-    // Allocate temporary buffers for values, filled-switch, and indices.
+    // Allocate temporary buffers for values/filled-switch and added.
     // We do not use stack buffers for this, since the expanded size may
     // be rather large (as it envelops a single expanded dense dimension).
     Value values = genAlloc(eltType);
     Value filled = genAlloc(boolType);
-    Value indices = genAlloc(idxType);
+    Value added = genAlloc(idxType);
     Value zero = constantZero(rewriter, loc, idxType);
     // Reset the values/filled-switch to all-zero/false. Note that this
     // introduces an O(N) operation into the computation, but this reset
     // operation is amortized over the innermost loops for the access
     // pattern expansion. As noted in the operation doc, we would like
     // to amortize this setup cost even between kernels.
     rewriter.create<linalg::FillOp>(
         loc, ValueRange{constantZero(rewriter, loc, eltType)},
         ValueRange{values});
     rewriter.create<linalg::FillOp>(
         loc, ValueRange{constantZero(rewriter, loc, boolType)},
         ValueRange{filled});
     // Replace expansion op with these buffers and initial index.
     assert(op.getNumResults() == 4);
-    rewriter.replaceOp(op, {values, filled, indices, zero});
+    rewriter.replaceOp(op, {values, filled, added, zero});
+    return success();
+  }
+};
+
+/// Sparse codegen rule for the compress operator.
+class SparseCompressConverter : public OpConversionPattern<CompressOp> {
+public:
+  using OpConversionPattern::OpConversionPattern;
+  LogicalResult
+  matchAndRewrite(CompressOp op, OpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override {
+    Location loc = op->getLoc();
+    ShapedType srcType = op.getTensor().getType().cast<ShapedType>();
+    Type eltType = srcType.getElementType();
+    Value values = adaptor.getValues();
+    Value filled = adaptor.getFilled();
+    Value added = adaptor.getAdded();
+    Value count = adaptor.getCount();
+
+    //
+    // TODO: need to implement "std::sort(added, added + count);" for ordered
+    //
+
+    // While performing the insertions, we also need to reset the elements
+    // of the values/filled-switch by only iterating over the set elements,
+    // to ensure that the runtime complexity remains proportional to the
+    // sparsity of the expanded access pattern.
+    //
+    // Generate
+    //    for (i = 0; i < count; i++) {
+    //      index = added[i];
+    //      value = values[index];
+    //
+    //      TODO: insert prev_indices, index, value
+    //
+    //      values[index] = 0;
+    //      filled[index] = false;
+    //    }
+    Value i = createFor(rewriter, loc, count).getInductionVar();
+    Value index = rewriter.create<memref::LoadOp>(loc, added, i);
+    rewriter.create<memref::LoadOp>(loc, values, index);
+    // TODO: insert
+    rewriter.create<memref::StoreOp>(loc, constantZero(rewriter, loc, eltType),
+                                     values, index);
+    rewriter.create<memref::StoreOp>(loc, constantI1(rewriter, loc, false),
+                                     filled, index);
+
+    // Deallocate the buffers on exit of the full loop nest.
+    Operation *parent = op;
+    for (; isa<scf::ForOp>(parent->getParentOp()) ||
+           isa<scf::WhileOp>(parent->getParentOp()) ||
+           isa<scf::ParallelOp>(parent->getParentOp()) ||
+           isa<scf::IfOp>(parent->getParentOp());
+         parent = parent->getParentOp())
+      ;
+    rewriter.setInsertionPointAfter(parent);
+    rewriter.create<memref::DeallocOp>(loc, values);
+    rewriter.create<memref::DeallocOp>(loc, filled);
+    rewriter.create<memref::DeallocOp>(loc, added);
+    rewriter.eraseOp(op);
     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:

 
 /// 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, SparseTensorAllocConverter,
                SparseTensorDeallocConverter, SparseTensorLoadConverter,
-               SparseExpandConverter, SparseToPointersConverter,
-               SparseToIndicesConverter, SparseToValuesConverter>(
-      typeConverter, patterns.getContext());
+               SparseExpandConverter, SparseCompressConverter,
+               SparseToPointersConverter, SparseToIndicesConverter,
+               SparseToValuesConverter>(typeConverter, patterns.getContext());
 }

mlir/test/Dialect/SparseTensor/codegen.mlir

 //       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>
 }
+
+// CHECK-LABEL: func @sparse_compression(
+//  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]]: memref<?xindex>,
+//  CHECK-SAME: %[[A6:.*6]]: memref<?xf64>,
+//  CHECK-SAME: %[[A7:.*7]]: memref<?xi1>,
+//  CHECK-SAME: %[[A8:.*8]]: memref<?xindex>,
+//  CHECK-SAME: %[[A9:.*9]]: index)
+//   CHECK-DAG: %[[B0:.*]] = arith.constant false
+//   CHECK-DAG: %[[F0:.*]] = arith.constant 0.000000e+00 : f64
+//   CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
+//   CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index
+//        TODO: sort
+//  CHECK-NEXT: scf.for %[[I:.*]] = %[[C0]] to %[[A9]] step %[[C1]] {
+//  CHECK-NEXT:   %[[INDEX:.*]] = memref.load %[[A8]][%[[I]]] : memref<?xindex>
+//        TODO:   insert
+//   CHECK-DAG:   memref.store %[[F0]], %[[A6]][%[[INDEX]]] : memref<?xf64>
+//   CHECK-DAG:   memref.store %[[B0]], %[[A7]][%[[INDEX]]] : memref<?xi1>
+//  CHECK-NEXT: }
+//   CHECK-DAG: memref.dealloc %[[A6]] : memref<?xf64>
+//   CHECK-DAG: memref.dealloc %[[A7]] : memref<?xi1>
+//   CHECK-DAG: memref.dealloc %[[A8]] : memref<?xindex>
+//       CHECK: return
+func.func @sparse_compression(%arg0: tensor<8x8xf64, #CSR>,
+                              %arg1: memref<?xindex>,
+                              %arg2: memref<?xf64>,
+                              %arg3: memref<?xi1>,
+                              %arg4: memref<?xindex>,
+                              %arg5: index) {
+  sparse_tensor.compress %arg0, %arg1, %arg2, %arg3, %arg4, %arg5
+    : tensor<8x8xf64, #CSR>, memref<?xindex>, memref<?xf64>, memref<?xi1>, memref<?xindex>, index
+  return
+}