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