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/Analysis/DataLayoutAnalysis.h"
11 #include "mlir/Dialect/DLTI/DLTI.h"
12 #include "mlir/IR/BuiltinAttributes.h"
13 #include "mlir/Pass/Pass.h"
14 
15 using namespace mlir;
16 
17 namespace {
18 
19 /// A pass that finds "test.data_layout_query" operations and attaches to them
20 /// attributes containing the results of data layout queries for operation
21 /// result types.
22 struct TestDataLayoutQuery
23     : public PassWrapper<TestDataLayoutQuery, FunctionPass> {
getArgument__anon6d2c1c400111::TestDataLayoutQuery24   StringRef getArgument() const final { return "test-data-layout-query"; }
getDescription__anon6d2c1c400111::TestDataLayoutQuery25   StringRef getDescription() const final { return "Test data layout queries"; }
runOnFunction__anon6d2c1c400111::TestDataLayoutQuery26   void runOnFunction() override {
27     FuncOp func = getFunction();
28     Builder builder(func.getContext());
29     const DataLayoutAnalysis &layouts = getAnalysis<DataLayoutAnalysis>();
30 
31     func.walk([&](test::DataLayoutQueryOp op) {
32       // Skip the ops with already processed in a deeper call.
33       if (op->getAttr("size"))
34         return;
35 
36       const DataLayout &layout = layouts.getAbove(op);
37       unsigned size = layout.getTypeSize(op.getType());
38       unsigned bitsize = layout.getTypeSizeInBits(op.getType());
39       unsigned alignment = layout.getTypeABIAlignment(op.getType());
40       unsigned preferred = layout.getTypePreferredAlignment(op.getType());
41       op->setAttrs(
42           {builder.getNamedAttr("size", builder.getIndexAttr(size)),
43            builder.getNamedAttr("bitsize", builder.getIndexAttr(bitsize)),
44            builder.getNamedAttr("alignment", builder.getIndexAttr(alignment)),
45            builder.getNamedAttr("preferred", builder.getIndexAttr(preferred))});
46     });
47   }
48 };
49 } // namespace
50 
51 namespace mlir {
52 namespace test {
registerTestDataLayoutQuery()53 void registerTestDataLayoutQuery() { PassRegistration<TestDataLayoutQuery>(); }
54 } // namespace test
55 } // namespace mlir
56