xref: /llvm-project/mlir/test/lib/IR/TestBuiltinDistinctAttributes.cpp (revision 728a8d5a81f35f24492fcd31928ef95e597a66af)
1 //===- TestBuiltinDistinctAttributes.cpp - Test DistinctAttributes --------===//
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 #include "TestDialect.h"
10 #include "mlir/IR/BuiltinAttributes.h"
11 #include "mlir/Pass/Pass.h"
12 
13 using namespace mlir;
14 
15 namespace {
16 /// This is a distinct attribute test pass that tests if distinct attributes can
17 /// be created in parallel in a deterministic way.
18 struct DistinctAttributesPass
19     : public PassWrapper<DistinctAttributesPass, OperationPass<func::FuncOp>> {
MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID__anon2fd4c7550111::DistinctAttributesPass20   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(DistinctAttributesPass)
21 
22   StringRef getArgument() const final { return "test-distinct-attrs"; }
getDescription__anon2fd4c7550111::DistinctAttributesPass23   StringRef getDescription() const final {
24     return "Test parallel creation of distinct attributes";
25   }
26 
runOnOperation__anon2fd4c7550111::DistinctAttributesPass27   void runOnOperation() override {
28     auto funcOp = getOperation();
29 
30     /// Walk all operations and create a distinct output attribute given a
31     /// distinct input attribute.
32     funcOp->walk([](Operation *op) {
33       auto distinctAttr = op->getAttrOfType<DistinctAttr>("distinct.input");
34       if (!distinctAttr)
35         return;
36       op->setAttr("distinct.output",
37                   DistinctAttr::create(distinctAttr.getReferencedAttr()));
38     });
39   }
40 };
41 } // namespace
42 
43 namespace mlir {
44 namespace test {
registerTestBuiltinDistinctAttributes()45 void registerTestBuiltinDistinctAttributes() {
46   PassRegistration<DistinctAttributesPass>();
47 }
48 } // namespace test
49 } // namespace mlir
50