1 //===- BroadcastShapeTest.cpp - broadcasting shape unit tests -------------===//
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/Dialect/Traits.h"
10 #include "llvm/ADT/SmallVector.h"
11 #include "gmock/gmock.h"
12 
13 using namespace mlir::OpTrait::util;
14 
15 using llvm::SmallVector;
16 using ::testing::ElementsAre;
17 
18 TEST(BroadcastShapeTest, CompatibleScalarAndScalar) {
19   SmallVector<int64_t, 4> result;
20   ASSERT_TRUE(getBroadcastedShape({}, {}, result));
21   EXPECT_TRUE(result.empty());
22 }
23 
24 TEST(BroadcastShapeTest, Compatible0DAnd1DTensor) {
25   SmallVector<int64_t, 4> result;
26   ASSERT_TRUE(getBroadcastedShape({}, {4}, result));
27   EXPECT_THAT(result, ElementsAre(4));
28 }
29 
30 TEST(BroadcastShapeTest, Compatible0DAnd3DTensor) {
31   SmallVector<int64_t, 4> result;
32   ASSERT_TRUE(getBroadcastedShape({}, {3, 5, 4}, result));
33   EXPECT_THAT(result, ElementsAre(3, 5, 4));
34 }
35 
36 TEST(BroadcastShapeTest, CompatibleTensorAndTensor) {
37   SmallVector<int64_t, 4> result;
38   ASSERT_TRUE(getBroadcastedShape({1, 7, 8, 9}, {8, 9}, result));
39   EXPECT_THAT(result, ElementsAre(1, 7, 8, 9));
40 }
41 
42 TEST(BroadcastShapeTest, InterleavingOnes) {
43   SmallVector<int64_t, 4> result;
44   ASSERT_TRUE(getBroadcastedShape({8, 1, 2, 1, 4}, {5, 1, 7, 1}, result));
45   EXPECT_THAT(result, ElementsAre(8, 5, 2, 7, 4));
46 }
47 
48 TEST(BroadcastShapeTest, InterleavingUnknowns) {
49   SmallVector<int64_t, 4> result;
50   ASSERT_TRUE(
51       getBroadcastedShape({1, 2, -1, -1, -1}, {-1, -1, -1, 4, 1}, result));
52   EXPECT_THAT(result, ElementsAre(-1, 2, -1, 4, -1));
53 }
54 
55 TEST(BroadcastShapeTest, IncompatibleLowDim) {
56   SmallVector<int64_t, 4> result;
57   ASSERT_FALSE(getBroadcastedShape({4, 3, 5, 5}, {3, 5, 4}, result));
58   EXPECT_TRUE(result.empty());
59 }
60 
61 TEST(BroadcastShapeTest, IncompatibleMiddleDim) {
62   SmallVector<int64_t, 4> result;
63   ASSERT_FALSE(getBroadcastedShape({4, 3, 5, 5}, {3, 7, 5}, result));
64   EXPECT_TRUE(result.empty());
65 }
66 
67 TEST(BroadcastShapeTest, IncompatibleHighDim) {
68   SmallVector<int64_t, 4> result;
69   ASSERT_FALSE(getBroadcastedShape({3, 5, 5}, {4, 5, 5}, result));
70   EXPECT_TRUE(result.empty());
71 }
72