This is an archive of the discontinued LLVM Phabricator instance.

Differential D147358

[mlir][sparse] support sparse bufferization.alloc_tensor with copy argument.
ClosedPublic

Authored by Peiming on Mar 31 2023, 2:45 PM.

Details

Reviewers
aartbik
nicolasvasilache
wrengr
bixia
Commits
rG7b86f7c5d4c4: [mlir][sparse] support sparse bufferization.alloc_tensor with copy argument.

Diff Detail

Event Timeline

Peiming created this revision.Mar 31 2023, 2:45 PM
Herald added a reviewer: aartbik. · View Herald TranscriptMar 31 2023, 2:45 PM
Herald added a project: Restricted Project. · View Herald Transcript
Herald added subscribers: hanchung, jsetoain, Moerafaat and 23 others. · View Herald Transcript
Peiming requested review of this revision.Mar 31 2023, 2:45 PM
Herald added a reviewer: nicolasvasilache. · View Herald TranscriptMar 31 2023, 2:45 PM
Herald added a project: Restricted Project. · View Herald Transcript
Herald added subscribers: stephenneuendorffer, nicolasvasilache. · View Herald Transcript
Peiming added reviewers: wrengr, bixia.Mar 31 2023, 2:46 PM
Harbormaster completed remote builds in B223080: Diff 510134.Mar 31 2023, 2:46 PM
Peiming updated this revision to Diff 510135.Mar 31 2023, 2:48 PM

rebase.

Harbormaster completed remote builds in B223081: Diff 510135.Mar 31 2023, 3:12 PM
aartbik accepted this revision.Mar 31 2023, 3:13 PM
This revision is now accepted and ready to land.Mar 31 2023, 3:13 PM
Closed by commit rG7b86f7c5d4c4: [mlir][sparse] support sparse bufferization.alloc_tensor with copy argument. (authored by Peiming). · Explain WhyMar 31 2023, 3:27 PM
This revision was automatically updated to reflect the committed changes.
Peiming added a commit: rG7b86f7c5d4c4: [mlir][sparse] support sparse bufferization.alloc_tensor with copy argument..

Revision Contents

PathSize
mlir/
lib/
Dialect/
SparseTensor/
Transforms/
6 lines
test/
Dialect/
SparseTensor/
44 lines

Diff 510134

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

Show First 20 Lines • Show All 750 Lines • ▼ Show 20 Lines if (!resType.hasEncoding())
return failure(); return failure();
// Construct allocation for each field. // Construct allocation for each field.
const Location loc = op.getLoc(); const Location loc = op.getLoc();
if (op.getCopy()) { if (op.getCopy()) {
auto desc = getDescriptorFromTensorTuple(adaptor.getCopy()); auto desc = getDescriptorFromTensorTuple(adaptor.getCopy());
SmallVector<Value> fields; SmallVector<Value> fields;
fields.reserve(desc.getNumFields()); fields.reserve(desc.getNumFields());
// Memcpy on memref fields.
for (auto field : desc.getMemRefFields()) { for (auto field : desc.getMemRefFields()) {
auto memrefTp = field.getType().cast<MemRefType>(); auto memrefTp = field.getType().cast<MemRefType>();
auto size = rewriter.create<memref::DimOp>(loc, field, 0); auto size = rewriter.create<memref::DimOp>(loc, field, 0);
auto copied = rewriter.create<memref::AllocOp>(loc, memrefTp, ValueRange{size}); auto copied =
rewriter.create<memref::AllocOp>(loc, memrefTp, ValueRange{size});
rewriter.create<memref::CopyOp>(loc, field, copied); rewriter.create<memref::CopyOp>(loc, field, copied);
fields.push_back(copied); fields.push_back(copied);
} }
// Reuses specifier.
fields.push_back(desc.getSpecifier()); fields.push_back(desc.getSpecifier());
assert(fields.size() == desc.getNumFields()); assert(fields.size() == desc.getNumFields());
rewriter.replaceOp(op, genTuple(rewriter, loc, resType, fields)); rewriter.replaceOp(op, genTuple(rewriter, loc, resType, fields));
return success(); return success();
} }
const Value sizeHint = op.getSizeHint(); const Value sizeHint = op.getSizeHint();
const ValueRange dynSizes = adaptor.getDynamicSizes(); const ValueRange dynSizes = adaptor.getDynamicSizes();
▲ Show 20 Lines • Show All 772 Lines • Show Last 20 Lines

