1 //===- LLVMDialect.cpp - MLIR SPIR-V dialect ------------------------------===//
2 //
3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM
4 // See https://llvm.org/LICENSE.txt for license information.
5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6 //
7 //===----------------------------------------------------------------------===//
8 //
9 // This file defines the SPIR-V dialect in MLIR.
10 //
11 //===----------------------------------------------------------------------===//
12 
13 #include "mlir/Dialect/SPIRV/IR/SPIRVDialect.h"
14 #include "mlir/Dialect/SPIRV/IR/ParserUtils.h"
15 #include "mlir/Dialect/SPIRV/IR/SPIRVOps.h"
16 #include "mlir/Dialect/SPIRV/IR/SPIRVTypes.h"
17 #include "mlir/Dialect/SPIRV/IR/TargetAndABI.h"
18 #include "mlir/IR/Builders.h"
19 #include "mlir/IR/BuiltinTypes.h"
20 #include "mlir/IR/DialectImplementation.h"
21 #include "mlir/IR/MLIRContext.h"
22 #include "mlir/Parser.h"
23 #include "mlir/Transforms/InliningUtils.h"
24 #include "llvm/ADT/DenseMap.h"
25 #include "llvm/ADT/Sequence.h"
26 #include "llvm/ADT/SetVector.h"
27 #include "llvm/ADT/StringExtras.h"
28 #include "llvm/ADT/StringMap.h"
29 #include "llvm/ADT/StringSwitch.h"
30 #include "llvm/ADT/TypeSwitch.h"
31 #include "llvm/Support/raw_ostream.h"
32 
33 using namespace mlir;
34 using namespace mlir::spirv;
35 
36 //===----------------------------------------------------------------------===//
37 // InlinerInterface
38 //===----------------------------------------------------------------------===//
39 
40 /// Returns true if the given region contains spv.Return or spv.ReturnValue ops.
41 static inline bool containsReturn(Region &region) {
42   return llvm::any_of(region, [](Block &block) {
43     Operation *terminator = block.getTerminator();
44     return isa<spirv::ReturnOp, spirv::ReturnValueOp>(terminator);
45   });
46 }
47 
48 namespace {
49 /// This class defines the interface for inlining within the SPIR-V dialect.
50 struct SPIRVInlinerInterface : public DialectInlinerInterface {
51   using DialectInlinerInterface::DialectInlinerInterface;
52 
53   /// All call operations within SPIRV can be inlined.
54   bool isLegalToInline(Operation *call, Operation *callable,
55                        bool wouldBeCloned) const final {
56     return true;
57   }
58 
59   /// Returns true if the given region 'src' can be inlined into the region
60   /// 'dest' that is attached to an operation registered to the current dialect.
61   bool isLegalToInline(Region *dest, Region *src, bool wouldBeCloned,
62                        BlockAndValueMapping &) const final {
63     // Return true here when inlining into spv.func, spv.mlir.selection, and
64     // spv.mlir.loop operations.
65     auto *op = dest->getParentOp();
66     return isa<spirv::FuncOp, spirv::SelectionOp, spirv::LoopOp>(op);
67   }
68 
69   /// Returns true if the given operation 'op', that is registered to this
70   /// dialect, can be inlined into the region 'dest' that is attached to an
71   /// operation registered to the current dialect.
72   bool isLegalToInline(Operation *op, Region *dest, bool wouldBeCloned,
73                        BlockAndValueMapping &) const final {
74     // TODO: Enable inlining structured control flows with return.
75     if ((isa<spirv::SelectionOp, spirv::LoopOp>(op)) &&
76         containsReturn(op->getRegion(0)))
77       return false;
78     // TODO: we need to filter OpKill here to avoid inlining it to
79     // a loop continue construct:
80     // https://github.com/KhronosGroup/SPIRV-Headers/issues/86
81     // However OpKill is fragment shader specific and we don't support it yet.
82     return true;
83   }
84 
85   /// Handle the given inlined terminator by replacing it with a new operation
86   /// as necessary.
87   void handleTerminator(Operation *op, Block *newDest) const final {
88     if (auto returnOp = dyn_cast<spirv::ReturnOp>(op)) {
89       OpBuilder(op).create<spirv::BranchOp>(op->getLoc(), newDest);
90       op->erase();
91     } else if (auto retValOp = dyn_cast<spirv::ReturnValueOp>(op)) {
92       llvm_unreachable("unimplemented spv.ReturnValue in inliner");
93     }
94   }
95 
96   /// Handle the given inlined terminator by replacing it with a new operation
97   /// as necessary.
98   void handleTerminator(Operation *op,
99                         ArrayRef<Value> valuesToRepl) const final {
100     // Only spv.ReturnValue needs to be handled here.
101     auto retValOp = dyn_cast<spirv::ReturnValueOp>(op);
102     if (!retValOp)
103       return;
104 
105     // Replace the values directly with the return operands.
106     assert(valuesToRepl.size() == 1 &&
107            "spv.ReturnValue expected to only handle one result");
108     valuesToRepl.front().replaceAllUsesWith(retValOp.value());
109   }
110 };
111 } // namespace
112 
113 //===----------------------------------------------------------------------===//
114 // SPIR-V Dialect
115 //===----------------------------------------------------------------------===//
116 
117 void SPIRVDialect::initialize() {
118   registerAttributes();
119   registerTypes();
120 
121   // Add SPIR-V ops.
122   addOperations<
123 #define GET_OP_LIST
124 #include "mlir/Dialect/SPIRV/IR/SPIRVOps.cpp.inc"
125       >();
126 
127   addInterfaces<SPIRVInlinerInterface>();
128 
129   // Allow unknown operations because SPIR-V is extensible.
130   allowUnknownOperations();
131 }
132 
133 std::string SPIRVDialect::getAttributeName(Decoration decoration) {
134   return llvm::convertToSnakeFromCamelCase(stringifyDecoration(decoration));
135 }
136 
137 //===----------------------------------------------------------------------===//
138 // Type Parsing
139 //===----------------------------------------------------------------------===//
140 
141 // Forward declarations.
142 template <typename ValTy>
143 static Optional<ValTy> parseAndVerify(SPIRVDialect const &dialect,
144                                       DialectAsmParser &parser);
145 template <>
146 Optional<Type> parseAndVerify<Type>(SPIRVDialect const &dialect,
147                                     DialectAsmParser &parser);
148 
149 template <>
150 Optional<unsigned> parseAndVerify<unsigned>(SPIRVDialect const &dialect,
151                                             DialectAsmParser &parser);
152 
153 static Type parseAndVerifyType(SPIRVDialect const &dialect,
154                                DialectAsmParser &parser) {
155   Type type;
156   llvm::SMLoc typeLoc = parser.getCurrentLocation();
157   if (parser.parseType(type))
158     return Type();
159 
160   // Allow SPIR-V dialect types
161   if (&type.getDialect() == &dialect)
162     return type;
163 
164   // Check other allowed types
165   if (auto t = type.dyn_cast<FloatType>()) {
166     if (type.isBF16()) {
167       parser.emitError(typeLoc, "cannot use 'bf16' to compose SPIR-V types");
168       return Type();
169     }
170   } else if (auto t = type.dyn_cast<IntegerType>()) {
171     if (!ScalarType::isValid(t)) {
172       parser.emitError(typeLoc,
173                        "only 1/8/16/32/64-bit integer type allowed but found ")
174           << type;
175       return Type();
176     }
177   } else if (auto t = type.dyn_cast<VectorType>()) {
178     if (t.getRank() != 1) {
179       parser.emitError(typeLoc, "only 1-D vector allowed but found ") << t;
180       return Type();
181     }
182     if (t.getNumElements() > 4) {
183       parser.emitError(
184           typeLoc, "vector length has to be less than or equal to 4 but found ")
185           << t.getNumElements();
186       return Type();
187     }
188   } else {
189     parser.emitError(typeLoc, "cannot use ")
190         << type << " to compose SPIR-V types";
191     return Type();
192   }
193 
194   return type;
195 }
196 
197 static Type parseAndVerifyMatrixType(SPIRVDialect const &dialect,
198                                      DialectAsmParser &parser) {
199   Type type;
200   llvm::SMLoc typeLoc = parser.getCurrentLocation();
201   if (parser.parseType(type))
202     return Type();
203 
204   if (auto t = type.dyn_cast<VectorType>()) {
205     if (t.getRank() != 1) {
206       parser.emitError(typeLoc, "only 1-D vector allowed but found ") << t;
207       return Type();
208     }
209     if (t.getNumElements() > 4 || t.getNumElements() < 2) {
210       parser.emitError(typeLoc,
211                        "matrix columns size has to be less than or equal "
212                        "to 4 and greater than or equal 2, but found ")
213           << t.getNumElements();
214       return Type();
215     }
216 
217     if (!t.getElementType().isa<FloatType>()) {
218       parser.emitError(typeLoc, "matrix columns' elements must be of "
219                                 "Float type, got ")
220           << t.getElementType();
221       return Type();
222     }
223   } else {
224     parser.emitError(typeLoc, "matrix must be composed using vector "
225                               "type, got ")
226         << type;
227     return Type();
228   }
229 
230   return type;
231 }
232 
233 static Type parseAndVerifySampledImageType(SPIRVDialect const &dialect,
234                                            DialectAsmParser &parser) {
235   Type type;
236   llvm::SMLoc typeLoc = parser.getCurrentLocation();
237   if (parser.parseType(type))
238     return Type();
239 
240   if (!type.isa<ImageType>()) {
241     parser.emitError(typeLoc,
242                      "sampled image must be composed using image type, got ")
243         << type;
244     return Type();
245   }
246 
247   return type;
248 }
249 
250 /// Parses an optional `, stride = N` assembly segment. If no parsing failure
251 /// occurs, writes `N` to `stride` if existing and writes 0 to `stride` if
252 /// missing.
253 static LogicalResult parseOptionalArrayStride(const SPIRVDialect &dialect,
254                                               DialectAsmParser &parser,
255                                               unsigned &stride) {
256   if (failed(parser.parseOptionalComma())) {
257     stride = 0;
258     return success();
259   }
260 
261   if (parser.parseKeyword("stride") || parser.parseEqual())
262     return failure();
263 
264   llvm::SMLoc strideLoc = parser.getCurrentLocation();
265   Optional<unsigned> optStride = parseAndVerify<unsigned>(dialect, parser);
266   if (!optStride)
267     return failure();
268 
269   if (!(stride = optStride.getValue())) {
270     parser.emitError(strideLoc, "ArrayStride must be greater than zero");
271     return failure();
272   }
273   return success();
274 }
275 
276 // element-type ::= integer-type
277 //                | floating-point-type
278 //                | vector-type
279 //                | spirv-type
280 //
281 // array-type ::= `!spv.array` `<` integer-literal `x` element-type
282 //                (`,` `stride` `=` integer-literal)? `>`
283 static Type parseArrayType(SPIRVDialect const &dialect,
284                            DialectAsmParser &parser) {
285   if (parser.parseLess())
286     return Type();
287 
288   SmallVector<int64_t, 1> countDims;
289   llvm::SMLoc countLoc = parser.getCurrentLocation();
290   if (parser.parseDimensionList(countDims, /*allowDynamic=*/false))
291     return Type();
292   if (countDims.size() != 1) {
293     parser.emitError(countLoc,
294                      "expected single integer for array element count");
295     return Type();
296   }
297 
298   // According to the SPIR-V spec:
299   // "Length is the number of elements in the array. It must be at least 1."
300   int64_t count = countDims[0];
301   if (count == 0) {
302     parser.emitError(countLoc, "expected array length greater than 0");
303     return Type();
304   }
305 
306   Type elementType = parseAndVerifyType(dialect, parser);
307   if (!elementType)
308     return Type();
309 
310   unsigned stride = 0;
311   if (failed(parseOptionalArrayStride(dialect, parser, stride)))
312     return Type();
313 
314   if (parser.parseGreater())
315     return Type();
316   return ArrayType::get(elementType, count, stride);
317 }
318 
319 // cooperative-matrix-type ::= `!spv.coopmatrix` `<` element-type ',' scope ','
320 //                                                   rows ',' columns>`
321 static Type parseCooperativeMatrixType(SPIRVDialect const &dialect,
322                                        DialectAsmParser &parser) {
323   if (parser.parseLess())
324     return Type();
325 
326   SmallVector<int64_t, 2> dims;
327   llvm::SMLoc countLoc = parser.getCurrentLocation();
328   if (parser.parseDimensionList(dims, /*allowDynamic=*/false))
329     return Type();
330 
331   if (dims.size() != 2) {
332     parser.emitError(countLoc, "expected rows and columns size");
333     return Type();
334   }
335 
336   auto elementTy = parseAndVerifyType(dialect, parser);
337   if (!elementTy)
338     return Type();
339 
340   Scope scope;
341   if (parser.parseComma() || parseEnumKeywordAttr(scope, parser, "scope <id>"))
342     return Type();
343 
344   if (parser.parseGreater())
345     return Type();
346   return CooperativeMatrixNVType::get(elementTy, scope, dims[0], dims[1]);
347 }
348 
349 // TODO: Reorder methods to be utilities first and parse*Type
350 // methods in alphabetical order
351 //
352 // storage-class ::= `UniformConstant`
353 //                 | `Uniform`
354 //                 | `Workgroup`
355 //                 | <and other storage classes...>
356 //
357 // pointer-type ::= `!spv.ptr<` element-type `,` storage-class `>`
358 static Type parsePointerType(SPIRVDialect const &dialect,
359                              DialectAsmParser &parser) {
360   if (parser.parseLess())
361     return Type();
362 
363   auto pointeeType = parseAndVerifyType(dialect, parser);
364   if (!pointeeType)
365     return Type();
366 
367   StringRef storageClassSpec;
368   llvm::SMLoc storageClassLoc = parser.getCurrentLocation();
369   if (parser.parseComma() || parser.parseKeyword(&storageClassSpec))
370     return Type();
371 
372   auto storageClass = symbolizeStorageClass(storageClassSpec);
373   if (!storageClass) {
374     parser.emitError(storageClassLoc, "unknown storage class: ")
375         << storageClassSpec;
376     return Type();
377   }
378   if (parser.parseGreater())
379     return Type();
380   return PointerType::get(pointeeType, *storageClass);
381 }
382 
383 // runtime-array-type ::= `!spv.rtarray` `<` element-type
384 //                        (`,` `stride` `=` integer-literal)? `>`
385 static Type parseRuntimeArrayType(SPIRVDialect const &dialect,
386                                   DialectAsmParser &parser) {
387   if (parser.parseLess())
388     return Type();
389 
390   Type elementType = parseAndVerifyType(dialect, parser);
391   if (!elementType)
392     return Type();
393 
394   unsigned stride = 0;
395   if (failed(parseOptionalArrayStride(dialect, parser, stride)))
396     return Type();
397 
398   if (parser.parseGreater())
399     return Type();
400   return RuntimeArrayType::get(elementType, stride);
401 }
402 
403 // matrix-type ::= `!spv.matrix` `<` integer-literal `x` element-type `>`
404 static Type parseMatrixType(SPIRVDialect const &dialect,
405                             DialectAsmParser &parser) {
406   if (parser.parseLess())
407     return Type();
408 
409   SmallVector<int64_t, 1> countDims;
410   llvm::SMLoc countLoc = parser.getCurrentLocation();
411   if (parser.parseDimensionList(countDims, /*allowDynamic=*/false))
412     return Type();
413   if (countDims.size() != 1) {
414     parser.emitError(countLoc, "expected single unsigned "
415                                "integer for number of columns");
416     return Type();
417   }
418 
419   int64_t columnCount = countDims[0];
420   // According to the specification, Matrices can have 2, 3, or 4 columns
421   if (columnCount < 2 || columnCount > 4) {
422     parser.emitError(countLoc, "matrix is expected to have 2, 3, or 4 "
423                                "columns");
424     return Type();
425   }
426 
427   Type columnType = parseAndVerifyMatrixType(dialect, parser);
428   if (!columnType)
429     return Type();
430 
431   if (parser.parseGreater())
432     return Type();
433 
434   return MatrixType::get(columnType, columnCount);
435 }
436 
437 // Specialize this function to parse each of the parameters that define an
438 // ImageType. By default it assumes this is an enum type.
439 template <typename ValTy>
440 static Optional<ValTy> parseAndVerify(SPIRVDialect const &dialect,
441                                       DialectAsmParser &parser) {
442   StringRef enumSpec;
443   llvm::SMLoc enumLoc = parser.getCurrentLocation();
444   if (parser.parseKeyword(&enumSpec)) {
445     return llvm::None;
446   }
447 
448   auto val = spirv::symbolizeEnum<ValTy>(enumSpec);
449   if (!val)
450     parser.emitError(enumLoc, "unknown attribute: '") << enumSpec << "'";
451   return val;
452 }
453 
454 template <>
455 Optional<Type> parseAndVerify<Type>(SPIRVDialect const &dialect,
456                                     DialectAsmParser &parser) {
457   // TODO: Further verify that the element type can be sampled
458   auto ty = parseAndVerifyType(dialect, parser);
459   if (!ty)
460     return llvm::None;
461   return ty;
462 }
463 
464 template <typename IntTy>
465 static Optional<IntTy> parseAndVerifyInteger(SPIRVDialect const &dialect,
466                                              DialectAsmParser &parser) {
467   IntTy offsetVal = std::numeric_limits<IntTy>::max();
468   if (parser.parseInteger(offsetVal))
469     return llvm::None;
470   return offsetVal;
471 }
472 
473 template <>
474 Optional<unsigned> parseAndVerify<unsigned>(SPIRVDialect const &dialect,
475                                             DialectAsmParser &parser) {
476   return parseAndVerifyInteger<unsigned>(dialect, parser);
477 }
478 
479 namespace {
480 // Functor object to parse a comma separated list of specs. The function
481 // parseAndVerify does the actual parsing and verification of individual
482 // elements. This is a functor since parsing the last element of the list
483 // (termination condition) needs partial specialization.
484 template <typename ParseType, typename... Args>
485 struct ParseCommaSeparatedList {
486   Optional<std::tuple<ParseType, Args...>>
487   operator()(SPIRVDialect const &dialect, DialectAsmParser &parser) const {
488     auto parseVal = parseAndVerify<ParseType>(dialect, parser);
489     if (!parseVal)
490       return llvm::None;
491 
492     auto numArgs = std::tuple_size<std::tuple<Args...>>::value;
493     if (numArgs != 0 && failed(parser.parseComma()))
494       return llvm::None;
495     auto remainingValues = ParseCommaSeparatedList<Args...>{}(dialect, parser);
496     if (!remainingValues)
497       return llvm::None;
498     return std::tuple_cat(std::tuple<ParseType>(parseVal.getValue()),
499                           remainingValues.getValue());
500   }
501 };
502 
503 // Partial specialization of the function to parse a comma separated list of
504 // specs to parse the last element of the list.
505 template <typename ParseType>
506 struct ParseCommaSeparatedList<ParseType> {
507   Optional<std::tuple<ParseType>> operator()(SPIRVDialect const &dialect,
508                                              DialectAsmParser &parser) const {
509     if (auto value = parseAndVerify<ParseType>(dialect, parser))
510       return std::tuple<ParseType>(value.getValue());
511     return llvm::None;
512   }
513 };
514 } // namespace
515 
516 // dim ::= `1D` | `2D` | `3D` | `Cube` | <and other SPIR-V Dim specifiers...>
517 //
518 // depth-info ::= `NoDepth` | `IsDepth` | `DepthUnknown`
519 //
520 // arrayed-info ::= `NonArrayed` | `Arrayed`
521 //
522 // sampling-info ::= `SingleSampled` | `MultiSampled`
523 //
524 // sampler-use-info ::= `SamplerUnknown` | `NeedSampler` |  `NoSampler`
525 //
526 // format ::= `Unknown` | `Rgba32f` | <and other SPIR-V Image formats...>
527 //
528 // image-type ::= `!spv.image<` element-type `,` dim `,` depth-info `,`
529 //                              arrayed-info `,` sampling-info `,`
530 //                              sampler-use-info `,` format `>`
531 static Type parseImageType(SPIRVDialect const &dialect,
532                            DialectAsmParser &parser) {
533   if (parser.parseLess())
534     return Type();
535 
536   auto value =
537       ParseCommaSeparatedList<Type, Dim, ImageDepthInfo, ImageArrayedInfo,
538                               ImageSamplingInfo, ImageSamplerUseInfo,
539                               ImageFormat>{}(dialect, parser);
540   if (!value)
541     return Type();
542 
543   if (parser.parseGreater())
544     return Type();
545   return ImageType::get(value.getValue());
546 }
547 
548 // sampledImage-type :: = `!spv.sampledImage<` image-type `>`
549 static Type parseSampledImageType(SPIRVDialect const &dialect,
550                                   DialectAsmParser &parser) {
551   if (parser.parseLess())
552     return Type();
553 
554   Type parsedType = parseAndVerifySampledImageType(dialect, parser);
555   if (!parsedType)
556     return Type();
557 
558   if (parser.parseGreater())
559     return Type();
560   return SampledImageType::get(parsedType);
561 }
562 
563 // Parse decorations associated with a member.
564 static ParseResult parseStructMemberDecorations(
565     SPIRVDialect const &dialect, DialectAsmParser &parser,
566     ArrayRef<Type> memberTypes,
567     SmallVectorImpl<StructType::OffsetInfo> &offsetInfo,
568     SmallVectorImpl<StructType::MemberDecorationInfo> &memberDecorationInfo) {
569 
570   // Check if the first element is offset.
571   llvm::SMLoc offsetLoc = parser.getCurrentLocation();
572   StructType::OffsetInfo offset = 0;
573   OptionalParseResult offsetParseResult = parser.parseOptionalInteger(offset);
574   if (offsetParseResult.hasValue()) {
575     if (failed(*offsetParseResult))
576       return failure();
577 
578     if (offsetInfo.size() != memberTypes.size() - 1) {
579       return parser.emitError(offsetLoc,
580                               "offset specification must be given for "
581                               "all members");
582     }
583     offsetInfo.push_back(offset);
584   }
585 
586   // Check for no spirv::Decorations.
587   if (succeeded(parser.parseOptionalRSquare()))
588     return success();
589 
590   // If there was an offset, make sure to parse the comma.
591   if (offsetParseResult.hasValue() && parser.parseComma())
592     return failure();
593 
594   // Check for spirv::Decorations.
595   do {
596     auto memberDecoration = parseAndVerify<spirv::Decoration>(dialect, parser);
597     if (!memberDecoration)
598       return failure();
599 
600     // Parse member decoration value if it exists.
601     if (succeeded(parser.parseOptionalEqual())) {
602       auto memberDecorationValue =
603           parseAndVerifyInteger<uint32_t>(dialect, parser);
604 
605       if (!memberDecorationValue)
606         return failure();
607 
608       memberDecorationInfo.emplace_back(
609           static_cast<uint32_t>(memberTypes.size() - 1), 1,
610           memberDecoration.getValue(), memberDecorationValue.getValue());
611     } else {
612       memberDecorationInfo.emplace_back(
613           static_cast<uint32_t>(memberTypes.size() - 1), 0,
614           memberDecoration.getValue(), 0);
615     }
616 
617   } while (succeeded(parser.parseOptionalComma()));
618 
619   return parser.parseRSquare();
620 }
621 
622 // struct-member-decoration ::= integer-literal? spirv-decoration*
623 // struct-type ::=
624 //             `!spv.struct<` (id `,`)?
625 //                          `(`
626 //                            (spirv-type (`[` struct-member-decoration `]`)?)*
627 //                          `)>`
628 static Type parseStructType(SPIRVDialect const &dialect,
629                             DialectAsmParser &parser) {
630   // TODO: This function is quite lengthy. Break it down into smaller chunks.
631 
632   // To properly resolve recursive references while parsing recursive struct
633   // types, we need to maintain a list of enclosing struct type names. This set
634   // maintains the names of struct types in which the type we are about to parse
635   // is nested.
636   //
637   // Note: This has to be thread_local to enable multiple threads to safely
638   // parse concurrently.
639   thread_local llvm::SetVector<StringRef> structContext;
640 
641   static auto removeIdentifierAndFail =
642       [](llvm::SetVector<StringRef> &structContext, StringRef identifier) {
643         if (!identifier.empty())
644           structContext.remove(identifier);
645 
646         return Type();
647       };
648 
649   if (parser.parseLess())
650     return Type();
651 
652   StringRef identifier;
653 
654   // Check if this is an identified struct type.
655   if (succeeded(parser.parseOptionalKeyword(&identifier))) {
656     // Check if this is a possible recursive reference.
657     if (succeeded(parser.parseOptionalGreater())) {
658       if (structContext.count(identifier) == 0) {
659         parser.emitError(
660             parser.getNameLoc(),
661             "recursive struct reference not nested in struct definition");
662 
663         return Type();
664       }
665 
666       return StructType::getIdentified(dialect.getContext(), identifier);
667     }
668 
669     if (failed(parser.parseComma()))
670       return Type();
671 
672     if (structContext.count(identifier) != 0) {
673       parser.emitError(parser.getNameLoc(),
674                        "identifier already used for an enclosing struct");
675 
676       return removeIdentifierAndFail(structContext, identifier);
677     }
678 
679     structContext.insert(identifier);
680   }
681 
682   if (failed(parser.parseLParen()))
683     return removeIdentifierAndFail(structContext, identifier);
684 
685   if (succeeded(parser.parseOptionalRParen()) &&
686       succeeded(parser.parseOptionalGreater())) {
687     if (!identifier.empty())
688       structContext.remove(identifier);
689 
690     return StructType::getEmpty(dialect.getContext(), identifier);
691   }
692 
693   StructType idStructTy;
694 
695   if (!identifier.empty())
696     idStructTy = StructType::getIdentified(dialect.getContext(), identifier);
697 
698   SmallVector<Type, 4> memberTypes;
699   SmallVector<StructType::OffsetInfo, 4> offsetInfo;
700   SmallVector<StructType::MemberDecorationInfo, 4> memberDecorationInfo;
701 
702   do {
703     Type memberType;
704     if (parser.parseType(memberType))
705       return removeIdentifierAndFail(structContext, identifier);
706     memberTypes.push_back(memberType);
707 
708     if (succeeded(parser.parseOptionalLSquare()))
709       if (parseStructMemberDecorations(dialect, parser, memberTypes, offsetInfo,
710                                        memberDecorationInfo))
711         return removeIdentifierAndFail(structContext, identifier);
712   } while (succeeded(parser.parseOptionalComma()));
713 
714   if (!offsetInfo.empty() && memberTypes.size() != offsetInfo.size()) {
715     parser.emitError(parser.getNameLoc(),
716                      "offset specification must be given for all members");
717     return removeIdentifierAndFail(structContext, identifier);
718   }
719 
720   if (failed(parser.parseRParen()) || failed(parser.parseGreater()))
721     return removeIdentifierAndFail(structContext, identifier);
722 
723   if (!identifier.empty()) {
724     if (failed(idStructTy.trySetBody(memberTypes, offsetInfo,
725                                      memberDecorationInfo)))
726       return Type();
727 
728     structContext.remove(identifier);
729     return idStructTy;
730   }
731 
732   return StructType::get(memberTypes, offsetInfo, memberDecorationInfo);
733 }
734 
735 // spirv-type ::= array-type
736 //              | element-type
737 //              | image-type
738 //              | pointer-type
739 //              | runtime-array-type
740 //              | sampled-image-type
741 //              | struct-type
742 Type SPIRVDialect::parseType(DialectAsmParser &parser) const {
743   StringRef keyword;
744   if (parser.parseKeyword(&keyword))
745     return Type();
746 
747   if (keyword == "array")
748     return parseArrayType(*this, parser);
749   if (keyword == "coopmatrix")
750     return parseCooperativeMatrixType(*this, parser);
751   if (keyword == "image")
752     return parseImageType(*this, parser);
753   if (keyword == "ptr")
754     return parsePointerType(*this, parser);
755   if (keyword == "rtarray")
756     return parseRuntimeArrayType(*this, parser);
757   if (keyword == "sampled_image")
758     return parseSampledImageType(*this, parser);
759   if (keyword == "struct")
760     return parseStructType(*this, parser);
761   if (keyword == "matrix")
762     return parseMatrixType(*this, parser);
763   parser.emitError(parser.getNameLoc(), "unknown SPIR-V type: ") << keyword;
764   return Type();
765 }
766 
767 //===----------------------------------------------------------------------===//
768 // Type Printing
769 //===----------------------------------------------------------------------===//
770 
771 static void print(ArrayType type, DialectAsmPrinter &os) {
772   os << "array<" << type.getNumElements() << " x " << type.getElementType();
773   if (unsigned stride = type.getArrayStride())
774     os << ", stride=" << stride;
775   os << ">";
776 }
777 
778 static void print(RuntimeArrayType type, DialectAsmPrinter &os) {
779   os << "rtarray<" << type.getElementType();
780   if (unsigned stride = type.getArrayStride())
781     os << ", stride=" << stride;
782   os << ">";
783 }
784 
785 static void print(PointerType type, DialectAsmPrinter &os) {
786   os << "ptr<" << type.getPointeeType() << ", "
787      << stringifyStorageClass(type.getStorageClass()) << ">";
788 }
789 
790 static void print(ImageType type, DialectAsmPrinter &os) {
791   os << "image<" << type.getElementType() << ", " << stringifyDim(type.getDim())
792      << ", " << stringifyImageDepthInfo(type.getDepthInfo()) << ", "
793      << stringifyImageArrayedInfo(type.getArrayedInfo()) << ", "
794      << stringifyImageSamplingInfo(type.getSamplingInfo()) << ", "
795      << stringifyImageSamplerUseInfo(type.getSamplerUseInfo()) << ", "
796      << stringifyImageFormat(type.getImageFormat()) << ">";
797 }
798 
799 static void print(SampledImageType type, DialectAsmPrinter &os) {
800   os << "sampled_image<" << type.getImageType() << ">";
801 }
802 
803 static void print(StructType type, DialectAsmPrinter &os) {
804   thread_local llvm::SetVector<StringRef> structContext;
805 
806   os << "struct<";
807 
808   if (type.isIdentified()) {
809     os << type.getIdentifier();
810 
811     if (structContext.count(type.getIdentifier())) {
812       os << ">";
813       return;
814     }
815 
816     os << ", ";
817     structContext.insert(type.getIdentifier());
818   }
819 
820   os << "(";
821 
822   auto printMember = [&](unsigned i) {
823     os << type.getElementType(i);
824     SmallVector<spirv::StructType::MemberDecorationInfo, 0> decorations;
825     type.getMemberDecorations(i, decorations);
826     if (type.hasOffset() || !decorations.empty()) {
827       os << " [";
828       if (type.hasOffset()) {
829         os << type.getMemberOffset(i);
830         if (!decorations.empty())
831           os << ", ";
832       }
833       auto eachFn = [&os](spirv::StructType::MemberDecorationInfo decoration) {
834         os << stringifyDecoration(decoration.decoration);
835         if (decoration.hasValue) {
836           os << "=" << decoration.decorationValue;
837         }
838       };
839       llvm::interleaveComma(decorations, os, eachFn);
840       os << "]";
841     }
842   };
843   llvm::interleaveComma(llvm::seq<unsigned>(0, type.getNumElements()), os,
844                         printMember);
845   os << ")>";
846 
847   if (type.isIdentified())
848     structContext.remove(type.getIdentifier());
849 }
850 
851 static void print(CooperativeMatrixNVType type, DialectAsmPrinter &os) {
852   os << "coopmatrix<" << type.getRows() << "x" << type.getColumns() << "x";
853   os << type.getElementType() << ", " << stringifyScope(type.getScope());
854   os << ">";
855 }
856 
857 static void print(MatrixType type, DialectAsmPrinter &os) {
858   os << "matrix<" << type.getNumColumns() << " x " << type.getColumnType();
859   os << ">";
860 }
861 
862 void SPIRVDialect::printType(Type type, DialectAsmPrinter &os) const {
863   TypeSwitch<Type>(type)
864       .Case<ArrayType, CooperativeMatrixNVType, PointerType, RuntimeArrayType,
865             ImageType, SampledImageType, StructType, MatrixType>(
866           [&](auto type) { print(type, os); })
867       .Default([](Type) { llvm_unreachable("unhandled SPIR-V type"); });
868 }
869 
870 //===----------------------------------------------------------------------===//
871 // Attribute Parsing
872 //===----------------------------------------------------------------------===//
873 
874 /// Parses a comma-separated list of keywords, invokes `processKeyword` on each
875 /// of the parsed keyword, and returns failure if any error occurs.
876 static ParseResult parseKeywordList(
877     DialectAsmParser &parser,
878     function_ref<LogicalResult(llvm::SMLoc, StringRef)> processKeyword) {
879   if (parser.parseLSquare())
880     return failure();
881 
882   // Special case for empty list.
883   if (succeeded(parser.parseOptionalRSquare()))
884     return success();
885 
886   // Keep parsing the keyword and an optional comma following it. If the comma
887   // is successfully parsed, then we have more keywords to parse.
888   do {
889     auto loc = parser.getCurrentLocation();
890     StringRef keyword;
891     if (parser.parseKeyword(&keyword) || failed(processKeyword(loc, keyword)))
892       return failure();
893   } while (succeeded(parser.parseOptionalComma()));
894 
895   if (parser.parseRSquare())
896     return failure();
897 
898   return success();
899 }
900 
901 /// Parses a spirv::InterfaceVarABIAttr.
902 static Attribute parseInterfaceVarABIAttr(DialectAsmParser &parser) {
903   if (parser.parseLess())
904     return {};
905 
906   Builder &builder = parser.getBuilder();
907 
908   if (parser.parseLParen())
909     return {};
910 
911   IntegerAttr descriptorSetAttr;
912   {
913     auto loc = parser.getCurrentLocation();
914     uint32_t descriptorSet = 0;
915     auto descriptorSetParseResult = parser.parseOptionalInteger(descriptorSet);
916 
917     if (!descriptorSetParseResult.hasValue() ||
918         failed(*descriptorSetParseResult)) {
919       parser.emitError(loc, "missing descriptor set");
920       return {};
921     }
922     descriptorSetAttr = builder.getI32IntegerAttr(descriptorSet);
923   }
924 
925   if (parser.parseComma())
926     return {};
927 
928   IntegerAttr bindingAttr;
929   {
930     auto loc = parser.getCurrentLocation();
931     uint32_t binding = 0;
932     auto bindingParseResult = parser.parseOptionalInteger(binding);
933 
934     if (!bindingParseResult.hasValue() || failed(*bindingParseResult)) {
935       parser.emitError(loc, "missing binding");
936       return {};
937     }
938     bindingAttr = builder.getI32IntegerAttr(binding);
939   }
940 
941   if (parser.parseRParen())
942     return {};
943 
944   IntegerAttr storageClassAttr;
945   {
946     if (succeeded(parser.parseOptionalComma())) {
947       auto loc = parser.getCurrentLocation();
948       StringRef storageClass;
949       if (parser.parseKeyword(&storageClass))
950         return {};
951 
952       if (auto storageClassSymbol =
953               spirv::symbolizeStorageClass(storageClass)) {
954         storageClassAttr = builder.getI32IntegerAttr(
955             static_cast<uint32_t>(*storageClassSymbol));
956       } else {
957         parser.emitError(loc, "unknown storage class: ") << storageClass;
958         return {};
959       }
960     }
961   }
962 
963   if (parser.parseGreater())
964     return {};
965 
966   return spirv::InterfaceVarABIAttr::get(descriptorSetAttr, bindingAttr,
967                                          storageClassAttr);
968 }
969 
970 static Attribute parseVerCapExtAttr(DialectAsmParser &parser) {
971   if (parser.parseLess())
972     return {};
973 
974   Builder &builder = parser.getBuilder();
975 
976   IntegerAttr versionAttr;
977   {
978     auto loc = parser.getCurrentLocation();
979     StringRef version;
980     if (parser.parseKeyword(&version) || parser.parseComma())
981       return {};
982 
983     if (auto versionSymbol = spirv::symbolizeVersion(version)) {
984       versionAttr =
985           builder.getI32IntegerAttr(static_cast<uint32_t>(*versionSymbol));
986     } else {
987       parser.emitError(loc, "unknown version: ") << version;
988       return {};
989     }
990   }
991 
992   ArrayAttr capabilitiesAttr;
993   {
994     SmallVector<Attribute, 4> capabilities;
995     llvm::SMLoc errorloc;
996     StringRef errorKeyword;
997 
998     auto processCapability = [&](llvm::SMLoc loc, StringRef capability) {
999       if (auto capSymbol = spirv::symbolizeCapability(capability)) {
1000         capabilities.push_back(
1001             builder.getI32IntegerAttr(static_cast<uint32_t>(*capSymbol)));
1002         return success();
1003       }
1004       return errorloc = loc, errorKeyword = capability, failure();
1005     };
1006     if (parseKeywordList(parser, processCapability) || parser.parseComma()) {
1007       if (!errorKeyword.empty())
1008         parser.emitError(errorloc, "unknown capability: ") << errorKeyword;
1009       return {};
1010     }
1011 
1012     capabilitiesAttr = builder.getArrayAttr(capabilities);
1013   }
1014 
1015   ArrayAttr extensionsAttr;
1016   {
1017     SmallVector<Attribute, 1> extensions;
1018     llvm::SMLoc errorloc;
1019     StringRef errorKeyword;
1020 
1021     auto processExtension = [&](llvm::SMLoc loc, StringRef extension) {
1022       if (spirv::symbolizeExtension(extension)) {
1023         extensions.push_back(builder.getStringAttr(extension));
1024         return success();
1025       }
1026       return errorloc = loc, errorKeyword = extension, failure();
1027     };
1028     if (parseKeywordList(parser, processExtension)) {
1029       if (!errorKeyword.empty())
1030         parser.emitError(errorloc, "unknown extension: ") << errorKeyword;
1031       return {};
1032     }
1033 
1034     extensionsAttr = builder.getArrayAttr(extensions);
1035   }
1036 
1037   if (parser.parseGreater())
1038     return {};
1039 
1040   return spirv::VerCapExtAttr::get(versionAttr, capabilitiesAttr,
1041                                    extensionsAttr);
1042 }
1043 
1044 /// Parses a spirv::TargetEnvAttr.
1045 static Attribute parseTargetEnvAttr(DialectAsmParser &parser) {
1046   if (parser.parseLess())
1047     return {};
1048 
1049   spirv::VerCapExtAttr tripleAttr;
1050   if (parser.parseAttribute(tripleAttr) || parser.parseComma())
1051     return {};
1052 
1053   // Parse [vendor[:device-type[:device-id]]]
1054   Vendor vendorID = Vendor::Unknown;
1055   DeviceType deviceType = DeviceType::Unknown;
1056   uint32_t deviceID = spirv::TargetEnvAttr::kUnknownDeviceID;
1057   {
1058     auto loc = parser.getCurrentLocation();
1059     StringRef vendorStr;
1060     if (succeeded(parser.parseOptionalKeyword(&vendorStr))) {
1061       if (auto vendorSymbol = spirv::symbolizeVendor(vendorStr)) {
1062         vendorID = *vendorSymbol;
1063       } else {
1064         parser.emitError(loc, "unknown vendor: ") << vendorStr;
1065       }
1066 
1067       if (succeeded(parser.parseOptionalColon())) {
1068         loc = parser.getCurrentLocation();
1069         StringRef deviceTypeStr;
1070         if (parser.parseKeyword(&deviceTypeStr))
1071           return {};
1072         if (auto deviceTypeSymbol = spirv::symbolizeDeviceType(deviceTypeStr)) {
1073           deviceType = *deviceTypeSymbol;
1074         } else {
1075           parser.emitError(loc, "unknown device type: ") << deviceTypeStr;
1076         }
1077 
1078         if (succeeded(parser.parseOptionalColon())) {
1079           loc = parser.getCurrentLocation();
1080           if (parser.parseInteger(deviceID))
1081             return {};
1082         }
1083       }
1084       if (parser.parseComma())
1085         return {};
1086     }
1087   }
1088 
1089   DictionaryAttr limitsAttr;
1090   {
1091     auto loc = parser.getCurrentLocation();
1092     if (parser.parseAttribute(limitsAttr))
1093       return {};
1094 
1095     if (!limitsAttr.isa<spirv::ResourceLimitsAttr>()) {
1096       parser.emitError(
1097           loc,
1098           "limits must be a dictionary attribute containing two 32-bit integer "
1099           "attributes 'max_compute_workgroup_invocations' and "
1100           "'max_compute_workgroup_size'");
1101       return {};
1102     }
1103   }
1104 
1105   if (parser.parseGreater())
1106     return {};
1107 
1108   return spirv::TargetEnvAttr::get(tripleAttr, vendorID, deviceType, deviceID,
1109                                    limitsAttr);
1110 }
1111 
1112 Attribute SPIRVDialect::parseAttribute(DialectAsmParser &parser,
1113                                        Type type) const {
1114   // SPIR-V attributes are dictionaries so they do not have type.
1115   if (type) {
1116     parser.emitError(parser.getNameLoc(), "unexpected type");
1117     return {};
1118   }
1119 
1120   // Parse the kind keyword first.
1121   StringRef attrKind;
1122   if (parser.parseKeyword(&attrKind))
1123     return {};
1124 
1125   if (attrKind == spirv::TargetEnvAttr::getKindName())
1126     return parseTargetEnvAttr(parser);
1127   if (attrKind == spirv::VerCapExtAttr::getKindName())
1128     return parseVerCapExtAttr(parser);
1129   if (attrKind == spirv::InterfaceVarABIAttr::getKindName())
1130     return parseInterfaceVarABIAttr(parser);
1131 
1132   parser.emitError(parser.getNameLoc(), "unknown SPIR-V attribute kind: ")
1133       << attrKind;
1134   return {};
1135 }
1136 
1137 //===----------------------------------------------------------------------===//
1138 // Attribute Printing
1139 //===----------------------------------------------------------------------===//
1140 
1141 static void print(spirv::VerCapExtAttr triple, DialectAsmPrinter &printer) {
1142   auto &os = printer.getStream();
1143   printer << spirv::VerCapExtAttr::getKindName() << "<"
1144           << spirv::stringifyVersion(triple.getVersion()) << ", [";
1145   llvm::interleaveComma(
1146       triple.getCapabilities(), os,
1147       [&](spirv::Capability cap) { os << spirv::stringifyCapability(cap); });
1148   printer << "], [";
1149   llvm::interleaveComma(triple.getExtensionsAttr(), os, [&](Attribute attr) {
1150     os << attr.cast<StringAttr>().getValue();
1151   });
1152   printer << "]>";
1153 }
1154 
1155 static void print(spirv::TargetEnvAttr targetEnv, DialectAsmPrinter &printer) {
1156   printer << spirv::TargetEnvAttr::getKindName() << "<#spv.";
1157   print(targetEnv.getTripleAttr(), printer);
1158   spirv::Vendor vendorID = targetEnv.getVendorID();
1159   spirv::DeviceType deviceType = targetEnv.getDeviceType();
1160   uint32_t deviceID = targetEnv.getDeviceID();
1161   if (vendorID != spirv::Vendor::Unknown) {
1162     printer << ", " << spirv::stringifyVendor(vendorID);
1163     if (deviceType != spirv::DeviceType::Unknown) {
1164       printer << ":" << spirv::stringifyDeviceType(deviceType);
1165       if (deviceID != spirv::TargetEnvAttr::kUnknownDeviceID)
1166         printer << ":" << deviceID;
1167     }
1168   }
1169   printer << ", " << targetEnv.getResourceLimits() << ">";
1170 }
1171 
1172 static void print(spirv::InterfaceVarABIAttr interfaceVarABIAttr,
1173                   DialectAsmPrinter &printer) {
1174   printer << spirv::InterfaceVarABIAttr::getKindName() << "<("
1175           << interfaceVarABIAttr.getDescriptorSet() << ", "
1176           << interfaceVarABIAttr.getBinding() << ")";
1177   auto storageClass = interfaceVarABIAttr.getStorageClass();
1178   if (storageClass)
1179     printer << ", " << spirv::stringifyStorageClass(*storageClass);
1180   printer << ">";
1181 }
1182 
1183 void SPIRVDialect::printAttribute(Attribute attr,
1184                                   DialectAsmPrinter &printer) const {
1185   if (auto targetEnv = attr.dyn_cast<TargetEnvAttr>())
1186     print(targetEnv, printer);
1187   else if (auto vceAttr = attr.dyn_cast<VerCapExtAttr>())
1188     print(vceAttr, printer);
1189   else if (auto interfaceVarABIAttr = attr.dyn_cast<InterfaceVarABIAttr>())
1190     print(interfaceVarABIAttr, printer);
1191   else
1192     llvm_unreachable("unhandled SPIR-V attribute kind");
1193 }
1194 
1195 //===----------------------------------------------------------------------===//
1196 // Constant
1197 //===----------------------------------------------------------------------===//
1198 
1199 Operation *SPIRVDialect::materializeConstant(OpBuilder &builder,
1200                                              Attribute value, Type type,
1201                                              Location loc) {
1202   if (!spirv::ConstantOp::isBuildableWith(type))
1203     return nullptr;
1204 
1205   return builder.create<spirv::ConstantOp>(loc, type, value);
1206 }
1207 
1208 //===----------------------------------------------------------------------===//
1209 // Shader Interface ABI
1210 //===----------------------------------------------------------------------===//
1211 
1212 LogicalResult SPIRVDialect::verifyOperationAttribute(Operation *op,
1213                                                      NamedAttribute attribute) {
1214   StringRef symbol = attribute.first.strref();
1215   Attribute attr = attribute.second;
1216 
1217   // TODO: figure out a way to generate the description from the
1218   // StructAttr definition.
1219   if (symbol == spirv::getEntryPointABIAttrName()) {
1220     if (!attr.isa<spirv::EntryPointABIAttr>())
1221       return op->emitError("'")
1222              << symbol
1223              << "' attribute must be a dictionary attribute containing one "
1224                 "32-bit integer elements attribute: 'local_size'";
1225   } else if (symbol == spirv::getTargetEnvAttrName()) {
1226     if (!attr.isa<spirv::TargetEnvAttr>())
1227       return op->emitError("'") << symbol << "' must be a spirv::TargetEnvAttr";
1228   } else {
1229     return op->emitError("found unsupported '")
1230            << symbol << "' attribute on operation";
1231   }
1232 
1233   return success();
1234 }
1235 
1236 /// Verifies the given SPIR-V `attribute` attached to a value of the given
1237 /// `valueType` is valid.
1238 static LogicalResult verifyRegionAttribute(Location loc, Type valueType,
1239                                            NamedAttribute attribute) {
1240   StringRef symbol = attribute.first.strref();
1241   Attribute attr = attribute.second;
1242 
1243   if (symbol != spirv::getInterfaceVarABIAttrName())
1244     return emitError(loc, "found unsupported '")
1245            << symbol << "' attribute on region argument";
1246 
1247   auto varABIAttr = attr.dyn_cast<spirv::InterfaceVarABIAttr>();
1248   if (!varABIAttr)
1249     return emitError(loc, "'")
1250            << symbol << "' must be a spirv::InterfaceVarABIAttr";
1251 
1252   if (varABIAttr.getStorageClass() && !valueType.isIntOrIndexOrFloat())
1253     return emitError(loc, "'") << symbol
1254                                << "' attribute cannot specify storage class "
1255                                   "when attaching to a non-scalar value";
1256 
1257   return success();
1258 }
1259 
1260 LogicalResult SPIRVDialect::verifyRegionArgAttribute(Operation *op,
1261                                                      unsigned regionIndex,
1262                                                      unsigned argIndex,
1263                                                      NamedAttribute attribute) {
1264   return verifyRegionAttribute(
1265       op->getLoc(), op->getRegion(regionIndex).getArgument(argIndex).getType(),
1266       attribute);
1267 }
1268 
1269 LogicalResult SPIRVDialect::verifyRegionResultAttribute(
1270     Operation *op, unsigned /*regionIndex*/, unsigned /*resultIndex*/,
1271     NamedAttribute attribute) {
1272   return op->emitError("cannot attach SPIR-V attributes to region result");
1273 }
1274