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