[mlir][sparse] code cleanup using the assumption that dim2lvl maps ar… (#72894)

…e simplified.
This commit is contained in:
Peiming Liu
2023-11-20 10:25:42 -08:00
committed by GitHub
parent 445f6f1373
commit 2cc4b3d07c
2 changed files with 16 additions and 20 deletions

View File

@@ -745,8 +745,8 @@ public:
const auto resType = getSparseTensorType(op);
if (!resType.hasEncoding())
return failure();
Location loc = op.getLoc();
Location loc = op.getLoc();
// Deal with copy.
if (op.getCopy()) {
auto desc = getDescriptorFromTensorTuple(adaptor.getCopy());
@@ -768,16 +768,14 @@ public:
return success();
}
// Construct the dim/lvl sizes and the (unused) dim2lvl/lvl2dim buffers.
SmallVector<Value> dimSizesValues;
if (!resType.isIdentity()) {
return rewriter.notifyMatchFailure(
op, "try run --sparse-reinterpret-map before codegen");
}
// Level size equals to dimension size since lvl2dim map is an identity map.
SmallVector<Value> lvlSizesValues;
Value dimSizesBuffer;
Value dim2lvlBuffer;
Value lvl2dimBuffer;
createDimSizes(rewriter, loc, resType, adaptor.getDynamicSizes(),
dimSizesValues);
genMapBuffers(rewriter, loc, resType, dimSizesValues, dimSizesBuffer,
lvlSizesValues, dim2lvlBuffer, lvl2dimBuffer);
/*dimSizesValues=*/lvlSizesValues);
// Construct allocation for each field.
Value sizeHint = op.getSizeHint();
@@ -809,19 +807,17 @@ public:
const auto resType = getSparseTensorType(op);
if (!resType.hasEncoding())
return failure();
if (!resType.isIdentity()) {
return rewriter.notifyMatchFailure(
op, "try run --sparse-reinterpret-map before codegen");
}
Location loc = op.getLoc();
// Construct the dim/lvl sizes and the (unused) dim2lvl/lvl2dim buffers.
SmallVector<Value> dimSizesValues;
// Level size equals to dimension size since lvl2dim map is an identity map.
SmallVector<Value> lvlSizesValues;
Value dimSizesBuffer;
Value dim2lvlBuffer;
Value lvl2dimBuffer;
createDimSizes(rewriter, loc, resType, adaptor.getDynamicSizes(),
dimSizesValues);
genMapBuffers(rewriter, loc, resType, dimSizesValues, dimSizesBuffer,
lvlSizesValues, dim2lvlBuffer, lvl2dimBuffer);
/*dimSizesValues=*/lvlSizesValues);
// Construct allocation for each field.
Value sizeHint; // none
SmallVector<Value> fields;