1 //===- TestDataLayoutQuery.cpp - Test Data Layout Queries -----------------===//
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/Dialect/DLTI/DLTI.h"
11 #include "mlir/IR/BuiltinAttributes.h"
12 #include "mlir/Pass/Pass.h"
13 
14 using namespace mlir;
15 
16 namespace {
17 
18 /// A pass that finds "test.data_layout_query" operations and attaches to them
19 /// attributes containing the results of data layout queries for operation
20 /// result types.
21 struct TestDataLayoutQuery
22     : public PassWrapper<TestDataLayoutQuery, FunctionPass> {
23   void runOnFunction() override {
24     FuncOp func = getFunction();
25     Builder builder(func.getContext());
26     DenseMap<Operation *, DataLayout> layouts;
27 
28     func.walk([&](test::DataLayoutQueryOp op) {
29       // Skip the ops with already processed in a deeper call.
30       if (op->getAttr("size"))
31         return;
32 
33       auto scope = op->getParentOfType<test::OpWithDataLayoutOp>();
34       if (!layouts.count(scope)) {
35         layouts.try_emplace(
36             scope, scope ? cast<DataLayoutOpInterface>(scope.getOperation())
37                          : nullptr);
38       }
39       auto module = op->getParentOfType<ModuleOp>();
40       if (!layouts.count(module))
41         layouts.try_emplace(module, module);
42 
43       Operation *closest = (scope && module && module->isProperAncestor(scope))
44                                ? scope.getOperation()
45                                : module.getOperation();
46 
47       const DataLayout &layout = layouts.find(closest)->getSecond();
48       unsigned size = layout.getTypeSize(op.getType());
49       unsigned bitsize = layout.getTypeSizeInBits(op.getType());
50       unsigned alignment = layout.getTypeABIAlignment(op.getType());
51       unsigned preferred = layout.getTypePreferredAlignment(op.getType());
52       op->setAttrs(
53           {builder.getNamedAttr("size", builder.getIndexAttr(size)),
54            builder.getNamedAttr("bitsize", builder.getIndexAttr(bitsize)),
55            builder.getNamedAttr("alignment", builder.getIndexAttr(alignment)),
56            builder.getNamedAttr("preferred", builder.getIndexAttr(preferred))});
57     });
58   }
59 };
60 } // namespace
61 
62 namespace mlir {
63 namespace test {
64 void registerTestDataLayoutQuery() {
65   PassRegistration<TestDataLayoutQuery>("test-data-layout-query",
66                                         "Test data layout queries");
67 }
68 } // namespace test
69 } // namespace mlir
70