1//===- InferTypeOpInterfaceTest.cpp - Unit Test for type interface --------===//
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 "mlir/Interfaces/InferTypeOpInterface.h"
10#include "mlir/Dialect/Arith/IR/Arith.h"
11#include "mlir/Dialect/Func/IR/FuncOps.h"
12#include "mlir/IR/Builders.h"
13#include "mlir/IR/BuiltinOps.h"
14#include "mlir/IR/Dialect.h"
15#include "mlir/IR/DialectImplementation.h"
16#include "mlir/IR/ImplicitLocOpBuilder.h"
17#include "mlir/IR/OpDefinition.h"
18#include "mlir/IR/OpImplementation.h"
19#include "mlir/Parser/Parser.h"
20
21#include <gtest/gtest.h>
22
23using namespace mlir;
24
25class ValueShapeRangeTest : public testing::Test {
26protected:
27 void SetUp() override {
28 const char *ir = R"MLIR(
29 func.func @map(%arg : tensor<1xi64>) {
30 %0 = arith.constant dense<[10]> : tensor<1xi64>
31 %1 = arith.addi %arg, %0 : tensor<1xi64>
32 return
33 }
34 )MLIR";
35
36 registry.insert<func::FuncDialect, arith::ArithDialect>();
37 ctx.appendDialectRegistry(registry);
38 module = parseSourceString<ModuleOp>(ir, &ctx);
39 assert(module);
40 mapFn = cast<func::FuncOp>(module->front());
41 }
42
43 // Create ValueShapeRange on the arith.addi operation.
44 ValueShapeRange addiRange() {
45 auto &fnBody = mapFn.getBody();
46 return std::next(fnBody.front().begin())->getOperands();
47 }
48
49 DialectRegistry registry;
50 MLIRContext ctx;
51 OwningOpRef<ModuleOp> module;
52 func::FuncOp mapFn;
53};
54
55TEST_F(ValueShapeRangeTest, ShapesFromValues) {
56 ValueShapeRange range = addiRange();
57
58 EXPECT_FALSE(range.getValueAsShape(0));
59 ASSERT_TRUE(range.getValueAsShape(1));
60 EXPECT_TRUE(range.getValueAsShape(1).hasRank());
61 EXPECT_EQ(range.getValueAsShape(1).getRank(), 1);
62 EXPECT_EQ(range.getValueAsShape(1).getDimSize(0), 10);
63 EXPECT_EQ(range.getShape(1).getRank(), 1);
64 EXPECT_EQ(range.getShape(1).getDimSize(0), 1);
65}
66
67TEST_F(ValueShapeRangeTest, MapValuesToShapes) {
68 ValueShapeRange range = addiRange();
69 ShapedTypeComponents fixed(SmallVector<int64_t>{30});
70 auto mapping = [&](Value val) -> ShapeAdaptor {
71 if (val == mapFn.getArgument(0))
72 return &fixed;
73 return nullptr;
74 };
75 range.setValueToShapeMapping(mapping);
76
77 ASSERT_TRUE(range.getValueAsShape(0));
78 EXPECT_TRUE(range.getValueAsShape(0).hasRank());
79 EXPECT_EQ(range.getValueAsShape(0).getRank(), 1);
80 EXPECT_EQ(range.getValueAsShape(0).getDimSize(0), 30);
81 ASSERT_TRUE(range.getValueAsShape(1));
82 EXPECT_TRUE(range.getValueAsShape(1).hasRank());
83 EXPECT_EQ(range.getValueAsShape(1).getRank(), 1);
84 EXPECT_EQ(range.getValueAsShape(1).getDimSize(0), 10);
85}
86
87TEST_F(ValueShapeRangeTest, SettingShapes) {
88 ShapedTypeComponents shape(SmallVector<int64_t>{10, 20});
89 ValueShapeRange range = addiRange();
90 auto mapping = [&](Value val) -> ShapeAdaptor {
91 if (val == mapFn.getArgument(0))
92 return &shape;
93 return nullptr;
94 };
95 range.setOperandShapeMapping(mapping);
96
97 ASSERT_TRUE(range.getShape(0));
98 EXPECT_EQ(range.getShape(0).getRank(), 2);
99 EXPECT_EQ(range.getShape(0).getDimSize(0), 10);
100 EXPECT_EQ(range.getShape(0).getDimSize(1), 20);
101 ASSERT_TRUE(range.getShape(1));
102 EXPECT_EQ(range.getShape(1).getRank(), 1);
103 EXPECT_EQ(range.getShape(1).getDimSize(0), 1);
104 EXPECT_FALSE(range.getShape(2));
105}
106

source code of mlir/unittests/Interfaces/InferTypeOpInterfaceTest.cpp