xref: /llvm-project/mlir/test/lib/Dialect/DLTI/TestDataLayoutQuery.cpp (revision e95e94adc6bb748de015ac3053e7f0786b65f351)
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 "TestOps.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, OperationPass<func::FuncOp>> {
MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID__anon6b9d4f3a0111::TestDataLayoutQuery24   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestDataLayoutQuery)
25 
26   StringRef getArgument() const final { return "test-data-layout-query"; }
getDescription__anon6b9d4f3a0111::TestDataLayoutQuery27   StringRef getDescription() const final { return "Test data layout queries"; }
runOnOperation__anon6b9d4f3a0111::TestDataLayoutQuery28   void runOnOperation() override {
29     func::FuncOp func = getOperation();
30     Builder builder(func.getContext());
31     const DataLayoutAnalysis &layouts = getAnalysis<DataLayoutAnalysis>();
32 
33     func.walk([&](test::DataLayoutQueryOp op) {
34       // Skip the ops with already processed in a deeper call.
35       if (op->getDiscardableAttr("size"))
36         return;
37 
38       const DataLayout &layout = layouts.getAbove(op);
39       llvm::TypeSize size = layout.getTypeSize(op.getType());
40       llvm::TypeSize bitsize = layout.getTypeSizeInBits(op.getType());
41       uint64_t alignment = layout.getTypeABIAlignment(op.getType());
42       uint64_t preferred = layout.getTypePreferredAlignment(op.getType());
43       uint64_t index = layout.getTypeIndexBitwidth(op.getType()).value_or(0);
44       Attribute endianness = layout.getEndianness();
45       Attribute allocaMemorySpace = layout.getAllocaMemorySpace();
46       Attribute programMemorySpace = layout.getProgramMemorySpace();
47       Attribute globalMemorySpace = layout.getGlobalMemorySpace();
48       uint64_t stackAlignment = layout.getStackAlignment();
49 
50       auto convertTypeSizeToAttr = [&](llvm::TypeSize typeSize) -> Attribute {
51         if (!typeSize.isScalable())
52           return builder.getIndexAttr(typeSize);
53 
54         return builder.getDictionaryAttr({
55             builder.getNamedAttr("scalable", builder.getUnitAttr()),
56             builder.getNamedAttr(
57                 "minimal_size",
58                 builder.getIndexAttr(typeSize.getKnownMinValue())),
59         });
60       };
61 
62       op->setAttrs(
63           {builder.getNamedAttr("size", convertTypeSizeToAttr(size)),
64            builder.getNamedAttr("bitsize", convertTypeSizeToAttr(bitsize)),
65            builder.getNamedAttr("alignment", builder.getIndexAttr(alignment)),
66            builder.getNamedAttr("preferred", builder.getIndexAttr(preferred)),
67            builder.getNamedAttr("index", builder.getIndexAttr(index)),
68            builder.getNamedAttr("endianness", endianness == Attribute()
69                                                   ? builder.getStringAttr("")
70                                                   : endianness),
71            builder.getNamedAttr("alloca_memory_space",
72                                 allocaMemorySpace == Attribute()
73                                     ? builder.getUI32IntegerAttr(0)
74                                     : allocaMemorySpace),
75            builder.getNamedAttr("program_memory_space",
76                                 programMemorySpace == Attribute()
77                                     ? builder.getUI32IntegerAttr(0)
78                                     : programMemorySpace),
79            builder.getNamedAttr("global_memory_space",
80                                 globalMemorySpace == Attribute()
81                                     ? builder.getUI32IntegerAttr(0)
82                                     : globalMemorySpace),
83            builder.getNamedAttr("stack_alignment",
84                                 builder.getIndexAttr(stackAlignment))});
85     });
86   }
87 };
88 } // namespace
89 
90 namespace mlir {
91 namespace test {
registerTestDataLayoutQuery()92 void registerTestDataLayoutQuery() { PassRegistration<TestDataLayoutQuery>(); }
93 } // namespace test
94 } // namespace mlir
95