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 @@ -756,14 +756,16 @@ 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}); + 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)); 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>) -> () +}