diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorCodegen.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorCodegen.cpp --- a/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorCodegen.cpp +++ b/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorCodegen.cpp @@ -178,9 +178,14 @@ /// Creates allocation operation. static Value createAllocation(OpBuilder &builder, Location loc, - MemRefType memRefType, Value sz, - bool enableInit) { - Value buffer = builder.create(loc, memRefType, sz); + MemRefType memRefType, Value sz, bool enableInit, + bool onStack = false) { + Value buffer; + if (onStack) + buffer = builder.create(loc, memRefType, sz); + else + buffer = builder.create(loc, memRefType, sz); + Type elemType = memRefType.getElementType(); if (enableInit) { Value fillValue = constantZero(builder, loc, elemType); @@ -1002,6 +1007,97 @@ } }; +struct SparsePackOpConverter : public OpConversionPattern { + using OpConversionPattern::OpConversionPattern; + LogicalResult + matchAndRewrite(PackOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + + auto rtp = op.getResult().getType().cast(); + if (!isUniqueCOOType(rtp)) + return op.emitError("Run canonicalizer before lowering PackOp"); + + SmallVector fields; + Location loc = op.getLoc(); + + foreachFieldAndTypeInSparseTensor( + rtp, + [&rewriter, &fields, &op, rtp, + loc](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(rewriter, loc, rtp); + break; + case SparseTensorFieldKind::PtrMemRef: + // TACO-style COO starts with a PtrBuffer + // FIXME: how should we free it? in unpack? + field = createAllocation(rewriter, loc, fType.cast(), + constantIndex(rewriter, loc, 2), false, + /*onStack=*/true); + break; + case SparseTensorFieldKind::IdxMemRef: { + auto tensorType = op.getIndices().getType(); + auto memrefType = MemRefType::get(tensorType.getShape(), + tensorType.getElementType()); + auto idxMemRef = rewriter.create( + op->getLoc(), memrefType, op.getIndices()); + ReassociationIndices reassociation; + for (int i = 0, e = tensorType.getRank(); i < e; i++) + reassociation.push_back(i); + + // Flattened the indices buffer to rank 1. + field = rewriter.create( + loc, idxMemRef, ArrayRef(reassociation)); + break; + } + case SparseTensorFieldKind::ValMemRef: { + auto tensorType = op.getData().getType(); + auto memrefType = MemRefType::get(tensorType.getShape(), + tensorType.getElementType()); + field = rewriter.create( + op->getLoc(), memrefType, op.getData()); + break; + } + } + + assert(field); + if (fType != field.getType()) + field = rewriter.create(loc, fType, field); + fields.push_back(field); + // Returns true to continue the iteration. + return true; + }); + + MutSparseTensorDescriptor desc(rtp, fields); + auto noe = linalg::createOrFoldDimOp(rewriter, loc, op.getData(), 0); + for (unsigned i = 0, e = rtp.getRank(); i < e; i++) { + int dim = rtp.getShape()[i]; + assert(!ShapedType::isDynamic(dim)); + desc.setDimSize(rewriter, loc, i, constantIndex(rewriter, loc, dim)); + if (i == 0) { + desc.setSpecifierField(rewriter, loc, StorageSpecifierKind::PtrMemSize, + i, constantIndex(rewriter, loc, 2)); + } + desc.setSpecifierField(rewriter, loc, StorageSpecifierKind::IdxMemSize, i, + noe); + } + desc.setSpecifierField(rewriter, loc, StorageSpecifierKind::ValMemSize, + std::nullopt, noe); + + // Set up the ptrMemref => [0, noe] + genStore(rewriter, loc, constantIndex(rewriter, loc, 0), + desc.getPtrMemRef(0), constantIndex(rewriter, loc, 0)); + genStore(rewriter, loc, noe, desc.getPtrMemRef(0), + constantIndex(rewriter, loc, 1)); + + rewriter.replaceOp(op, genTuple(rewriter, loc, desc)); + return success(); + } +}; + } // namespace //===----------------------------------------------------------------------===// @@ -1013,14 +1109,15 @@ void mlir::populateSparseTensorCodegenPatterns( TypeConverter &typeConverter, RewritePatternSet &patterns, bool enableBufferInitialization) { - patterns.add( - typeConverter, patterns.getContext()); + patterns.add(typeConverter, + patterns.getContext()); patterns.add(typeConverter, patterns.getContext(), enableBufferInitialization); } diff --git a/mlir/test/Dialect/SparseTensor/sparse_pack.mlir b/mlir/test/Dialect/SparseTensor/sparse_pack.mlir new file mode 100644 --- /dev/null +++ b/mlir/test/Dialect/SparseTensor/sparse_pack.mlir @@ -0,0 +1,40 @@ +// RUN: mlir-opt %s --canonicalize --post-sparsification-rewrite="enable-runtime-library=false" --sparse-tensor-codegen -cse | FileCheck %s + +#CSR = #sparse_tensor.encoding<{ + dimLevelType = ["dense", "compressed"], + indexBitWidth=32 +}> + +// CHECK-LABEL: func.func private @_insert_dense_compressed_100_100_f64_32_0( + +// CHECK-LABEL: func.func @sparse_pack( +// CHECK-SAME: %[[VAL_0:.*]]: tensor<6xf64>, +// CHECK-SAME: %[[VAL_1:.*]]: tensor<6x2xi32>) +// CHECK: %[[VAL_2:.*]] = arith.constant 0 : index +// CHECK: %[[VAL_3:.*]] = arith.constant 1 : index +// CHECK: %[[VAL_4:.*]] = arith.constant 2 : index +// CHECK: %[[VAL_5:.*]] = memref.alloca(%[[VAL_4]]) : memref +// CHECK: %[[VAL_6:.*]] = bufferization.to_memref %[[VAL_1]] : memref<6x2xi32> +// CHECK: %[[VAL_7:.*]] = memref.collapse_shape %[[VAL_6]] {{\[\[}}0, 1]] : memref<6x2xi32> into memref<12xi32> +// CHECK: %[[VAL_8:.*]] = memref.cast %[[VAL_7]] : memref<12xi32> to memref +// CHECK: %[[VAL_9:.*]] = bufferization.to_memref %[[VAL_0]] : memref<6xf64> +// CHECK: %[[VAL_10:.*]] = memref.cast %[[VAL_9]] : memref<6xf64> to memref +// CHECK: %[[VAL_11:.*]] = sparse_tensor.storage_specifier.init : !sparse_tensor.storage_specifier<#sparse_tensor.encoding<{ dimLevelType = [ "compressed", "singleton" ], indexBitWidth = 32 }>> +// CHECK: %[[VAL_12:.*]] = arith.constant 6 : index +// CHECK: %[[VAL_13:.*]] = arith.constant 100 : index +// CHECK: %[[VAL_14:.*]] = arith.index_cast %[[VAL_13]] : index to i32 +// CHECK: %[[VAL_15:.*]] = sparse_tensor.storage_specifier.set %[[VAL_11]] dim_sz at 0 with %[[VAL_14]] : i32, !sparse_tensor.storage_specifier<#sparse_tensor.encoding<{ dimLevelType = [ "compressed", "singleton" ], indexBitWidth = 32 }>> +// CHECK: %[[VAL_16:.*]] = arith.index_cast %[[VAL_4]] : index to i32 +// CHECK: %[[VAL_17:.*]] = sparse_tensor.storage_specifier.set %[[VAL_15]] ptr_mem_sz at 0 with %[[VAL_16]] : i32, !sparse_tensor.storage_specifier<#sparse_tensor.encoding<{ dimLevelType = [ "compressed", "singleton" ], indexBitWidth = 32 }>> +// CHECK: %[[VAL_18:.*]] = arith.index_cast %[[VAL_12]] : index to i32 +// CHECK: %[[VAL_19:.*]] = sparse_tensor.storage_specifier.set %[[VAL_17]] idx_mem_sz at 0 with %[[VAL_18]] : i32, !sparse_tensor.storage_specifier<#sparse_tensor.encoding<{ dimLevelType = [ "compressed", "singleton" ], indexBitWidth = 32 }>> +// CHECK: %[[VAL_20:.*]] = sparse_tensor.storage_specifier.set %[[VAL_19]] dim_sz at 1 with %[[VAL_14]] : i32, !sparse_tensor.storage_specifier<#sparse_tensor.encoding<{ dimLevelType = [ "compressed", "singleton" ], indexBitWidth = 32 }>> +// CHECK: %[[VAL_21:.*]] = sparse_tensor.storage_specifier.set %[[VAL_20]] idx_mem_sz at 1 with %[[VAL_18]] : i32, !sparse_tensor.storage_specifier<#sparse_tensor.encoding<{ dimLevelType = [ "compressed", "singleton" ], indexBitWidth = 32 }>> +// CHECK: %[[VAL_22:.*]] = sparse_tensor.storage_specifier.set %[[VAL_21]] val_mem_sz with %[[VAL_18]] : i32, !sparse_tensor.storage_specifier<#sparse_tensor.encoding<{ dimLevelType = [ "compressed", "singleton" ], indexBitWidth = 32 }>> +// +func.func @sparse_pack(%data: tensor<6xf64>, %index: tensor<6x2xi32>) + -> tensor<100x100xf64, #CSR> { + %0 = sparse_tensor.pack %data, %index : tensor<6xf64>, tensor<6x2xi32> + to tensor<100x100xf64, #CSR> + return %0 : tensor<100x100xf64, #CSR> +} diff --git a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_pack.mlir b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_pack.mlir new file mode 100644 --- /dev/null +++ b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_pack.mlir @@ -0,0 +1,55 @@ +// DEFINE: %{option} = enable-runtime-library=false +// DEFINE: %{command} = mlir-opt %s --sparse-compiler=%{option} | \ +// DEFINE: mlir-cpu-runner \ +// DEFINE: -e entry -entry-point-result=void \ +// DEFINE: -shared-libs=%mlir_lib_dir/libmlir_c_runner_utils%shlibext | \ +// DEFINE: FileCheck %s +// +// RUN: %{command} +// + +// TODO: Pack only support CodeGen Path + +#SortedCOO = #sparse_tensor.encoding<{ + dimLevelType = [ "compressed-nu", "singleton" ] +}> + +module { + // + // Main driver. + // + func.func @entry() { + // + // Initialize a 3-dim dense tensor. + // + %data = arith.constant dense< + [ 1.0, 2.0, 3.0] + > : tensor<3xf64> + + %index = arith.constant dense< + [[ 1, 2], + [ 5, 6], + [ 7, 8]] + > : tensor<3x2xindex> + + %s4 = sparse_tensor.pack %data, %index : tensor<3xf64>, tensor<3x2xindex> to tensor<10x10xf64, #SortedCOO> + // CHECK:1 + // CHECK-NEXT:2 + // CHECK-NEXT:1 + // + // CHECK-NEXT:5 + // CHECK-NEXT:6 + // CHECK-NEXT:2 + // + // CHECK-NEXT:7 + // CHECK-NEXT:8 + // CHECK-NEXT:3 + sparse_tensor.foreach in %s4 : tensor<10x10xf64, #SortedCOO> do { + ^bb0(%1: index, %2: index, %v: f64) : + vector.print %1: index + vector.print %2: index + vector.print %v: f64 + } + return + } +}