100 lines · c
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-exception6//7//===----------------------------------------------------------------------===//8 9#ifndef MLIR_DIALECT_SPARSETENSOR_IR_DETAIL_DIMLVLMAPPARSER_H10#define MLIR_DIALECT_SPARSETENSOR_IR_DETAIL_DIMLVLMAPPARSER_H11 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-vars25/// {l0, ...} // optional forward decl lvl-vars26/// (27/// d0 = ..., // dim-var = dim-exp28/// ...29/// ) -> (30/// l0 = ..., // lvl-var = lvl-exp31/// ...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 map41 // 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, if53 /// `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 of59 /// forward-declared variables.60 FailureOr<VarInfo::ID> parseVarBinding(VarKind vk, bool requireKnown = false);61 62 /// Parses an optional variable binding. When the next token is63 /// 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, and69 /// the `{dims,lvls}AndSymbols` lists (which will be passed70 /// to `AsmParser::parseAffineExpr`). This method is already called by the71 /// `parseVarBinding`/`parseOptionalVarBinding` methods, therefore should72 /// 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 avoid87 // the O(n^2) cost of repeatedly constructing them inside of the88 // `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_detail96} // namespace sparse_tensor97} // namespace mlir98 99#endif // MLIR_DIALECT_SPARSETENSOR_IR_DETAIL_DIMLVLMAPPARSER_H100