1 //===- TestDialect.h - MLIR Dialect for testing -----------------*- 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 defines a fake 'test' dialect that can be used for testing things
10 // that do not have a respective counterpart in the main source directories.
11 //
12 //===----------------------------------------------------------------------===//
13 
14 #ifndef MLIR_TESTDIALECT_H
15 #define MLIR_TESTDIALECT_H
16 
17 #include "TestAttributes.h"
18 #include "TestInterfaces.h"
19 #include "mlir/Dialect/DLTI/DLTI.h"
20 #include "mlir/Dialect/DLTI/Traits.h"
21 #include "mlir/Dialect/Func/IR/FuncOps.h"
22 #include "mlir/Dialect/Linalg/IR/Linalg.h"
23 #include "mlir/Dialect/Traits.h"
24 #include "mlir/IR/AsmState.h"
25 #include "mlir/IR/BuiltinOps.h"
26 #include "mlir/IR/BuiltinTypes.h"
27 #include "mlir/IR/Dialect.h"
28 #include "mlir/IR/ExtensibleDialect.h"
29 #include "mlir/IR/OpDefinition.h"
30 #include "mlir/IR/OpImplementation.h"
31 #include "mlir/IR/RegionKindInterface.h"
32 #include "mlir/IR/SymbolTable.h"
33 #include "mlir/Interfaces/CallInterfaces.h"
34 #include "mlir/Interfaces/ControlFlowInterfaces.h"
35 #include "mlir/Interfaces/CopyOpInterface.h"
36 #include "mlir/Interfaces/DerivedAttributeOpInterface.h"
37 #include "mlir/Interfaces/InferIntRangeInterface.h"
38 #include "mlir/Interfaces/InferTypeOpInterface.h"
39 #include "mlir/Interfaces/LoopLikeInterface.h"
40 #include "mlir/Interfaces/SideEffectInterfaces.h"
41 #include "mlir/Interfaces/ViewLikeInterface.h"
42 
43 namespace mlir {
44 class DLTIDialect;
45 class RewritePatternSet;
46 } // namespace mlir
47 
48 namespace test {
49 class TestDialect;
50 
51 //===----------------------------------------------------------------------===//
52 // External Elements Data
53 //===----------------------------------------------------------------------===//
54 
55 /// This class represents a single external elements instance. It keeps track of
56 /// the data, and deallocates when destructed.
57 class TestExternalElementsData : public mlir::AsmResourceBlob {
58 public:
59   using mlir::AsmResourceBlob::AsmResourceBlob;
TestExternalElementsData(mlir::AsmResourceBlob && blob)60   TestExternalElementsData(mlir::AsmResourceBlob &&blob)
61       : mlir::AsmResourceBlob(std::move(blob)) {}
62 
63   /// Return the data of this external elements instance.
64   llvm::ArrayRef<uint64_t> getData() const;
65 
66   /// Allocate a new external elements instance with the given number of
67   /// elements.
68   static TestExternalElementsData allocate(size_t numElements);
69 };
70 
71 /// A handle used to reference external elements instances.
72 struct TestExternalElementsDataHandle
73     : public mlir::AsmDialectResourceHandleBase<
74           TestExternalElementsDataHandle,
75           llvm::StringMapEntry<std::unique_ptr<TestExternalElementsData>>,
76           TestDialect> {
77   using AsmDialectResourceHandleBase::AsmDialectResourceHandleBase;
78 
79   /// Return a key to use for this handle.
getKeyTestExternalElementsDataHandle80   llvm::StringRef getKey() const { return getResource()->getKey(); }
81 
82   /// Return the data referenced by this handle.
getDataTestExternalElementsDataHandle83   TestExternalElementsData *getData() const {
84     return getResource()->getValue().get();
85   }
86 };
87 
88 /// This class acts as a manager for external elements data. It provides API
89 /// for creating and accessing registered elements data.
90 class TestExternalElementsDataManager {
91   using DataMap = llvm::StringMap<std::unique_ptr<TestExternalElementsData>>;
92 
93 public:
94   /// Return the data registered for the given name, or nullptr if no data is
95   /// registered.
96   const TestExternalElementsData *getData(llvm::StringRef name) const;
97 
98   /// Register an entry with the provided name, which may be modified if another
99   /// entry was already inserted with that name. Returns the inserted entry.
100   std::pair<DataMap::iterator, bool> insert(llvm::StringRef name);
101 
102   /// Set the data for the given entry, which is expected to exist.
103   void setData(llvm::StringRef name, TestExternalElementsData &&data);
104 
105 private:
106   llvm::StringMap<std::unique_ptr<TestExternalElementsData>> dataMap;
107 };
108 } // namespace test
109 
110 //===----------------------------------------------------------------------===//
111 // TestDialect
112 //===----------------------------------------------------------------------===//
113 
114 #include "TestOpInterfaces.h.inc"
115 #include "TestOpsDialect.h.inc"
116 
117 #define GET_OP_CLASSES
118 #include "TestOps.h.inc"
119 
120 namespace test {
121 void registerTestDialect(::mlir::DialectRegistry &registry);
122 void populateTestReductionPatterns(::mlir::RewritePatternSet &patterns);
123 } // namespace test
124 
125 #endif // MLIR_TESTDIALECT_H
126