caio/clang-p2996
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

[mlir][sparse] migrate more to new surface syntax

Browse Source 
Replaced the "NEW_SYNTAX" with the more readable "map"
(which we may, or may not keep). Minor improvement in
keyword parsing, migrated a few more examples over.

Reviewed By: Peiming, yinying-lisa-li

Differential Revision: https://reviews.llvm.org/D158325
This commit is contained in:
Aart Bik
2023-08-18 17:03:44 -07:00
parent 58fe7b751d
commit bb44a6b7bb
3 changed files with 57 additions and 95 deletions
Download Patch File Download Diff File Expand all files Collapse all files

55
mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp
View File

@@ -451,24 +451,22 @@ Attribute SparseTensorEncodingAttr::parse(AsmParser &parser, Type type) {
AffineMap dimToLvl = {};
unsigned posWidth = 0;
unsigned crdWidth = 0;
StringRef attrName;
// Exactly 6 keys.
SmallVector<StringRef, 6> keys = {"lvlTypes", "dimToLvl", "posWidth",
"crdWidth", "dimSlices", "NEW_SYNTAX"};
"crdWidth", "dimSlices", "map"};
while (succeeded(parser.parseOptionalKeyword(&attrName))) {
if (!llvm::is_contained(keys, attrName)) {
// Detect admissible keyword.
auto *it = find(keys, attrName);
if (it == keys.end()) {
parser.emitError(parser.getNameLoc(), "unexpected key: ") << attrName;
return {};
}
unsigned keyWordIndex = it - keys.begin();
// Consume the `=` after keys
RETURN_ON_FAIL(parser.parseEqual())
// FIXME: using `operator==` below duplicates the string comparison
// cost of the `is_contained` check above. Should instead use some
// "find" function that returns the index into `keys` so that we can
// dispatch on that instead.
if (attrName == "lvlTypes") {
// Dispatch on keyword.
switch (keyWordIndex) {
case 0: { // lvlTypes
Attribute attr;
RETURN_ON_FAIL(parser.parseAttribute(attr));
auto arrayAttr = llvm::dyn_cast<ArrayAttr>(attr);
@@ -485,25 +483,33 @@ Attribute SparseTensorEncodingAttr::parse(AsmParser &parser, Type type) {
return {};
}
}
} else if (attrName == "dimToLvl") {
break;
}
case 1: { // dimToLvl
Attribute attr;
RETURN_ON_FAIL(parser.parseAttribute(attr))
auto affineAttr = llvm::dyn_cast<AffineMapAttr>(attr);
ERROR_IF(!affineAttr, "expected an affine map for dimToLvl")
dimToLvl = affineAttr.getValue();
} else if (attrName == "posWidth") {
break;
}
case 2: { // posWidth
Attribute attr;
RETURN_ON_FAIL(parser.parseAttribute(attr))
auto intAttr = llvm::dyn_cast<IntegerAttr>(attr);
ERROR_IF(!intAttr, "expected an integral position bitwidth")
posWidth = intAttr.getInt();
} else if (attrName == "crdWidth") {
break;
}
case 3: { // crdWidth
Attribute attr;
RETURN_ON_FAIL(parser.parseAttribute(attr))
auto intAttr = llvm::dyn_cast<IntegerAttr>(attr);
ERROR_IF(!intAttr, "expected an integral index bitwidth")
crdWidth = intAttr.getInt();
} else if (attrName == "dimSlices") {
break;
}
case 4: { // dimSlices
RETURN_ON_FAIL(parser.parseLSquare())
// Dispatches to DimSliceAttr to skip mnemonic
bool finished = false;
@@ -519,13 +525,9 @@ Attribute SparseTensorEncodingAttr::parse(AsmParser &parser, Type type) {
if (!finished)
return {};
RETURN_ON_FAIL(parser.parseRSquare())
} else if (attrName == "NEW_SYNTAX") {
// Note that we are in the process of migrating to a new STEA surface
// syntax. While this is ongoing we use the temporary "NEW_SYNTAX = ...."
// to switch to the new parser. This allows us to gradually migrate
// examples over to the new surface syntax before making the complete
// switch once work is completed.
// TODO: replace everything here with new STEA surface syntax parser
break;
}
case 5: { // map (new STEA surface syntax)
ir_detail::DimLvlMapParser cParser(parser);
auto res = cParser.parseDimLvlMap();
RETURN_ON_FAIL(res);
@@ -533,12 +535,12 @@ Attribute SparseTensorEncodingAttr::parse(AsmParser &parser, Type type) {
// than converting things over.
const auto &dlm = *res;
ERROR_IF(!lvlTypes.empty(), "Cannot mix `lvlTypes` with `NEW_SYNTAX`")
ERROR_IF(!lvlTypes.empty(), "Cannot mix `lvlTypes` with `map`")
const Level lvlRank = dlm.getLvlRank();
for (Level lvl = 0; lvl < lvlRank; lvl++)
lvlTypes.push_back(dlm.getLvlType(lvl));
ERROR_IF(!dimSlices.empty(), "Cannot mix `dimSlices` with `NEW_SYNTAX`")
ERROR_IF(!dimSlices.empty(), "Cannot mix `dimSlices` with `map`")
const Dimension dimRank = dlm.getDimRank();
for (Dimension dim = 0; dim < dimRank; dim++)
dimSlices.push_back(dlm.getDimSlice(dim));
@@ -558,11 +560,12 @@ Attribute SparseTensorEncodingAttr::parse(AsmParser &parser, Type type) {
dimSlices.clear();
}
ERROR_IF(dimToLvl, "Cannot mix `dimToLvl` with `NEW_SYNTAX`")
ERROR_IF(dimToLvl, "Cannot mix `dimToLvl` with `map`")
dimToLvl = dlm.getDimToLvlMap(parser.getContext());
break;
}
// Only the last item can omit the comma
} // switch
// Only last item can omit the comma.
if (parser.parseOptionalComma().failed())
break;
}