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> struct ParseCommaSeparatedList {
485   Optional<std::tuple<ParseType, Args...>>
486   operator()(SPIRVDialect const &dialect, DialectAsmParser &parser) const {
487     auto parseVal = parseAndVerify<ParseType>(dialect, parser);
488     if (!parseVal)
489       return llvm::None;
490 
491     auto numArgs = std::tuple_size<std::tuple<Args...>>::value;
492     if (numArgs != 0 && failed(parser.parseComma()))
493       return llvm::None;
494     auto remainingValues = ParseCommaSeparatedList<Args...>{}(dialect, parser);
495     if (!remainingValues)
496       return llvm::None;
497     return std::tuple_cat(std::tuple<ParseType>(parseVal.getValue()),
498                           remainingValues.getValue());
499   }
500 };
501 
502 // Partial specialization of the function to parse a comma separated list of
503 // specs to parse the last element of the list.
504 template <typename ParseType> struct ParseCommaSeparatedList<ParseType> {
505   Optional<std::tuple<ParseType>> operator()(SPIRVDialect const &dialect,
506                                              DialectAsmParser &parser) const {
507     if (auto value = parseAndVerify<ParseType>(dialect, parser))
508       return std::tuple<ParseType>(value.getValue());
509     return llvm::None;
510   }
511 };
512 } // namespace
513 
514 // dim ::= `1D` | `2D` | `3D` | `Cube` | <and other SPIR-V Dim specifiers...>
515 //
516 // depth-info ::= `NoDepth` | `IsDepth` | `DepthUnknown`
517 //
518 // arrayed-info ::= `NonArrayed` | `Arrayed`
519 //
520 // sampling-info ::= `SingleSampled` | `MultiSampled`
521 //
522 // sampler-use-info ::= `SamplerUnknown` | `NeedSampler` |  `NoSampler`
523 //
524 // format ::= `Unknown` | `Rgba32f` | <and other SPIR-V Image formats...>
525 //
526 // image-type ::= `!spv.image<` element-type `,` dim `,` depth-info `,`
527 //                              arrayed-info `,` sampling-info `,`
528 //                              sampler-use-info `,` format `>`
529 static Type parseImageType(SPIRVDialect const &dialect,
530                            DialectAsmParser &parser) {
531   if (parser.parseLess())
532     return Type();
533 
534   auto value =
535       ParseCommaSeparatedList<Type, Dim, ImageDepthInfo, ImageArrayedInfo,
536                               ImageSamplingInfo, ImageSamplerUseInfo,
537                               ImageFormat>{}(dialect, parser);
538   if (!value)
539     return Type();
540 
541   if (parser.parseGreater())
542     return Type();
543   return ImageType::get(value.getValue());
544 }
545 
546 // sampledImage-type :: = `!spv.sampledImage<` image-type `>`
547 static Type parseSampledImageType(SPIRVDialect const &dialect,
548                                   DialectAsmParser &parser) {
549   if (parser.parseLess())
550     return Type();
551 
552   Type parsedType = parseAndVerifySampledImageType(dialect, parser);
553   if (!parsedType)
554     return Type();
555 
556   if (parser.parseGreater())
557     return Type();
558   return SampledImageType::get(parsedType);
559 }
560 
561 // Parse decorations associated with a member.
562 static ParseResult parseStructMemberDecorations(
563     SPIRVDialect const &dialect, DialectAsmParser &parser,
564     ArrayRef<Type> memberTypes,
565     SmallVectorImpl<StructType::OffsetInfo> &offsetInfo,
566     SmallVectorImpl<StructType::MemberDecorationInfo> &memberDecorationInfo) {
567 
568   // Check if the first element is offset.
569   llvm::SMLoc offsetLoc = parser.getCurrentLocation();
570   StructType::OffsetInfo offset = 0;
571   OptionalParseResult offsetParseResult = parser.parseOptionalInteger(offset);
572   if (offsetParseResult.hasValue()) {
573     if (failed(*offsetParseResult))
574       return failure();
575 
576     if (offsetInfo.size() != memberTypes.size() - 1) {
577       return parser.emitError(offsetLoc,
578                               "offset specification must be given for "
579                               "all members");
580     }
581     offsetInfo.push_back(offset);
582   }
583 
584   // Check for no spirv::Decorations.
585   if (succeeded(parser.parseOptionalRSquare()))
586     return success();
587 
588   // If there was an offset, make sure to parse the comma.
589   if (offsetParseResult.hasValue() && parser.parseComma())
590     return failure();
591 
592   // Check for spirv::Decorations.
593   do {
594     auto memberDecoration = parseAndVerify<spirv::Decoration>(dialect, parser);
595     if (!memberDecoration)
596       return failure();
597 
598     // Parse member decoration value if it exists.
599     if (succeeded(parser.parseOptionalEqual())) {
600       auto memberDecorationValue =
601           parseAndVerifyInteger<uint32_t>(dialect, parser);
602 
603       if (!memberDecorationValue)
604         return failure();
605 
606       memberDecorationInfo.emplace_back(
607           static_cast<uint32_t>(memberTypes.size() - 1), 1,
608           memberDecoration.getValue(), memberDecorationValue.getValue());
609     } else {
610       memberDecorationInfo.emplace_back(
611           static_cast<uint32_t>(memberTypes.size() - 1), 0,
612           memberDecoration.getValue(), 0);
613     }
614 
615   } while (succeeded(parser.parseOptionalComma()));
616 
617   return parser.parseRSquare();
618 }
619 
620 // struct-member-decoration ::= integer-literal? spirv-decoration*
621 // struct-type ::=
622 //             `!spv.struct<` (id `,`)?
623 //                          `(`
624 //                            (spirv-type (`[` struct-member-decoration `]`)?)*
625 //                          `)>`
626 static Type parseStructType(SPIRVDialect const &dialect,
627                             DialectAsmParser &parser) {
628   // TODO: This function is quite lengthy. Break it down into smaller chunks.
629 
630   // To properly resolve recursive references while parsing recursive struct
631   // types, we need to maintain a list of enclosing struct type names. This set
632   // maintains the names of struct types in which the type we are about to parse
633   // is nested.
634   //
635   // Note: This has to be thread_local to enable multiple threads to safely
636   // parse concurrently.
637   thread_local SetVector<StringRef> structContext;
638 
639   static auto removeIdentifierAndFail = [](SetVector<StringRef> &structContext,
640                                            StringRef identifier) {
641     if (!identifier.empty())
642       structContext.remove(identifier);
643 
644     return Type();
645   };
646 
647   if (parser.parseLess())
648     return Type();
649 
650   StringRef identifier;
651 
652   // Check if this is an identified struct type.
653   if (succeeded(parser.parseOptionalKeyword(&identifier))) {
654     // Check if this is a possible recursive reference.
655     if (succeeded(parser.parseOptionalGreater())) {
656       if (structContext.count(identifier) == 0) {
657         parser.emitError(
658             parser.getNameLoc(),
659             "recursive struct reference not nested in struct definition");
660 
661         return Type();
662       }
663 
664       return StructType::getIdentified(dialect.getContext(), identifier);
665     }
666 
667     if (failed(parser.parseComma()))
668       return Type();
669 
670     if (structContext.count(identifier) != 0) {
671       parser.emitError(parser.getNameLoc(),
672                        "identifier already used for an enclosing struct");
673 
674       return removeIdentifierAndFail(structContext, identifier);
675     }
676 
677     structContext.insert(identifier);
678   }
679 
680   if (failed(parser.parseLParen()))
681     return removeIdentifierAndFail(structContext, identifier);
682 
683   if (succeeded(parser.parseOptionalRParen()) &&
684       succeeded(parser.parseOptionalGreater())) {
685     if (!identifier.empty())
686       structContext.remove(identifier);
687 
688     return StructType::getEmpty(dialect.getContext(), identifier);
689   }
690 
691   StructType idStructTy;
692 
693   if (!identifier.empty())
694     idStructTy = StructType::getIdentified(dialect.getContext(), identifier);
695 
696   SmallVector<Type, 4> memberTypes;
697   SmallVector<StructType::OffsetInfo, 4> offsetInfo;
698   SmallVector<StructType::MemberDecorationInfo, 4> memberDecorationInfo;
699 
700   do {
701     Type memberType;
702     if (parser.parseType(memberType))
703       return removeIdentifierAndFail(structContext, identifier);
704     memberTypes.push_back(memberType);
705 
706     if (succeeded(parser.parseOptionalLSquare()))
707       if (parseStructMemberDecorations(dialect, parser, memberTypes, offsetInfo,
708                                        memberDecorationInfo))
709         return removeIdentifierAndFail(structContext, identifier);
710   } while (succeeded(parser.parseOptionalComma()));
711 
712   if (!offsetInfo.empty() && memberTypes.size() != offsetInfo.size()) {
713     parser.emitError(parser.getNameLoc(),
714                      "offset specification must be given for all members");
715     return removeIdentifierAndFail(structContext, identifier);
716   }
717 
718   if (failed(parser.parseRParen()) || failed(parser.parseGreater()))
719     return removeIdentifierAndFail(structContext, identifier);
720 
721   if (!identifier.empty()) {
722     if (failed(idStructTy.trySetBody(memberTypes, offsetInfo,
723                                      memberDecorationInfo)))
724       return Type();
725 
726     structContext.remove(identifier);
727     return idStructTy;
728   }
729 
730   return StructType::get(memberTypes, offsetInfo, memberDecorationInfo);
731 }
732 
733 // spirv-type ::= array-type
734 //              | element-type
735 //              | image-type
736 //              | pointer-type
737 //              | runtime-array-type
738 //              | sampled-image-type
739 //              | struct-type
740 Type SPIRVDialect::parseType(DialectAsmParser &parser) const {
741   StringRef keyword;
742   if (parser.parseKeyword(&keyword))
743     return Type();
744 
745   if (keyword == "array")
746     return parseArrayType(*this, parser);
747   if (keyword == "coopmatrix")
748     return parseCooperativeMatrixType(*this, parser);
749   if (keyword == "image")
750     return parseImageType(*this, parser);
751   if (keyword == "ptr")
752     return parsePointerType(*this, parser);
753   if (keyword == "rtarray")
754     return parseRuntimeArrayType(*this, parser);
755   if (keyword == "sampled_image")
756     return parseSampledImageType(*this, parser);
757   if (keyword == "struct")
758     return parseStructType(*this, parser);
759   if (keyword == "matrix")
760     return parseMatrixType(*this, parser);
761   parser.emitError(parser.getNameLoc(), "unknown SPIR-V type: ") << keyword;
762   return Type();
763 }
764 
765 //===----------------------------------------------------------------------===//
766 // Type Printing
767 //===----------------------------------------------------------------------===//
768 
769 static void print(ArrayType type, DialectAsmPrinter &os) {
770   os << "array<" << type.getNumElements() << " x " << type.getElementType();
771   if (unsigned stride = type.getArrayStride())
772     os << ", stride=" << stride;
773   os << ">";
774 }
775 
776 static void print(RuntimeArrayType type, DialectAsmPrinter &os) {
777   os << "rtarray<" << type.getElementType();
778   if (unsigned stride = type.getArrayStride())
779     os << ", stride=" << stride;
780   os << ">";
781 }
782 
783 static void print(PointerType type, DialectAsmPrinter &os) {
784   os << "ptr<" << type.getPointeeType() << ", "
785      << stringifyStorageClass(type.getStorageClass()) << ">";
786 }
787 
788 static void print(ImageType type, DialectAsmPrinter &os) {
789   os << "image<" << type.getElementType() << ", " << stringifyDim(type.getDim())
790      << ", " << stringifyImageDepthInfo(type.getDepthInfo()) << ", "
791      << stringifyImageArrayedInfo(type.getArrayedInfo()) << ", "
792      << stringifyImageSamplingInfo(type.getSamplingInfo()) << ", "
793      << stringifyImageSamplerUseInfo(type.getSamplerUseInfo()) << ", "
794      << stringifyImageFormat(type.getImageFormat()) << ">";
795 }
796 
797 static void print(SampledImageType type, DialectAsmPrinter &os) {
798   os << "sampled_image<" << type.getImageType() << ">";
799 }
800 
801 static void print(StructType type, DialectAsmPrinter &os) {
802   thread_local SetVector<StringRef> structContext;
803 
804   os << "struct<";
805 
806   if (type.isIdentified()) {
807     os << type.getIdentifier();
808 
809     if (structContext.count(type.getIdentifier())) {
810       os << ">";
811       return;
812     }
813 
814     os << ", ";
815     structContext.insert(type.getIdentifier());
816   }
817 
818   os << "(";
819 
820   auto printMember = [&](unsigned i) {
821     os << type.getElementType(i);
822     SmallVector<spirv::StructType::MemberDecorationInfo, 0> decorations;
823     type.getMemberDecorations(i, decorations);
824     if (type.hasOffset() || !decorations.empty()) {
825       os << " [";
826       if (type.hasOffset()) {
827         os << type.getMemberOffset(i);
828         if (!decorations.empty())
829           os << ", ";
830       }
831       auto eachFn = [&os](spirv::StructType::MemberDecorationInfo decoration) {
832         os << stringifyDecoration(decoration.decoration);
833         if (decoration.hasValue) {
834           os << "=" << decoration.decorationValue;
835         }
836       };
837       llvm::interleaveComma(decorations, os, eachFn);
838       os << "]";
839     }
840   };
841   llvm::interleaveComma(llvm::seq<unsigned>(0, type.getNumElements()), os,
842                         printMember);
843   os << ")>";
844 
845   if (type.isIdentified())
846     structContext.remove(type.getIdentifier());
847 }
848 
849 static void print(CooperativeMatrixNVType type, DialectAsmPrinter &os) {
850   os << "coopmatrix<" << type.getRows() << "x" << type.getColumns() << "x";
851   os << type.getElementType() << ", " << stringifyScope(type.getScope());
852   os << ">";
853 }
854 
855 static void print(MatrixType type, DialectAsmPrinter &os) {
856   os << "matrix<" << type.getNumColumns() << " x " << type.getColumnType();
857   os << ">";
858 }
859 
860 void SPIRVDialect::printType(Type type, DialectAsmPrinter &os) const {
861   TypeSwitch<Type>(type)
862       .Case<ArrayType, CooperativeMatrixNVType, PointerType, RuntimeArrayType,
863             ImageType, SampledImageType, StructType, MatrixType>(
864           [&](auto type) { print(type, os); })
865       .Default([](Type) { llvm_unreachable("unhandled SPIR-V type"); });
866 }
867 
868 //===----------------------------------------------------------------------===//
869 // Attribute Parsing
870 //===----------------------------------------------------------------------===//
871 
872 /// Parses a comma-separated list of keywords, invokes `processKeyword` on each
873 /// of the parsed keyword, and returns failure if any error occurs.
874 static ParseResult parseKeywordList(
875     DialectAsmParser &parser,
876     function_ref<LogicalResult(llvm::SMLoc, StringRef)> processKeyword) {
877   if (parser.parseLSquare())
878     return failure();
879 
880   // Special case for empty list.
881   if (succeeded(parser.parseOptionalRSquare()))
882     return success();
883 
884   // Keep parsing the keyword and an optional comma following it. If the comma
885   // is successfully parsed, then we have more keywords to parse.
886   do {
887     auto loc = parser.getCurrentLocation();
888     StringRef keyword;
889     if (parser.parseKeyword(&keyword) || failed(processKeyword(loc, keyword)))
890       return failure();
891   } while (succeeded(parser.parseOptionalComma()));
892 
893   if (parser.parseRSquare())
894     return failure();
895 
896   return success();
897 }
898 
899 /// Parses a spirv::InterfaceVarABIAttr.
900 static Attribute parseInterfaceVarABIAttr(DialectAsmParser &parser) {
901   if (parser.parseLess())
902     return {};
903 
904   Builder &builder = parser.getBuilder();
905 
906   if (parser.parseLParen())
907     return {};
908 
909   IntegerAttr descriptorSetAttr;
910   {
911     auto loc = parser.getCurrentLocation();
912     uint32_t descriptorSet = 0;
913     auto descriptorSetParseResult = parser.parseOptionalInteger(descriptorSet);
914 
915     if (!descriptorSetParseResult.hasValue() ||
916         failed(*descriptorSetParseResult)) {
917       parser.emitError(loc, "missing descriptor set");
918       return {};
919     }
920     descriptorSetAttr = builder.getI32IntegerAttr(descriptorSet);
921   }
922 
923   if (parser.parseComma())
924     return {};
925 
926   IntegerAttr bindingAttr;
927   {
928     auto loc = parser.getCurrentLocation();
929     uint32_t binding = 0;
930     auto bindingParseResult = parser.parseOptionalInteger(binding);
931 
932     if (!bindingParseResult.hasValue() || failed(*bindingParseResult)) {
933       parser.emitError(loc, "missing binding");
934       return {};
935     }
936     bindingAttr = builder.getI32IntegerAttr(binding);
937   }
938 
939   if (parser.parseRParen())
940     return {};
941 
942   IntegerAttr storageClassAttr;
943   {
944     if (succeeded(parser.parseOptionalComma())) {
945       auto loc = parser.getCurrentLocation();
946       StringRef storageClass;
947       if (parser.parseKeyword(&storageClass))
948         return {};
949 
950       if (auto storageClassSymbol =
951               spirv::symbolizeStorageClass(storageClass)) {
952         storageClassAttr = builder.getI32IntegerAttr(
953             static_cast<uint32_t>(*storageClassSymbol));
954       } else {
955         parser.emitError(loc, "unknown storage class: ") << storageClass;
956         return {};
957       }
958     }
959   }
960 
961   if (parser.parseGreater())
962     return {};
963 
964   return spirv::InterfaceVarABIAttr::get(descriptorSetAttr, bindingAttr,
965                                          storageClassAttr);
966 }
967 
968 static Attribute parseVerCapExtAttr(DialectAsmParser &parser) {
969   if (parser.parseLess())
970     return {};
971 
972   Builder &builder = parser.getBuilder();
973 
974   IntegerAttr versionAttr;
975   {
976     auto loc = parser.getCurrentLocation();
977     StringRef version;
978     if (parser.parseKeyword(&version) || parser.parseComma())
979       return {};
980 
981     if (auto versionSymbol = spirv::symbolizeVersion(version)) {
982       versionAttr =
983           builder.getI32IntegerAttr(static_cast<uint32_t>(*versionSymbol));
984     } else {
985       parser.emitError(loc, "unknown version: ") << version;
986       return {};
987     }
988   }
989 
990   ArrayAttr capabilitiesAttr;
991   {
992     SmallVector<Attribute, 4> capabilities;
993     llvm::SMLoc errorloc;
994     StringRef errorKeyword;
995 
996     auto processCapability = [&](llvm::SMLoc loc, StringRef capability) {
997       if (auto capSymbol = spirv::symbolizeCapability(capability)) {
998         capabilities.push_back(
999             builder.getI32IntegerAttr(static_cast<uint32_t>(*capSymbol)));
1000         return success();
1001       }
1002       return errorloc = loc, errorKeyword = capability, failure();
1003     };
1004     if (parseKeywordList(parser, processCapability) || parser.parseComma()) {
1005       if (!errorKeyword.empty())
1006         parser.emitError(errorloc, "unknown capability: ") << errorKeyword;
1007       return {};
1008     }
1009 
1010     capabilitiesAttr = builder.getArrayAttr(capabilities);
1011   }
1012 
1013   ArrayAttr extensionsAttr;
1014   {
1015     SmallVector<Attribute, 1> extensions;
1016     llvm::SMLoc errorloc;
1017     StringRef errorKeyword;
1018 
1019     auto processExtension = [&](llvm::SMLoc loc, StringRef extension) {
1020       if (spirv::symbolizeExtension(extension)) {
1021         extensions.push_back(builder.getStringAttr(extension));
1022         return success();
1023       }
1024       return errorloc = loc, errorKeyword = extension, failure();
1025     };
1026     if (parseKeywordList(parser, processExtension)) {
1027       if (!errorKeyword.empty())
1028         parser.emitError(errorloc, "unknown extension: ") << errorKeyword;
1029       return {};
1030     }
1031 
1032     extensionsAttr = builder.getArrayAttr(extensions);
1033   }
1034 
1035   if (parser.parseGreater())
1036     return {};
1037 
1038   return spirv::VerCapExtAttr::get(versionAttr, capabilitiesAttr,
1039                                    extensionsAttr);
1040 }
1041 
1042 /// Parses a spirv::TargetEnvAttr.
1043 static Attribute parseTargetEnvAttr(DialectAsmParser &parser) {
1044   if (parser.parseLess())
1045     return {};
1046 
1047   spirv::VerCapExtAttr tripleAttr;
1048   if (parser.parseAttribute(tripleAttr) || parser.parseComma())
1049     return {};
1050 
1051   // Parse [vendor[:device-type[:device-id]]]
1052   Vendor vendorID = Vendor::Unknown;
1053   DeviceType deviceType = DeviceType::Unknown;
1054   uint32_t deviceID = spirv::TargetEnvAttr::kUnknownDeviceID;
1055   {
1056     auto loc = parser.getCurrentLocation();
1057     StringRef vendorStr;
1058     if (succeeded(parser.parseOptionalKeyword(&vendorStr))) {
1059       if (auto vendorSymbol = spirv::symbolizeVendor(vendorStr)) {
1060         vendorID = *vendorSymbol;
1061       } else {
1062         parser.emitError(loc, "unknown vendor: ") << vendorStr;
1063       }
1064 
1065       if (succeeded(parser.parseOptionalColon())) {
1066         loc = parser.getCurrentLocation();
1067         StringRef deviceTypeStr;
1068         if (parser.parseKeyword(&deviceTypeStr))
1069           return {};
1070         if (auto deviceTypeSymbol = spirv::symbolizeDeviceType(deviceTypeStr)) {
1071           deviceType = *deviceTypeSymbol;
1072         } else {
1073           parser.emitError(loc, "unknown device type: ") << deviceTypeStr;
1074         }
1075 
1076         if (succeeded(parser.parseOptionalColon())) {
1077           loc = parser.getCurrentLocation();
1078           if (parser.parseInteger(deviceID))
1079             return {};
1080         }
1081       }
1082       if (parser.parseComma())
1083         return {};
1084     }
1085   }
1086 
1087   DictionaryAttr limitsAttr;
1088   {
1089     auto loc = parser.getCurrentLocation();
1090     if (parser.parseAttribute(limitsAttr))
1091       return {};
1092 
1093     if (!limitsAttr.isa<spirv::ResourceLimitsAttr>()) {
1094       parser.emitError(
1095           loc,
1096           "limits must be a dictionary attribute containing two 32-bit integer "
1097           "attributes 'max_compute_workgroup_invocations' and "
1098           "'max_compute_workgroup_size'");
1099       return {};
1100     }
1101   }
1102 
1103   if (parser.parseGreater())
1104     return {};
1105 
1106   return spirv::TargetEnvAttr::get(tripleAttr, vendorID, deviceType, deviceID,
1107                                    limitsAttr);
1108 }
1109 
1110 Attribute SPIRVDialect::parseAttribute(DialectAsmParser &parser,
1111                                        Type type) const {
1112   // SPIR-V attributes are dictionaries so they do not have type.
1113   if (type) {
1114     parser.emitError(parser.getNameLoc(), "unexpected type");
1115     return {};
1116   }
1117 
1118   // Parse the kind keyword first.
1119   StringRef attrKind;
1120   if (parser.parseKeyword(&attrKind))
1121     return {};
1122 
1123   if (attrKind == spirv::TargetEnvAttr::getKindName())
1124     return parseTargetEnvAttr(parser);
1125   if (attrKind == spirv::VerCapExtAttr::getKindName())
1126     return parseVerCapExtAttr(parser);
1127   if (attrKind == spirv::InterfaceVarABIAttr::getKindName())
1128     return parseInterfaceVarABIAttr(parser);
1129 
1130   parser.emitError(parser.getNameLoc(), "unknown SPIR-V attribute kind: ")
1131       << attrKind;
1132   return {};
1133 }
1134 
1135 //===----------------------------------------------------------------------===//
1136 // Attribute Printing
1137 //===----------------------------------------------------------------------===//
1138 
1139 static void print(spirv::VerCapExtAttr triple, DialectAsmPrinter &printer) {
1140   auto &os = printer.getStream();
1141   printer << spirv::VerCapExtAttr::getKindName() << "<"
1142           << spirv::stringifyVersion(triple.getVersion()) << ", [";
1143   llvm::interleaveComma(
1144       triple.getCapabilities(), os,
1145       [&](spirv::Capability cap) { os << spirv::stringifyCapability(cap); });
1146   printer << "], [";
1147   llvm::interleaveComma(triple.getExtensionsAttr(), os, [&](Attribute attr) {
1148     os << attr.cast<StringAttr>().getValue();
1149   });
1150   printer << "]>";
1151 }
1152 
1153 static void print(spirv::TargetEnvAttr targetEnv, DialectAsmPrinter &printer) {
1154   printer << spirv::TargetEnvAttr::getKindName() << "<#spv.";
1155   print(targetEnv.getTripleAttr(), printer);
1156   spirv::Vendor vendorID = targetEnv.getVendorID();
1157   spirv::DeviceType deviceType = targetEnv.getDeviceType();
1158   uint32_t deviceID = targetEnv.getDeviceID();
1159   if (vendorID != spirv::Vendor::Unknown) {
1160     printer << ", " << spirv::stringifyVendor(vendorID);
1161     if (deviceType != spirv::DeviceType::Unknown) {
1162       printer << ":" << spirv::stringifyDeviceType(deviceType);
1163       if (deviceID != spirv::TargetEnvAttr::kUnknownDeviceID)
1164         printer << ":" << deviceID;
1165     }
1166   }
1167   printer << ", " << targetEnv.getResourceLimits() << ">";
1168 }
1169 
1170 static void print(spirv::InterfaceVarABIAttr interfaceVarABIAttr,
1171                   DialectAsmPrinter &printer) {
1172   printer << spirv::InterfaceVarABIAttr::getKindName() << "<("
1173           << interfaceVarABIAttr.getDescriptorSet() << ", "
1174           << interfaceVarABIAttr.getBinding() << ")";
1175   auto storageClass = interfaceVarABIAttr.getStorageClass();
1176   if (storageClass)
1177     printer << ", " << spirv::stringifyStorageClass(*storageClass);
1178   printer << ">";
1179 }
1180 
1181 void SPIRVDialect::printAttribute(Attribute attr,
1182                                   DialectAsmPrinter &printer) const {
1183   if (auto targetEnv = attr.dyn_cast<TargetEnvAttr>())
1184     print(targetEnv, printer);
1185   else if (auto vceAttr = attr.dyn_cast<VerCapExtAttr>())
1186     print(vceAttr, printer);
1187   else if (auto interfaceVarABIAttr = attr.dyn_cast<InterfaceVarABIAttr>())
1188     print(interfaceVarABIAttr, printer);
1189   else
1190     llvm_unreachable("unhandled SPIR-V attribute kind");
1191 }
1192 
1193 //===----------------------------------------------------------------------===//
1194 // Constant
1195 //===----------------------------------------------------------------------===//
1196 
1197 Operation *SPIRVDialect::materializeConstant(OpBuilder &builder,
1198                                              Attribute value, Type type,
1199                                              Location loc) {
1200   if (!spirv::ConstantOp::isBuildableWith(type))
1201     return nullptr;
1202 
1203   return builder.create<spirv::ConstantOp>(loc, type, value);
1204 }
1205 
1206 //===----------------------------------------------------------------------===//
1207 // Shader Interface ABI
1208 //===----------------------------------------------------------------------===//
1209 
1210 LogicalResult SPIRVDialect::verifyOperationAttribute(Operation *op,
1211                                                      NamedAttribute attribute) {
1212   StringRef symbol = attribute.first.strref();
1213   Attribute attr = attribute.second;
1214 
1215   // TODO: figure out a way to generate the description from the
1216   // StructAttr definition.
1217   if (symbol == spirv::getEntryPointABIAttrName()) {
1218     if (!attr.isa<spirv::EntryPointABIAttr>())
1219       return op->emitError("'")
1220              << symbol
1221              << "' attribute must be a dictionary attribute containing one "
1222                 "32-bit integer elements attribute: 'local_size'";
1223   } else if (symbol == spirv::getTargetEnvAttrName()) {
1224     if (!attr.isa<spirv::TargetEnvAttr>())
1225       return op->emitError("'") << symbol << "' must be a spirv::TargetEnvAttr";
1226   } else {
1227     return op->emitError("found unsupported '")
1228            << symbol << "' attribute on operation";
1229   }
1230 
1231   return success();
1232 }
1233 
1234 /// Verifies the given SPIR-V `attribute` attached to a value of the given
1235 /// `valueType` is valid.
1236 static LogicalResult verifyRegionAttribute(Location loc, Type valueType,
1237                                            NamedAttribute attribute) {
1238   StringRef symbol = attribute.first.strref();
1239   Attribute attr = attribute.second;
1240 
1241   if (symbol != spirv::getInterfaceVarABIAttrName())
1242     return emitError(loc, "found unsupported '")
1243            << symbol << "' attribute on region argument";
1244 
1245   auto varABIAttr = attr.dyn_cast<spirv::InterfaceVarABIAttr>();
1246   if (!varABIAttr)
1247     return emitError(loc, "'")
1248            << symbol << "' must be a spirv::InterfaceVarABIAttr";
1249 
1250   if (varABIAttr.getStorageClass() && !valueType.isIntOrIndexOrFloat())
1251     return emitError(loc, "'") << symbol
1252                                << "' attribute cannot specify storage class "
1253                                   "when attaching to a non-scalar value";
1254 
1255   return success();
1256 }
1257 
1258 LogicalResult SPIRVDialect::verifyRegionArgAttribute(Operation *op,
1259                                                      unsigned regionIndex,
1260                                                      unsigned argIndex,
1261                                                      NamedAttribute attribute) {
1262   return verifyRegionAttribute(
1263       op->getLoc(), op->getRegion(regionIndex).getArgument(argIndex).getType(),
1264       attribute);
1265 }
1266 
1267 LogicalResult SPIRVDialect::verifyRegionResultAttribute(
1268     Operation *op, unsigned /*regionIndex*/, unsigned /*resultIndex*/,
1269     NamedAttribute attribute) {
1270   return op->emitError("cannot attach SPIR-V attributes to region result");
1271 }
1272