derek.holloway
7502bf64a4
CodeQL / Analyze (actions) (push) Canceled after 0s
CodeQL / Analyze (python) (push) Canceled after 0s
pre-commit / pre-commit (push) Canceled after 0s
Multi-Arch CI / setup (push) Canceled after 0s
Publish build_manylinux_x86_64 images / publish_build_manylinux_x86_64 (push) Canceled after 0s
Setup (Multi-Arch) / setup (push) Canceled after 0s
Multi-Arch CI / Linux::${{ fromJSON(needs.setup.outputs.linux_build_config || '{}').build_variant_label || 'skip' }} (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Publish no_rocm_image_ubi10 image / publish_no_rocm_image_ubi10 (push) Canceled after 0s
Multi-Arch CI / Windows::${{ fromJSON(needs.setup.outputs.windows_build_config || '{}').build_variant_label || 'skip' }} (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Multi-Arch CI / Manifest Diff (push) Canceled after 0s
TheRock Manifest Diff Report / Generate Manifest Diff Report (push) Canceled after 0s
Multi-Arch CI / CI Summary (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04 (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04_ocl_rt images / publish_no_rocm_image_ubuntu24_04_ocl_rt (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04_rocgdb (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04_rocgdb images / publish_no_rocm_image_ubuntu24_04_rocgdb (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04_ocl_rt (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Unit Tests / Unit Tests :: ubuntu-24.04 (push) Canceled after 0s
Unit Tests / Unit Tests :: windows-2022 (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04_media (push) Canceled after 0s
Unit Tests / Unit Tests Summary (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04_openmpi (push) Canceled after 0s
204 lines
9.8 KiB
C++
204 lines
9.8 KiB
C++
//===- ReshapeOpsUtilsTest.cpp - ReshapeOpsUtils unit tests ---------------===//
|
|
//
|
|
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
|
|
// See https://llvm.org/LICENSE.txt for license information.
|
|
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
|
|
//
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
#include "mlir/Dialect/Utils/ReshapeOpsUtils.h"
|
|
#include "mlir/IR/BuiltinTypeInterfaces.h"
|
|
#include "llvm/ADT/STLExtras.h"
|
|
#include "gtest/gtest.h"
|
|
#include <optional>
|
|
|
|
using namespace mlir;
|
|
|
|
/// Helper to make constructing
|
|
/// `std::optional<SmallVector<ReassociationIndices>>` more readable.
|
|
static std::optional<SmallVector<ReassociationIndices>>
|
|
makeOptionalIndices(std::initializer_list<ReassociationIndices> list) {
|
|
return std::optional<SmallVector<ReassociationIndices>>(list);
|
|
}
|
|
|
|
TEST(ReassociationIndicesForCollapse, ScalarTest) {
|
|
EXPECT_EQ(getReassociationIndicesForCollapse({1}, {}),
|
|
makeOptionalIndices({}));
|
|
EXPECT_EQ(getReassociationIndicesForCollapse({1, 1}, {}),
|
|
makeOptionalIndices({}));
|
|
EXPECT_EQ(getReassociationIndicesForCollapse({ShapedType::kDynamic}, {}),
|
|
makeOptionalIndices({}));
|
|
EXPECT_EQ(getReassociationIndicesForCollapse({1, ShapedType::kDynamic,
|
|
ShapedType::kDynamic, 1,
|
|
ShapedType::kDynamic},
|
|
{}),
|
|
makeOptionalIndices({}));
|
|
}
|
|
|
|
TEST(ReassociationIndicesForCollapse, ScalarTestFailure) {
|
|
EXPECT_EQ(getReassociationIndicesForCollapse({}, {}), std::nullopt);
|
|
EXPECT_EQ(getReassociationIndicesForCollapse({}, {1}), std::nullopt);
|
|
EXPECT_EQ(getReassociationIndicesForCollapse({2}, {}), std::nullopt);
|
|
EXPECT_EQ(
|
|
getReassociationIndicesForCollapse({1, 2, ShapedType::kDynamic, 1}, {}),
|
|
std::nullopt);
|
|
}
|
|
|
|
TEST(ReassociationIndicesForCollapse, StaticTest) {
|
|
EXPECT_EQ(getReassociationIndicesForCollapse({10, 20}, {200}),
|
|
makeOptionalIndices({{0, 1}}));
|
|
EXPECT_EQ(getReassociationIndicesForCollapse({10, 20, 30}, {10, 600}),
|
|
makeOptionalIndices({{0}, {1, 2}}));
|
|
EXPECT_EQ(getReassociationIndicesForCollapse({10, 20, 30}, {200, 30}),
|
|
makeOptionalIndices({{0, 1}, {2}}));
|
|
}
|
|
|
|
TEST(ReassociationIndicesForCollapse, StaticTestFailure) {
|
|
// No-op reassociation
|
|
EXPECT_EQ(getReassociationIndicesForCollapse({10, 20}, {10, 20}),
|
|
std::nullopt);
|
|
// Invalid static reassociations
|
|
EXPECT_EQ(getReassociationIndicesForCollapse({10, 20}, {10}), std::nullopt);
|
|
EXPECT_EQ(getReassociationIndicesForCollapse({10, 20, 30}, {200, 300}),
|
|
std::nullopt);
|
|
// Non-collapsing (expanding) reassociation
|
|
EXPECT_EQ(getReassociationIndicesForCollapse({10, 20, 30}, {1, 10, 20, 30}),
|
|
std::nullopt);
|
|
}
|
|
|
|
TEST(ReassociationIndicesForCollapse, StaticTestUnitDims) {
|
|
EXPECT_EQ(getReassociationIndicesForCollapse({10, 1}, {10}),
|
|
makeOptionalIndices({{0, 1}}));
|
|
EXPECT_EQ(getReassociationIndicesForCollapse({1, 20, 30}, {600}),
|
|
makeOptionalIndices({{0, 1, 2}}));
|
|
EXPECT_EQ(getReassociationIndicesForCollapse({1, 1, 1}, {1}),
|
|
makeOptionalIndices({{0, 1, 2}}));
|
|
EXPECT_EQ(getReassociationIndicesForCollapse({1, 1, 1, 1}, {1, 1, 1}),
|
|
makeOptionalIndices({{0}, {1}, {2, 3}}));
|
|
}
|
|
|
|
TEST(ReassociationIndicesForCollapse, DynamicTest) {
|
|
EXPECT_EQ(getReassociationIndicesForCollapse({ShapedType::kDynamic, 1},
|
|
{ShapedType::kDynamic}),
|
|
makeOptionalIndices({{0, 1}}));
|
|
EXPECT_EQ(getReassociationIndicesForCollapse({ShapedType::kDynamic, 1, 1},
|
|
{ShapedType::kDynamic}),
|
|
makeOptionalIndices({{0, 1, 2}}));
|
|
EXPECT_EQ(getReassociationIndicesForCollapse(
|
|
{1, ShapedType::kDynamic, 1, ShapedType::kDynamic, 1},
|
|
{ShapedType::kDynamic, ShapedType::kDynamic}),
|
|
makeOptionalIndices({{0, 1}, {2, 3, 4}}));
|
|
EXPECT_EQ(
|
|
getReassociationIndicesForCollapse(
|
|
{ShapedType::kDynamic, ShapedType::kDynamic}, {ShapedType::kDynamic}),
|
|
makeOptionalIndices({{0, 1}}));
|
|
EXPECT_EQ(getReassociationIndicesForCollapse(
|
|
{1, ShapedType::kDynamic, ShapedType::kDynamic},
|
|
{1, ShapedType::kDynamic}),
|
|
makeOptionalIndices({{0}, {1, 2}}));
|
|
|
|
EXPECT_EQ(getReassociationIndicesForCollapse({ShapedType::kDynamic, 10},
|
|
{ShapedType::kDynamic}),
|
|
makeOptionalIndices({{0, 1}}));
|
|
EXPECT_EQ(getReassociationIndicesForCollapse(
|
|
{1, ShapedType::kDynamic, ShapedType::kDynamic},
|
|
{ShapedType::kDynamic}),
|
|
makeOptionalIndices({{0, 1, 2}}));
|
|
EXPECT_EQ(getReassociationIndicesForCollapse({10, ShapedType::kDynamic},
|
|
{ShapedType::kDynamic}),
|
|
makeOptionalIndices({{0, 1}}));
|
|
EXPECT_EQ(getReassociationIndicesForCollapse(
|
|
{ShapedType::kDynamic, 1, 2, ShapedType::kDynamic, 10},
|
|
{ShapedType::kDynamic, 10}),
|
|
makeOptionalIndices({{0, 1, 2, 3}, {4}}));
|
|
EXPECT_EQ(getReassociationIndicesForCollapse({ShapedType::kDynamic, 10, 20},
|
|
{ShapedType::kDynamic, 20}),
|
|
makeOptionalIndices({{0, 1}, {2}}));
|
|
EXPECT_EQ(getReassociationIndicesForCollapse({10, ShapedType::kDynamic, 20},
|
|
{ShapedType::kDynamic, 20}),
|
|
makeOptionalIndices({{0, 1}, {2}}));
|
|
EXPECT_EQ(getReassociationIndicesForCollapse(
|
|
{ShapedType::kDynamic, 3, 2, 5, 2}, {ShapedType::kDynamic, 20}),
|
|
makeOptionalIndices({{0, 1}, {2, 3, 4}}));
|
|
EXPECT_EQ(getReassociationIndicesForCollapse(
|
|
{10, ShapedType::kDynamic, 20, ShapedType::kDynamic, 1},
|
|
{ShapedType::kDynamic, 20, ShapedType::kDynamic}),
|
|
makeOptionalIndices({{0, 1}, {2}, {3, 4}}));
|
|
EXPECT_EQ(getReassociationIndicesForCollapse({1, ShapedType::kDynamic, 1},
|
|
{ShapedType::kDynamic}),
|
|
makeOptionalIndices({{0, 1, 2}}));
|
|
EXPECT_EQ(getReassociationIndicesForCollapse(
|
|
{ShapedType::kDynamic, ShapedType::kDynamic, 1},
|
|
{ShapedType::kDynamic, ShapedType::kDynamic}),
|
|
makeOptionalIndices({{0}, {1, 2}}));
|
|
EXPECT_EQ(getReassociationIndicesForCollapse(
|
|
{1, ShapedType::kDynamic, ShapedType::kDynamic},
|
|
{ShapedType::kDynamic, ShapedType::kDynamic}),
|
|
makeOptionalIndices({{0, 1}, {2}}));
|
|
EXPECT_EQ(getReassociationIndicesForCollapse(
|
|
{ShapedType::kDynamic, 1, ShapedType::kDynamic},
|
|
{ShapedType::kDynamic, ShapedType::kDynamic}),
|
|
makeOptionalIndices({{0}, {1, 2}}));
|
|
}
|
|
|
|
TEST(ReassociationIndicesForCollapse, DynamicTestFailure) {
|
|
EXPECT_EQ(getReassociationIndicesForCollapse({ShapedType::kDynamic, 10, 20},
|
|
{ShapedType::kDynamic, 10}),
|
|
std::nullopt);
|
|
EXPECT_EQ(getReassociationIndicesForCollapse(
|
|
{ShapedType::kDynamic, 10, ShapedType::kDynamic},
|
|
{ShapedType::kDynamic, ShapedType::kDynamic}),
|
|
std::nullopt);
|
|
EXPECT_EQ(getReassociationIndicesForCollapse(
|
|
{20, ShapedType::kDynamic, 10, ShapedType::kDynamic},
|
|
{ShapedType::kDynamic, ShapedType::kDynamic}),
|
|
std::nullopt);
|
|
EXPECT_EQ(getReassociationIndicesForCollapse(
|
|
{ShapedType::kDynamic, 5, 3, 2, 2}, {ShapedType::kDynamic, 20}),
|
|
std::nullopt);
|
|
EXPECT_EQ(
|
|
getReassociationIndicesForCollapse(
|
|
{ShapedType::kDynamic, ShapedType::kDynamic, ShapedType::kDynamic},
|
|
{ShapedType::kDynamic, ShapedType::kDynamic}),
|
|
std::nullopt);
|
|
EXPECT_EQ(getReassociationIndicesForCollapse(
|
|
{ShapedType::kDynamic, ShapedType::kDynamic, 10, 1,
|
|
ShapedType::kDynamic},
|
|
{ShapedType::kDynamic, ShapedType::kDynamic}),
|
|
std::nullopt);
|
|
EXPECT_EQ(getReassociationIndicesForCollapse(
|
|
{ShapedType::kDynamic, 10, 10, 10, ShapedType::kDynamic},
|
|
{ShapedType::kDynamic, 10, ShapedType::kDynamic}),
|
|
std::nullopt);
|
|
EXPECT_EQ(getReassociationIndicesForCollapse(
|
|
{ShapedType::kDynamic, 10, 10, 10, ShapedType::kDynamic},
|
|
{ShapedType::kDynamic, 2, 2, ShapedType::kDynamic}),
|
|
std::nullopt);
|
|
EXPECT_EQ(getReassociationIndicesForCollapse(
|
|
{ShapedType::kDynamic, 3, 4, 3, ShapedType::kDynamic},
|
|
{ShapedType::kDynamic, 12, ShapedType::kDynamic}),
|
|
std::nullopt);
|
|
EXPECT_EQ(getReassociationIndicesForCollapse(
|
|
{ShapedType::kDynamic, 8, 4, 2, 16, ShapedType::kDynamic},
|
|
{ShapedType::kDynamic, 32, ShapedType::kDynamic}),
|
|
std::nullopt);
|
|
|
|
//===----------------------------------------------------------------------===//
|
|
// TODO: Reassociation for the following examples can be computed, but isn't
|
|
// supported by `getReassociationIndicesForCollapse`.
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
// TODO: Fails because there's no backtracking when some source dimensions
|
|
// remain unmatched at either edge.
|
|
EXPECT_EQ(getReassociationIndicesForCollapse(
|
|
{ShapedType::kDynamic, 10, ShapedType::kDynamic, 10},
|
|
{ShapedType::kDynamic, 10}),
|
|
std::nullopt);
|
|
EXPECT_EQ(getReassociationIndicesForCollapse({1, ShapedType::kDynamic, 2, 2},
|
|
{1, ShapedType::kDynamic, 2}),
|
|
std::nullopt);
|
|
EXPECT_EQ(getReassociationIndicesForCollapse({2, 2, ShapedType::kDynamic, 1},
|
|
{2, ShapedType::kDynamic}),
|
|
std::nullopt);
|
|
}
|