73 lines · cpp
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-exception6//7//===----------------------------------------------------------------------===//8 9#include "mlir/Dialect/Traits.h"10#include "llvm/ADT/SmallVector.h"11#include "gmock/gmock.h"12 13using namespace mlir::OpTrait::util;14 15using llvm::SmallVector;16using ::testing::ElementsAre;17 18TEST(BroadcastShapeTest, CompatibleScalarAndScalar) {19 SmallVector<int64_t, 4> result;20 ASSERT_TRUE(getBroadcastedShape({}, {}, result));21 EXPECT_TRUE(result.empty());22}23 24TEST(BroadcastShapeTest, Compatible0DAnd1DTensor) {25 SmallVector<int64_t, 4> result;26 ASSERT_TRUE(getBroadcastedShape({}, {4}, result));27 EXPECT_THAT(result, ElementsAre(4));28}29 30TEST(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 36TEST(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 42TEST(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 48TEST(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 56TEST(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 62TEST(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 68TEST(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