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
TEST(BroadcastShapeTest,CompatibleScalarAndScalar)18 TEST(BroadcastShapeTest, CompatibleScalarAndScalar) {
19 SmallVector<int64_t, 4> result;
20 ASSERT_TRUE(getBroadcastedShape({}, {}, result));
21 EXPECT_TRUE(result.empty());
22 }
23
TEST(BroadcastShapeTest,Compatible0DAnd1DTensor)24 TEST(BroadcastShapeTest, Compatible0DAnd1DTensor) {
25 SmallVector<int64_t, 4> result;
26 ASSERT_TRUE(getBroadcastedShape({}, {4}, result));
27 EXPECT_THAT(result, ElementsAre(4));
28 }
29
TEST(BroadcastShapeTest,Compatible0DAnd3DTensor)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
TEST(BroadcastShapeTest,CompatibleTensorAndTensor)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
TEST(BroadcastShapeTest,InterleavingOnes)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
TEST(BroadcastShapeTest,InterleavingUnknowns)48 TEST(BroadcastShapeTest, InterleavingUnknowns) {
49 SmallVector<int64_t, 4> result;
50 int64_t dyn = mlir::ShapedType::kDynamic;
51 ASSERT_TRUE(getBroadcastedShape({1, 2, dyn, dyn, dyn}, {dyn, dyn, dyn, 4, 1},
52 result));
53 EXPECT_THAT(result, ElementsAre(dyn, 2, dyn, 4, dyn));
54 }
55
TEST(BroadcastShapeTest,IncompatibleLowDim)56 TEST(BroadcastShapeTest, IncompatibleLowDim) {
57 SmallVector<int64_t, 4> result;
58 ASSERT_FALSE(getBroadcastedShape({4, 3, 5, 5}, {3, 5, 4}, result));
59 EXPECT_TRUE(result.empty());
60 }
61
TEST(BroadcastShapeTest,IncompatibleMiddleDim)62 TEST(BroadcastShapeTest, IncompatibleMiddleDim) {
63 SmallVector<int64_t, 4> result;
64 ASSERT_FALSE(getBroadcastedShape({4, 3, 5, 5}, {3, 7, 5}, result));
65 EXPECT_TRUE(result.empty());
66 }
67
TEST(BroadcastShapeTest,IncompatibleHighDim)68 TEST(BroadcastShapeTest, IncompatibleHighDim) {
69 SmallVector<int64_t, 4> result;
70 ASSERT_FALSE(getBroadcastedShape({3, 5, 5}, {4, 5, 5}, result));
71 EXPECT_TRUE(result.empty());
72 }
73