mlir/unittests/Dialect/BroadcastShapeTest.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-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   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
  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
  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
  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 }