1 | //===- TestTypes.h - MLIR Test Dialect Types --------------------*- C++ -*-===// |
2 | // |
3 | // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. |
4 | // See https://llvm.org/LICENSE.txt for license information. |
5 | // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception |
6 | // |
7 | //===----------------------------------------------------------------------===// |
8 | // |
9 | // This file contains types defined by the TestDialect for testing various |
10 | // features of MLIR. |
11 | // |
12 | //===----------------------------------------------------------------------===// |
13 | |
14 | #ifndef MLIR_TESTTYPES_H |
15 | #define MLIR_TESTTYPES_H |
16 | |
17 | #include <optional> |
18 | #include <tuple> |
19 | |
20 | #include "TestTraits.h" |
21 | #include "mlir/Dialect/Bufferization/IR/BufferizationTypeInterfaces.h" |
22 | #include "mlir/IR/Diagnostics.h" |
23 | #include "mlir/IR/Dialect.h" |
24 | #include "mlir/IR/DialectImplementation.h" |
25 | #include "mlir/IR/Operation.h" |
26 | #include "mlir/IR/Types.h" |
27 | #include "mlir/Interfaces/DataLayoutInterfaces.h" |
28 | |
29 | namespace test { |
30 | class TestAttrWithFormatAttr; |
31 | |
32 | /// FieldInfo represents a field in the StructType data type. It is used as a |
33 | /// parameter in TestTypeDefs.td. |
34 | struct FieldInfo { |
35 | llvm::StringRef name; |
36 | mlir::Type type; |
37 | |
38 | // Custom allocation called from generated constructor code |
39 | FieldInfo allocateInto(mlir::TypeStorageAllocator &alloc) const { |
40 | return FieldInfo{.name: alloc.copyInto(str: name), .type: type}; |
41 | } |
42 | }; |
43 | |
44 | /// A custom type for a test type parameter. |
45 | struct CustomParam { |
46 | int value; |
47 | |
48 | bool operator==(const CustomParam &other) const { |
49 | return other.value == value; |
50 | } |
51 | }; |
52 | |
53 | inline llvm::hash_code hash_value(const test::CustomParam ¶m) { |
54 | return llvm::hash_value(value: param.value); |
55 | } |
56 | |
57 | } // namespace test |
58 | |
59 | namespace mlir { |
60 | template <> |
61 | struct FieldParser<test::CustomParam> { |
62 | static FailureOr<test::CustomParam> parse(AsmParser &parser) { |
63 | auto value = FieldParser<int>::parse(parser); |
64 | if (failed(Result: value)) |
65 | return failure(); |
66 | return test::CustomParam{.value: *value}; |
67 | } |
68 | }; |
69 | |
70 | inline mlir::AsmPrinter &operator<<(mlir::AsmPrinter &printer, |
71 | test::CustomParam param) { |
72 | return printer << param.value; |
73 | } |
74 | |
75 | /// Overload the attribute parameter parser for optional integers. |
76 | template <> |
77 | struct FieldParser<std::optional<int>> { |
78 | static FailureOr<std::optional<int>> parse(AsmParser &parser) { |
79 | std::optional<int> value; |
80 | value.emplace(); |
81 | OptionalParseResult result = parser.parseOptionalInteger(result&: *value); |
82 | if (result.has_value()) { |
83 | if (succeeded(Result: *result)) |
84 | return value; |
85 | return failure(); |
86 | } |
87 | value.reset(); |
88 | return value; |
89 | } |
90 | }; |
91 | } // namespace mlir |
92 | |
93 | #include "TestTypeInterfaces.h.inc" |
94 | |
95 | namespace test { |
96 | |
97 | /// Storage for simple named recursive types, where the type is identified by |
98 | /// its name and can "contain" another type, including itself. |
99 | struct TestRecursiveTypeStorage : public ::mlir::TypeStorage { |
100 | using KeyTy = ::llvm::StringRef; |
101 | |
102 | explicit TestRecursiveTypeStorage(::llvm::StringRef key) : name(key) {} |
103 | |
104 | bool operator==(const KeyTy &other) const { return name == other; } |
105 | |
106 | static TestRecursiveTypeStorage * |
107 | construct(::mlir::TypeStorageAllocator &allocator, const KeyTy &key) { |
108 | return new (allocator.allocate<TestRecursiveTypeStorage>()) |
109 | TestRecursiveTypeStorage(allocator.copyInto(str: key)); |
110 | } |
111 | |
112 | ::llvm::LogicalResult mutate(::mlir::TypeStorageAllocator &allocator, |
113 | ::mlir::Type newBody) { |
114 | // Cannot set a different body than before. |
115 | if (body && body != newBody) |
116 | return ::mlir::failure(); |
117 | |
118 | body = newBody; |
119 | return ::mlir::success(); |
120 | } |
121 | |
122 | ::llvm::StringRef name; |
123 | ::mlir::Type body; |
124 | }; |
125 | |
126 | /// Simple recursive type identified by its name and pointing to another named |
127 | /// type, potentially itself. This requires the body to be mutated separately |
128 | /// from type creation. |
129 | class TestRecursiveType |
130 | : public ::mlir::Type::TypeBase<TestRecursiveType, ::mlir::Type, |
131 | TestRecursiveTypeStorage, |
132 | ::mlir::TypeTrait::IsMutable> { |
133 | public: |
134 | using Base::Base; |
135 | |
136 | static constexpr ::mlir::StringLiteral name = "test.recursive" ; |
137 | |
138 | static TestRecursiveType get(::mlir::MLIRContext *ctx, |
139 | ::llvm::StringRef name) { |
140 | return Base::get(ctx, name); |
141 | } |
142 | |
143 | /// Body getter and setter. |
144 | ::llvm::LogicalResult setBody(Type body) { return Base::mutate(body); } |
145 | ::mlir::Type getBody() const { return getImpl()->body; } |
146 | |
147 | /// Name/key getter. |
148 | ::llvm::StringRef getName() { return getImpl()->name; } |
149 | }; |
150 | |
151 | } // namespace test |
152 | |
153 | #define GET_TYPEDEF_CLASSES |
154 | #include "TestTypeDefs.h.inc" |
155 | |
156 | #endif // MLIR_TESTTYPES_H |
157 | |