mlir/test/Dialect/SparseTensor/codegen_sparse_alloc.mlir

  • This file was added.
// 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<?xindex>,
// CHECK-SAME: %[[VAL_1:.*1]]: memref<?xindex>,
// CHECK-SAME: %[[VAL_2:.*2]]: memref<?xf32>,
// CHECK-SAME: %[[VAL_3:.*]]: !sparse_tensor.storage_specifier<#{{.*}}>) -> (memref<?xindex>, memref<?xindex>, memref<?xf32>, !sparse_tensor.storage_specifier<#{{.*}}>) {
// CHECK: %[[VAL_4:.*]] = arith.constant 0 : index
// CHECK: %[[VAL_5:.*]] = memref.dim %[[VAL_0]], %[[VAL_4]] : memref<?xindex>
// CHECK: %[[VAL_6:.*]] = memref.alloc(%[[VAL_5]]) : memref<?xindex>
// CHECK: memref.copy %[[VAL_0]], %[[VAL_6]] : memref<?xindex> to memref<?xindex>
// CHECK: %[[VAL_7:.*]] = memref.dim %[[VAL_1]], %[[VAL_4]] : memref<?xindex>
// CHECK: %[[VAL_8:.*]] = memref.alloc(%[[VAL_7]]) : memref<?xindex>
// CHECK: memref.copy %[[VAL_1]], %[[VAL_8]] : memref<?xindex> to memref<?xindex>
// CHECK: %[[VAL_9:.*]] = memref.dim %[[VAL_2]], %[[VAL_4]] : memref<?xf32>
// CHECK: %[[VAL_10:.*]] = memref.alloc(%[[VAL_9]]) : memref<?xf32>
// CHECK: memref.copy %[[VAL_2]], %[[VAL_10]] : memref<?xf32> to memref<?xf32>
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<?xindex>,
// CHECK-SAME: %[[VAL_1:.*1]]: memref<?xindex>,
// CHECK-SAME: %[[VAL_2:.*2]]: memref<?xf32>,
// CHECK-SAME: %[[VAL_3:.*]]: !sparse_tensor.storage_specifier<#{{.*}}>) -> (memref<?xindex>, memref<?xindex>, memref<?xf32>, !sparse_tensor.storage_specifier<#{{.*}}>) {
// CHECK: %[[VAL_4:.*]] = arith.constant 0 : index
// CHECK: %[[VAL_5:.*]] = memref.dim %[[VAL_0]], %[[VAL_4]] : memref<?xindex>
// CHECK: %[[VAL_6:.*]] = memref.alloc(%[[VAL_5]]) : memref<?xindex>
// CHECK: memref.copy %[[VAL_0]], %[[VAL_6]] : memref<?xindex> to memref<?xindex>
// CHECK: %[[VAL_7:.*]] = memref.dim %[[VAL_1]], %[[VAL_4]] : memref<?xindex>
// CHECK: %[[VAL_8:.*]] = memref.alloc(%[[VAL_7]]) : memref<?xindex>
// CHECK: memref.copy %[[VAL_1]], %[[VAL_8]] : memref<?xindex> to memref<?xindex>
// CHECK: %[[VAL_9:.*]] = memref.dim %[[VAL_2]], %[[VAL_4]] : memref<?xf32>
// CHECK: %[[VAL_10:.*]] = memref.alloc(%[[VAL_9]]) : memref<?xf32>
// CHECK: memref.copy %[[VAL_2]], %[[VAL_10]] : memref<?xf32> to memref<?xf32>
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>) -> ()
}