mlir/test/lib/Dialect/DLTI/TestDataLayoutQuery.cpp

   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>> {
  24   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestDataLayoutQuery)
  25
  26   StringRef getArgument() const final { return "test-data-layout-query"; }
  27   StringRef getDescription() const final { return "Test data layout queries"; }
  28   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 {
  92 void registerTestDataLayoutQuery() { PassRegistration<TestDataLayoutQuery>(); }
  93 } // namespace test
  94 } // namespace mlir