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 @@ -749,11 +749,29 @@ const auto resType = getSparseTensorType(op); if (!resType.hasEncoding()) return failure(); - if (op.getCopy()) - return rewriter.notifyMatchFailure(op, "tensor copy not implemented"); // Construct allocation for each field. const Location loc = op.getLoc(); + if (op.getCopy()) { + auto desc = getDescriptorFromTensorTuple(adaptor.getCopy()); + SmallVector fields; + fields.reserve(desc.getNumFields()); + // Memcpy on memref fields. + for (auto field : desc.getMemRefFields()) { + auto memrefTp = field.getType().cast(); + auto size = rewriter.create(loc, field, 0); + auto copied = + rewriter.create(loc, memrefTp, ValueRange{size}); + rewriter.create(loc, field, copied); + fields.push_back(copied); + } + // Reuses specifier. + fields.push_back(desc.getSpecifier()); + assert(fields.size() == desc.getNumFields()); + rewriter.replaceOp(op, genTuple(rewriter, loc, resType, fields)); + return success(); + } + const Value sizeHint = op.getSizeHint(); const ValueRange dynSizes = adaptor.getDynamicSizes(); const size_t found = dynSizes.size(); diff --git a/mlir/test/Dialect/SparseTensor/codegen_sparse_alloc.mlir b/mlir/test/Dialect/SparseTensor/codegen_sparse_alloc.mlir new file mode 100644 --- /dev/null +++ b/mlir/test/Dialect/SparseTensor/codegen_sparse_alloc.mlir @@ -0,0 +1,44 @@ +// RUN: mlir-opt %s --sparse-tensor-codegen --canonicalize --cse | FileCheck %s + +#CSR = #sparse_tensor.encoding<{ dimLevelType = ["dense", "compressed"]}> +#COO = #sparse_tensor.encoding<{ dimLevelType = ["compressed-nu", "singleton"]}> + +// CHECK-LABEL: func.func @sparse_alloc_copy_CSR( +// CHECK-SAME: %[[VAL_0:.*0]]: memref, +// CHECK-SAME: %[[VAL_1:.*1]]: memref, +// CHECK-SAME: %[[VAL_2:.*2]]: memref, +// CHECK-SAME: %[[VAL_3:.*]]: !sparse_tensor.storage_specifier<#{{.*}}>) -> (memref, memref, memref, !sparse_tensor.storage_specifier<#{{.*}}>) { +// CHECK: %[[VAL_4:.*]] = arith.constant 0 : index +// CHECK: %[[VAL_5:.*]] = memref.dim %[[VAL_0]], %[[VAL_4]] : memref +// CHECK: %[[VAL_6:.*]] = memref.alloc(%[[VAL_5]]) : memref +// CHECK: memref.copy %[[VAL_0]], %[[VAL_6]] : memref to memref +// CHECK: %[[VAL_7:.*]] = memref.dim %[[VAL_1]], %[[VAL_4]] : memref +// CHECK: %[[VAL_8:.*]] = memref.alloc(%[[VAL_7]]) : memref +// CHECK: memref.copy %[[VAL_1]], %[[VAL_8]] : memref to memref +// CHECK: %[[VAL_9:.*]] = memref.dim %[[VAL_2]], %[[VAL_4]] : memref +// CHECK: %[[VAL_10:.*]] = memref.alloc(%[[VAL_9]]) : memref +// CHECK: memref.copy %[[VAL_2]], %[[VAL_10]] : memref to memref +func.func @sparse_alloc_copy_CSR(%arg0: tensor<2x2xf32, #CSR>) -> tensor<2x2xf32, #CSR> { + %0 = bufferization.alloc_tensor() copy(%arg0) : tensor<2x2xf32, #CSR> + "test.sink"(%0) : (tensor<2x2xf32, #CSR>) -> () +} + +// CHECK-LABEL: func.func @sparse_alloc_copy_COO( +// CHECK-SAME: %[[VAL_0:.*0]]: memref, +// CHECK-SAME: %[[VAL_1:.*1]]: memref, +// CHECK-SAME: %[[VAL_2:.*2]]: memref, +// CHECK-SAME: %[[VAL_3:.*]]: !sparse_tensor.storage_specifier<#{{.*}}>) -> (memref, memref, memref, !sparse_tensor.storage_specifier<#{{.*}}>) { +// CHECK: %[[VAL_4:.*]] = arith.constant 0 : index +// CHECK: %[[VAL_5:.*]] = memref.dim %[[VAL_0]], %[[VAL_4]] : memref +// CHECK: %[[VAL_6:.*]] = memref.alloc(%[[VAL_5]]) : memref +// CHECK: memref.copy %[[VAL_0]], %[[VAL_6]] : memref to memref +// CHECK: %[[VAL_7:.*]] = memref.dim %[[VAL_1]], %[[VAL_4]] : memref +// CHECK: %[[VAL_8:.*]] = memref.alloc(%[[VAL_7]]) : memref +// CHECK: memref.copy %[[VAL_1]], %[[VAL_8]] : memref to memref +// CHECK: %[[VAL_9:.*]] = memref.dim %[[VAL_2]], %[[VAL_4]] : memref +// CHECK: %[[VAL_10:.*]] = memref.alloc(%[[VAL_9]]) : memref +// CHECK: memref.copy %[[VAL_2]], %[[VAL_10]] : memref to memref +func.func @sparse_alloc_copy_COO(%arg0: tensor<2x2xf32, #COO>) -> tensor<2x2xf32, #COO> { + %0 = bufferization.alloc_tensor() copy(%arg0) : tensor<2x2xf32, #COO> + "test.sink"(%0) : (tensor<2x2xf32, #COO>) -> () +}