1//===- DimLvlMapParser.h - `DimLvlMap` parser -------------------*- C++ -*-===//
2//
3// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4// See https://llvm.org/LICENSE.txt for license information.
5// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6//
7//===----------------------------------------------------------------------===//
8
9#ifndef MLIR_DIALECT_SPARSETENSOR_IR_DETAIL_DIMLVLMAPPARSER_H
10#define MLIR_DIALECT_SPARSETENSOR_IR_DETAIL_DIMLVLMAPPARSER_H
11
12#include "DimLvlMap.h"
13#include "LvlTypeParser.h"
14
15namespace mlir {
16namespace sparse_tensor {
17namespace ir_detail {
18
19///
20/// Parses the Sparse Tensor Encoding Attribute (STEA).
21///
22/// General syntax is as follows,
23///
24/// [s0, ...] // optional forward decl sym-vars
25/// {l0, ...} // optional forward decl lvl-vars
26/// (
27/// d0 = ..., // dim-var = dim-exp
28/// ...
29/// ) -> (
30/// l0 = ..., // lvl-var = lvl-exp
31/// ...
32/// )
33///
34/// with simplifications when variables are implicit.
35///
36class DimLvlMapParser final {
37public:
38 explicit DimLvlMapParser(AsmParser &parser) : parser(parser) {}
39
40 // Parses the input for a sparse tensor dimension-level map
41 // and returns the map on success.
42 FailureOr<DimLvlMap> parseDimLvlMap();
43
44private:
45 /// Client code should prefer using `parseVarUsage`
46 /// and `parseVarBinding` rather than calling this method directly.
47 OptionalParseResult parseVar(VarKind vk, bool isOptional,
48 Policy creationPolicy, VarInfo::ID &id,
49 bool &didCreate);
50
51 /// Parses a variable occurence which is a *use* of that variable.
52 /// When a valid variable name is currently unused, if
53 /// `requireKnown=true`, an error is raised; if `requireKnown=false`,
54 /// a new unbound variable will be created.
55 FailureOr<VarInfo::ID> parseVarUsage(VarKind vk, bool requireKnown);
56
57 /// Parses a variable occurence which is a *binding* of that variable.
58 /// The `requireKnown` parameter is for handling the binding of
59 /// forward-declared variables.
60 FailureOr<VarInfo::ID> parseVarBinding(VarKind vk, bool requireKnown = false);
61
62 /// Parses an optional variable binding. When the next token is
63 /// not a valid variable name, this will bind a new unnamed variable.
64 /// The returned `bool` indicates whether a variable name was parsed.
65 FailureOr<std::pair<Var, bool>>
66 parseOptionalVarBinding(VarKind vk, bool requireKnown = false);
67
68 /// Binds the given variable: both updating the `VarEnv` itself, and
69 /// the `{dims,lvls}AndSymbols` lists (which will be passed
70 /// to `AsmParser::parseAffineExpr`). This method is already called by the
71 /// `parseVarBinding`/`parseOptionalVarBinding` methods, therefore should
72 /// not need to be called elsewhere.
73 Var bindVar(llvm::SMLoc loc, VarInfo::ID id);
74
75 ParseResult parseSymbolBindingList();
76 ParseResult parseLvlVarBindingList();
77 ParseResult parseDimSpec();
78 ParseResult parseDimSpecList();
79 FailureOr<LvlVar> parseLvlVarBinding(bool requireLvlVarBinding);
80 ParseResult parseLvlSpec(bool requireLvlVarBinding);
81 ParseResult parseLvlSpecList();
82
83 AsmParser &parser;
84 LvlTypeParser lvlTypeParser;
85 VarEnv env;
86 // The parser maintains the `{dims,lvls}AndSymbols` lists to avoid
87 // the O(n^2) cost of repeatedly constructing them inside of the
88 // `parse{Dim,Lvl}Spec` methods.
89 SmallVector<std::pair<StringRef, AffineExpr>, 4> dimsAndSymbols;
90 SmallVector<std::pair<StringRef, AffineExpr>, 4> lvlsAndSymbols;
91 SmallVector<DimSpec> dimSpecs;
92 SmallVector<LvlSpec> lvlSpecs;
93};
94
95} // namespace ir_detail
96} // namespace sparse_tensor
97} // namespace mlir
98
99#endif // MLIR_DIALECT_SPARSETENSOR_IR_DETAIL_DIMLVLMAPPARSER_H
100

source code of mlir/lib/Dialect/SparseTensor/IR/Detail/DimLvlMapParser.h