::= '{' 'replicated'? 'maximal'? ('device=' int)? shape? ('devices=' ('[' dims ']')* device_list)? '}' dims ::= int_list device_list ::= int_list
| 1978 | // ('devices=' ('[' dims ']')* device_list)? '}' |
| 1979 | // dims ::= int_list device_list ::= int_list |
| 1980 | bool HloParserImpl::ParseSingleSharding(OpSharding* sharding, |
| 1981 | bool lbrace_pre_lexed) { |
| 1982 | if (!lbrace_pre_lexed && |
| 1983 | !ParseToken(TokKind::kLbrace, |
| 1984 | "expected '{' to start sharding attribute")) { |
| 1985 | return false; |
| 1986 | } |
| 1987 | |
| 1988 | LocTy loc = lexer_.GetLoc(); |
| 1989 | bool maximal = false; |
| 1990 | bool replicated = false; |
| 1991 | std::vector<int64> devices; |
| 1992 | std::vector<int64> tile_assignment_dimensions; |
| 1993 | while (lexer_.GetKind() != TokKind::kRbrace) { |
| 1994 | switch (lexer_.GetKind()) { |
| 1995 | case TokKind::kw_maximal: |
| 1996 | maximal = true; |
| 1997 | lexer_.Lex(); |
| 1998 | break; |
| 1999 | case TokKind::kw_replicated: |
| 2000 | replicated = true; |
| 2001 | lexer_.Lex(); |
| 2002 | break; |
| 2003 | case TokKind::kAttributeName: { |
| 2004 | if (lexer_.GetStrVal() == "device") { |
| 2005 | if (lexer_.Lex() != TokKind::kInt) { |
| 2006 | return TokenError("device= attribute must be an integer"); |
| 2007 | } |
| 2008 | devices = {lexer_.GetInt64Val()}; |
| 2009 | lexer_.Lex(); |
| 2010 | } else if (lexer_.GetStrVal() == "devices") { |
| 2011 | lexer_.Lex(); |
| 2012 | if (!ParseToken(TokKind::kLsquare, |
| 2013 | "expected '[' to start sharding devices shape")) { |
| 2014 | return false; |
| 2015 | } |
| 2016 | |
| 2017 | do { |
| 2018 | int64 dim; |
| 2019 | if (!ParseInt64(&dim)) { |
| 2020 | return false; |
| 2021 | } |
| 2022 | tile_assignment_dimensions.push_back(dim); |
| 2023 | } while (EatIfPresent(TokKind::kComma)); |
| 2024 | |
| 2025 | if (!ParseToken(TokKind::kRsquare, |
| 2026 | "expected ']' to start sharding devices shape")) { |
| 2027 | return false; |
| 2028 | } |
| 2029 | do { |
| 2030 | int64 device; |
| 2031 | if (!ParseInt64(&device)) { |
| 2032 | return false; |
| 2033 | } |
| 2034 | devices.push_back(device); |
| 2035 | } while (EatIfPresent(TokKind::kComma)); |
| 2036 | } else { |
| 2037 | return TokenError( |