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
416 lines
19 KiB
C++
416 lines
19 KiB
C++
// Copyright © Advanced Micro Devices, Inc., or its affiliates.
|
|
// SPDX-License-Identifier: MIT
|
|
|
|
#include <gtest/gtest.h>
|
|
#include <hip/hip_runtime.h>
|
|
#include <memory>
|
|
#include <unordered_map>
|
|
#include <vector>
|
|
|
|
#include <hipdnn_flatbuffers_sdk/data_objects/graph_generated.h>
|
|
#include <hipdnn_flatbuffers_sdk/data_objects/rmsnorm_backward_attributes_generated.h>
|
|
#include <hipdnn_frontend.hpp>
|
|
#include <hipdnn_test_sdk/constants/RMSNormBackwardConstants.hpp>
|
|
#include <hipdnn_test_sdk/utilities/IntegrationTestFixture.hpp>
|
|
#include <hipdnn_test_sdk/utilities/LoweringTestHelpers.hpp>
|
|
#include <hipdnn_test_sdk/utilities/TestUtilities.hpp>
|
|
#include <hipdnn_test_sdk/utilities/TestableGraph.hpp>
|
|
#include <hipdnn_test_sdk/utilities/ToVec.hpp>
|
|
|
|
using namespace hipdnn_frontend;
|
|
using namespace hipdnn_frontend::graph;
|
|
using namespace hipdnn_tests::constants;
|
|
using hipdnn_tests::buildTensorMap;
|
|
using hipdnn_tests::IntegrationTestFixture;
|
|
using hipdnn_tests::lowerAndDeserialize;
|
|
using hipdnn_tests::TestableGraphLowering;
|
|
using hipdnn_tests::toVec;
|
|
using DataTypeSdk = hipdnn_flatbuffers_sdk::data_objects::DataType;
|
|
using NodeAttrType = hipdnn_flatbuffers_sdk::data_objects::NodeAttributes;
|
|
|
|
namespace
|
|
{
|
|
|
|
// Lowers a frontend graph via build_operation_graph_via_descriptors, then
|
|
// retrieves the serialized graph and deserializes it for verification.
|
|
class IntegrationRMSNormBackwardDescriptorLowering : public IntegrationTestFixture
|
|
{
|
|
};
|
|
|
|
// Lowering round-trip: builds a graph, lowers via descriptors, and verifies
|
|
// the deserialized FlatBuffer attributes match.
|
|
TEST_F(IntegrationRMSNormBackwardDescriptorLowering, RMSNormBackwardRoundTrip)
|
|
{
|
|
auto graph = std::make_shared<TestableGraphLowering>();
|
|
graph->set_name("TestRMSNormBackwardGraph")
|
|
.set_io_data_type(DataType::FLOAT)
|
|
.set_intermediate_data_type(DataType::FLOAT)
|
|
.set_compute_data_type(DataType::FLOAT);
|
|
|
|
auto dy = std::make_shared<TensorAttributes>();
|
|
dy->set_uid(K_RMSNORMBACKWARD_TENSOR_DY_UID).set_name("DY").set_data_type(DataType::FLOAT);
|
|
dy->set_dim(toVec(K_RMSNORMBACKWARD_TENSOR_X_DIMS))
|
|
.set_stride(toVec(K_RMSNORMBACKWARD_TENSOR_X_STRIDES));
|
|
|
|
auto x = std::make_shared<TensorAttributes>();
|
|
x->set_uid(K_RMSNORMBACKWARD_TENSOR_X_UID).set_name("X").set_data_type(DataType::FLOAT);
|
|
x->set_dim(toVec(K_RMSNORMBACKWARD_TENSOR_X_DIMS))
|
|
.set_stride(toVec(K_RMSNORMBACKWARD_TENSOR_X_STRIDES));
|
|
|
|
auto scale = std::make_shared<TensorAttributes>();
|
|
scale->set_uid(K_RMSNORMBACKWARD_TENSOR_SCALE_UID)
|
|
.set_name("Scale")
|
|
.set_data_type(DataType::FLOAT);
|
|
scale->set_dim(toVec(K_RMSNORMBACKWARD_TENSOR_SCALE_DIMS))
|
|
.set_stride(toVec(K_RMSNORMBACKWARD_TENSOR_SCALE_STRIDES));
|
|
|
|
auto invRms = std::make_shared<TensorAttributes>();
|
|
invRms->set_uid(K_RMSNORMBACKWARD_TENSOR_INV_RMS_UID)
|
|
.set_name("inv_rms")
|
|
.set_data_type(DataType::FLOAT);
|
|
invRms->set_dim(toVec(K_RMSNORMBACKWARD_TENSOR_INV_RMS_DIMS))
|
|
.set_stride(toVec(K_RMSNORMBACKWARD_TENSOR_INV_RMS_STRIDES));
|
|
|
|
RMSNormBackwardAttributes attrs;
|
|
attrs.set_name("rmsnorm_bwd_op");
|
|
|
|
auto outputs = graph->rmsnorm_backward(dy, x, scale, invRms, std::move(attrs));
|
|
const auto& dx = outputs[0];
|
|
const auto& dScale = outputs[1];
|
|
const auto& dBias = outputs[2];
|
|
|
|
dx->set_uid(K_RMSNORMBACKWARD_TENSOR_DX_UID).set_output(true).set_name("DX");
|
|
dScale->set_uid(K_RMSNORMBACKWARD_TENSOR_DSCALE_UID).set_output(true).set_name("DScale");
|
|
EXPECT_EQ(dBias, nullptr);
|
|
|
|
auto graphT = lowerAndDeserialize(*graph, _handle);
|
|
|
|
// -- Verify graph-level attributes --
|
|
EXPECT_EQ(graphT.compute_data_type, DataTypeSdk::FLOAT);
|
|
EXPECT_EQ(graphT.intermediate_data_type, DataTypeSdk::FLOAT);
|
|
EXPECT_EQ(graphT.io_data_type, DataTypeSdk::FLOAT);
|
|
|
|
// Verify tensors: dy, x, scale, inv_rms, dx, dscale = 6 tensors
|
|
ASSERT_EQ(graphT.tensors.size(), 6u);
|
|
|
|
auto tensorMap = buildTensorMap(graphT);
|
|
|
|
// Verify DY tensor
|
|
ASSERT_NE(tensorMap.count(K_RMSNORMBACKWARD_TENSOR_DY_UID), 0u);
|
|
auto* dyT = tensorMap[K_RMSNORMBACKWARD_TENSOR_DY_UID];
|
|
EXPECT_EQ(dyT->name, "DY");
|
|
EXPECT_EQ(dyT->dims, toVec(K_RMSNORMBACKWARD_TENSOR_DY_DIMS));
|
|
EXPECT_EQ(dyT->strides, toVec(K_RMSNORMBACKWARD_TENSOR_DY_STRIDES));
|
|
EXPECT_EQ(dyT->data_type, DataTypeSdk::FLOAT);
|
|
EXPECT_FALSE(dyT->virtual_);
|
|
|
|
// Verify X tensor
|
|
ASSERT_NE(tensorMap.count(K_RMSNORMBACKWARD_TENSOR_X_UID), 0u);
|
|
auto* xT = tensorMap[K_RMSNORMBACKWARD_TENSOR_X_UID];
|
|
EXPECT_EQ(xT->name, "X");
|
|
EXPECT_EQ(xT->dims, toVec(K_RMSNORMBACKWARD_TENSOR_X_DIMS));
|
|
EXPECT_EQ(xT->strides, toVec(K_RMSNORMBACKWARD_TENSOR_X_STRIDES));
|
|
EXPECT_EQ(xT->data_type, DataTypeSdk::FLOAT);
|
|
EXPECT_FALSE(xT->virtual_);
|
|
|
|
// Verify Scale tensor
|
|
ASSERT_NE(tensorMap.count(K_RMSNORMBACKWARD_TENSOR_SCALE_UID), 0u);
|
|
auto* scaleT = tensorMap[K_RMSNORMBACKWARD_TENSOR_SCALE_UID];
|
|
EXPECT_EQ(scaleT->name, "Scale");
|
|
EXPECT_EQ(scaleT->dims, toVec(K_RMSNORMBACKWARD_TENSOR_SCALE_DIMS));
|
|
EXPECT_EQ(scaleT->strides, toVec(K_RMSNORMBACKWARD_TENSOR_SCALE_STRIDES));
|
|
EXPECT_EQ(scaleT->data_type, DataTypeSdk::FLOAT);
|
|
EXPECT_FALSE(scaleT->virtual_);
|
|
|
|
// Verify inv rms tensor
|
|
ASSERT_NE(tensorMap.count(K_RMSNORMBACKWARD_TENSOR_INV_RMS_UID), 0u);
|
|
auto* invRMST = tensorMap[K_RMSNORMBACKWARD_TENSOR_INV_RMS_UID];
|
|
EXPECT_EQ(invRMST->name, "inv_rms");
|
|
EXPECT_EQ(invRMST->dims, toVec(K_RMSNORMBACKWARD_TENSOR_INV_RMS_DIMS));
|
|
EXPECT_EQ(invRMST->strides, toVec(K_RMSNORMBACKWARD_TENSOR_INV_RMS_STRIDES));
|
|
EXPECT_EQ(invRMST->data_type, DataTypeSdk::FLOAT);
|
|
EXPECT_FALSE(invRMST->virtual_);
|
|
|
|
// Verify DX tensor
|
|
ASSERT_NE(tensorMap.count(K_RMSNORMBACKWARD_TENSOR_DX_UID), 0u);
|
|
auto* dxT = tensorMap[K_RMSNORMBACKWARD_TENSOR_DX_UID];
|
|
EXPECT_EQ(dxT->name, "DX");
|
|
EXPECT_EQ(dxT->dims, toVec(K_RMSNORMBACKWARD_TENSOR_DX_DIMS));
|
|
EXPECT_EQ(dxT->strides, toVec(K_RMSNORMBACKWARD_TENSOR_DX_STRIDES));
|
|
EXPECT_EQ(dxT->data_type, DataTypeSdk::FLOAT);
|
|
EXPECT_FALSE(dxT->virtual_);
|
|
|
|
// Verify Dscale tensor
|
|
ASSERT_NE(tensorMap.count(K_RMSNORMBACKWARD_TENSOR_DSCALE_UID), 0u);
|
|
auto* dscaleT = tensorMap[K_RMSNORMBACKWARD_TENSOR_DSCALE_UID];
|
|
EXPECT_EQ(dscaleT->name, "DScale");
|
|
EXPECT_EQ(tensorMap[K_RMSNORMBACKWARD_TENSOR_DSCALE_UID]->dims,
|
|
toVec(K_RMSNORMBACKWARD_TENSOR_DSCALE_DIMS));
|
|
EXPECT_EQ(dscaleT->strides, toVec(K_RMSNORMBACKWARD_TENSOR_DSCALE_STRIDES));
|
|
EXPECT_EQ(dscaleT->data_type, DataTypeSdk::FLOAT);
|
|
EXPECT_FALSE(dscaleT->virtual_);
|
|
|
|
// Verify operation node
|
|
ASSERT_EQ(graphT.nodes.size(), 1u);
|
|
auto& node = graphT.nodes[0];
|
|
EXPECT_EQ(node->compute_data_type, DataTypeSdk::FLOAT);
|
|
EXPECT_EQ(node->attributes.type, NodeAttrType::RMSNormBackwardAttributes);
|
|
|
|
auto* opNode = node->attributes.AsRMSNormBackwardAttributes();
|
|
ASSERT_NE(opNode, nullptr);
|
|
|
|
// Verify required tensor UIDs
|
|
EXPECT_EQ(opNode->dy_tensor_uid, K_RMSNORMBACKWARD_TENSOR_DY_UID);
|
|
EXPECT_EQ(opNode->x_tensor_uid, K_RMSNORMBACKWARD_TENSOR_X_UID);
|
|
EXPECT_EQ(opNode->scale_tensor_uid, K_RMSNORMBACKWARD_TENSOR_SCALE_UID);
|
|
EXPECT_EQ(opNode->inv_rms_tensor_uid, K_RMSNORMBACKWARD_TENSOR_INV_RMS_UID);
|
|
EXPECT_EQ(opNode->dx_tensor_uid, K_RMSNORMBACKWARD_TENSOR_DX_UID);
|
|
EXPECT_EQ(opNode->dscale_tensor_uid, K_RMSNORMBACKWARD_TENSOR_DSCALE_UID);
|
|
|
|
// Verify operation name preserved through lowering
|
|
EXPECT_EQ(node->name, "rmsnorm_bwd_op");
|
|
}
|
|
|
|
// Roundtrip with optional dbias tensor set, verifying it appears in the serialized graph.
|
|
TEST_F(IntegrationRMSNormBackwardDescriptorLowering, RMSNormBackwardWithDBiasRoundTrip)
|
|
{
|
|
auto graph = std::make_shared<TestableGraphLowering>();
|
|
graph->set_name("TestRMSNormBackwardGraph")
|
|
.set_io_data_type(DataType::FLOAT)
|
|
.set_intermediate_data_type(DataType::FLOAT)
|
|
.set_compute_data_type(DataType::FLOAT);
|
|
|
|
auto dy = std::make_shared<TensorAttributes>();
|
|
dy->set_uid(K_RMSNORMBACKWARD_TENSOR_DY_UID).set_name("DY").set_data_type(DataType::FLOAT);
|
|
dy->set_dim(toVec(K_RMSNORMBACKWARD_TENSOR_X_DIMS))
|
|
.set_stride(toVec(K_RMSNORMBACKWARD_TENSOR_X_STRIDES));
|
|
|
|
auto x = std::make_shared<TensorAttributes>();
|
|
x->set_uid(K_RMSNORMBACKWARD_TENSOR_X_UID).set_name("X").set_data_type(DataType::FLOAT);
|
|
x->set_dim(toVec(K_RMSNORMBACKWARD_TENSOR_X_DIMS))
|
|
.set_stride(toVec(K_RMSNORMBACKWARD_TENSOR_X_STRIDES));
|
|
|
|
auto scale = std::make_shared<TensorAttributes>();
|
|
scale->set_uid(K_RMSNORMBACKWARD_TENSOR_SCALE_UID)
|
|
.set_name("Scale")
|
|
.set_data_type(DataType::FLOAT);
|
|
scale->set_dim(toVec(K_RMSNORMBACKWARD_TENSOR_SCALE_DIMS))
|
|
.set_stride(toVec(K_RMSNORMBACKWARD_TENSOR_SCALE_STRIDES));
|
|
|
|
auto invRms = std::make_shared<TensorAttributes>();
|
|
invRms->set_uid(K_RMSNORMBACKWARD_TENSOR_INV_RMS_UID)
|
|
.set_name("inv_rms")
|
|
.set_data_type(DataType::FLOAT);
|
|
invRms->set_dim(toVec(K_RMSNORMBACKWARD_TENSOR_INV_RMS_DIMS))
|
|
.set_stride(toVec(K_RMSNORMBACKWARD_TENSOR_INV_RMS_STRIDES));
|
|
|
|
RMSNormBackwardAttributes attrs;
|
|
attrs.set_name("rmsnorm_bwd_op");
|
|
attrs.set_compute_dbias(true);
|
|
|
|
auto outputs = graph->rmsnorm_backward(dy, x, scale, invRms, std::move(attrs));
|
|
const auto& dx = outputs[0];
|
|
const auto& dScale = outputs[1];
|
|
const auto& dBias = outputs[2];
|
|
|
|
dx->set_uid(K_RMSNORMBACKWARD_TENSOR_DX_UID).set_output(true).set_name("DX");
|
|
dScale->set_uid(K_RMSNORMBACKWARD_TENSOR_DSCALE_UID).set_output(true).set_name("DScale");
|
|
dBias->set_uid(K_RMSNORMBACKWARD_TENSOR_DBIAS_UID).set_output(true).set_name("DBias");
|
|
|
|
auto graphT = lowerAndDeserialize(*graph, _handle);
|
|
|
|
// -- Verify graph-level attributes --
|
|
EXPECT_EQ(graphT.compute_data_type, DataTypeSdk::FLOAT);
|
|
EXPECT_EQ(graphT.intermediate_data_type, DataTypeSdk::FLOAT);
|
|
EXPECT_EQ(graphT.io_data_type, DataTypeSdk::FLOAT);
|
|
|
|
// Verify tensors: dy, x, scale, inv_rms, dx, dscale, dbias = 7 tensors
|
|
ASSERT_EQ(graphT.tensors.size(), 7u);
|
|
|
|
auto tensorMap = buildTensorMap(graphT);
|
|
|
|
// Verify DY tensor
|
|
ASSERT_NE(tensorMap.count(K_RMSNORMBACKWARD_TENSOR_DY_UID), 0u);
|
|
auto* dyT = tensorMap[K_RMSNORMBACKWARD_TENSOR_DY_UID];
|
|
EXPECT_EQ(dyT->name, "DY");
|
|
EXPECT_EQ(dyT->dims, toVec(K_RMSNORMBACKWARD_TENSOR_DY_DIMS));
|
|
EXPECT_EQ(dyT->strides, toVec(K_RMSNORMBACKWARD_TENSOR_DY_STRIDES));
|
|
EXPECT_EQ(dyT->data_type, DataTypeSdk::FLOAT);
|
|
EXPECT_FALSE(dyT->virtual_);
|
|
|
|
// Verify X tensor
|
|
ASSERT_NE(tensorMap.count(K_RMSNORMBACKWARD_TENSOR_X_UID), 0u);
|
|
auto* xT = tensorMap[K_RMSNORMBACKWARD_TENSOR_X_UID];
|
|
EXPECT_EQ(xT->name, "X");
|
|
EXPECT_EQ(xT->dims, toVec(K_RMSNORMBACKWARD_TENSOR_X_DIMS));
|
|
EXPECT_EQ(xT->strides, toVec(K_RMSNORMBACKWARD_TENSOR_X_STRIDES));
|
|
EXPECT_EQ(xT->data_type, DataTypeSdk::FLOAT);
|
|
EXPECT_FALSE(xT->virtual_);
|
|
|
|
// Verify Scale tensor
|
|
ASSERT_NE(tensorMap.count(K_RMSNORMBACKWARD_TENSOR_SCALE_UID), 0u);
|
|
auto* scaleT = tensorMap[K_RMSNORMBACKWARD_TENSOR_SCALE_UID];
|
|
EXPECT_EQ(scaleT->name, "Scale");
|
|
EXPECT_EQ(scaleT->dims, toVec(K_RMSNORMBACKWARD_TENSOR_SCALE_DIMS));
|
|
EXPECT_EQ(scaleT->strides, toVec(K_RMSNORMBACKWARD_TENSOR_SCALE_STRIDES));
|
|
EXPECT_EQ(scaleT->data_type, DataTypeSdk::FLOAT);
|
|
EXPECT_FALSE(scaleT->virtual_);
|
|
|
|
// Verify inv rms tensor
|
|
ASSERT_NE(tensorMap.count(K_RMSNORMBACKWARD_TENSOR_INV_RMS_UID), 0u);
|
|
auto* invRMST = tensorMap[K_RMSNORMBACKWARD_TENSOR_INV_RMS_UID];
|
|
EXPECT_EQ(invRMST->name, "inv_rms");
|
|
EXPECT_EQ(invRMST->dims, toVec(K_RMSNORMBACKWARD_TENSOR_INV_RMS_DIMS));
|
|
EXPECT_EQ(invRMST->strides, toVec(K_RMSNORMBACKWARD_TENSOR_INV_RMS_STRIDES));
|
|
EXPECT_EQ(invRMST->data_type, DataTypeSdk::FLOAT);
|
|
EXPECT_FALSE(invRMST->virtual_);
|
|
|
|
// Verify DX tensor
|
|
ASSERT_NE(tensorMap.count(K_RMSNORMBACKWARD_TENSOR_DX_UID), 0u);
|
|
auto* dxT = tensorMap[K_RMSNORMBACKWARD_TENSOR_DX_UID];
|
|
EXPECT_EQ(dxT->name, "DX");
|
|
EXPECT_EQ(dxT->dims, toVec(K_RMSNORMBACKWARD_TENSOR_DX_DIMS));
|
|
EXPECT_EQ(dxT->strides, toVec(K_RMSNORMBACKWARD_TENSOR_DX_STRIDES));
|
|
EXPECT_EQ(dxT->data_type, DataTypeSdk::FLOAT);
|
|
EXPECT_FALSE(dxT->virtual_);
|
|
|
|
// Verify Dscale tensor
|
|
ASSERT_NE(tensorMap.count(K_RMSNORMBACKWARD_TENSOR_DSCALE_UID), 0u);
|
|
auto* dscaleT = tensorMap[K_RMSNORMBACKWARD_TENSOR_DSCALE_UID];
|
|
EXPECT_EQ(dscaleT->name, "DScale");
|
|
EXPECT_EQ(tensorMap[K_RMSNORMBACKWARD_TENSOR_DSCALE_UID]->dims,
|
|
toVec(K_RMSNORMBACKWARD_TENSOR_DSCALE_DIMS));
|
|
EXPECT_EQ(dscaleT->strides, toVec(K_RMSNORMBACKWARD_TENSOR_DSCALE_STRIDES));
|
|
EXPECT_EQ(dscaleT->data_type, DataTypeSdk::FLOAT);
|
|
EXPECT_FALSE(dscaleT->virtual_);
|
|
|
|
// Verify Dbias tensor
|
|
ASSERT_NE(tensorMap.count(K_RMSNORMBACKWARD_TENSOR_DBIAS_UID), 0u);
|
|
auto* dbiasT = tensorMap[K_RMSNORMBACKWARD_TENSOR_DBIAS_UID];
|
|
EXPECT_EQ(dbiasT->name, "DBias");
|
|
EXPECT_EQ(dbiasT->dims, toVec(K_RMSNORMBACKWARD_TENSOR_DBIAS_DIMS));
|
|
EXPECT_EQ(dbiasT->strides, toVec(K_RMSNORMBACKWARD_TENSOR_DBIAS_STRIDES));
|
|
EXPECT_EQ(dbiasT->data_type, DataTypeSdk::FLOAT);
|
|
EXPECT_FALSE(dbiasT->virtual_);
|
|
|
|
// Verify operation node
|
|
ASSERT_EQ(graphT.nodes.size(), 1u);
|
|
auto& node = graphT.nodes[0];
|
|
EXPECT_EQ(node->compute_data_type, DataTypeSdk::FLOAT);
|
|
EXPECT_EQ(node->attributes.type, NodeAttrType::RMSNormBackwardAttributes);
|
|
|
|
auto* opNode = node->attributes.AsRMSNormBackwardAttributes();
|
|
ASSERT_NE(opNode, nullptr);
|
|
|
|
// Verify required tensor UIDs
|
|
EXPECT_EQ(opNode->dy_tensor_uid, K_RMSNORMBACKWARD_TENSOR_DY_UID);
|
|
EXPECT_EQ(opNode->x_tensor_uid, K_RMSNORMBACKWARD_TENSOR_X_UID);
|
|
EXPECT_EQ(opNode->scale_tensor_uid, K_RMSNORMBACKWARD_TENSOR_SCALE_UID);
|
|
EXPECT_EQ(opNode->inv_rms_tensor_uid, K_RMSNORMBACKWARD_TENSOR_INV_RMS_UID);
|
|
EXPECT_EQ(opNode->dx_tensor_uid, K_RMSNORMBACKWARD_TENSOR_DX_UID);
|
|
EXPECT_EQ(opNode->dscale_tensor_uid, K_RMSNORMBACKWARD_TENSOR_DSCALE_UID);
|
|
EXPECT_EQ(opNode->dbias_tensor_uid, K_RMSNORMBACKWARD_TENSOR_DBIAS_UID);
|
|
|
|
// Verify operation name preserved through lowering
|
|
EXPECT_EQ(node->name, "rmsnorm_bwd_op");
|
|
}
|
|
|
|
// Verifies that tensor UIDs auto-assigned by the frontend are preserved
|
|
// through the lowering round-trip.
|
|
TEST_F(IntegrationRMSNormBackwardDescriptorLowering, AutoAssignedUidsPreservedInRoundTrip)
|
|
{
|
|
auto graph = std::make_shared<TestableGraphLowering>();
|
|
graph->set_name("AutoUidRMSNormBwdGraph")
|
|
.set_io_data_type(DataType::FLOAT)
|
|
.set_intermediate_data_type(DataType::FLOAT)
|
|
.set_compute_data_type(DataType::FLOAT);
|
|
|
|
auto dy = std::make_shared<TensorAttributes>();
|
|
dy->set_name("DY").set_data_type(DataType::FLOAT);
|
|
dy->set_dim(toVec(K_RMSNORMBACKWARD_TENSOR_X_DIMS))
|
|
.set_stride(toVec(K_RMSNORMBACKWARD_TENSOR_X_STRIDES));
|
|
|
|
auto x = std::make_shared<TensorAttributes>();
|
|
x->set_name("X").set_data_type(DataType::FLOAT);
|
|
x->set_dim(toVec(K_RMSNORMBACKWARD_TENSOR_X_DIMS))
|
|
.set_stride(toVec(K_RMSNORMBACKWARD_TENSOR_X_STRIDES));
|
|
|
|
auto scale = std::make_shared<TensorAttributes>();
|
|
scale->set_name("Scale").set_data_type(DataType::FLOAT);
|
|
scale->set_dim(toVec(K_RMSNORMBACKWARD_TENSOR_SCALE_DIMS))
|
|
.set_stride(toVec(K_RMSNORMBACKWARD_TENSOR_SCALE_STRIDES));
|
|
|
|
auto invRms = std::make_shared<TensorAttributes>();
|
|
invRms->set_uid(K_RMSNORMBACKWARD_TENSOR_INV_RMS_UID)
|
|
.set_name("inv_rms")
|
|
.set_data_type(DataType::FLOAT);
|
|
invRms->set_dim(toVec(K_RMSNORMBACKWARD_TENSOR_INV_RMS_DIMS))
|
|
.set_stride(toVec(K_RMSNORMBACKWARD_TENSOR_INV_RMS_STRIDES));
|
|
|
|
RMSNormBackwardAttributes attrs;
|
|
attrs.set_name("rmsnorm_bwd_op");
|
|
attrs.set_compute_dbias(true);
|
|
|
|
auto outputs = graph->rmsnorm_backward(dy, x, scale, invRms, std::move(attrs));
|
|
const auto& dx = outputs[0];
|
|
const auto& dScale = outputs[1];
|
|
const auto& dBias = outputs[2];
|
|
|
|
dx->set_uid(K_RMSNORMBACKWARD_TENSOR_DX_UID).set_output(true);
|
|
dScale->set_uid(K_RMSNORMBACKWARD_TENSOR_DSCALE_UID).set_output(true);
|
|
dBias->set_uid(K_RMSNORMBACKWARD_TENSOR_DBIAS_UID).set_output(true);
|
|
|
|
auto graphT = lowerAndDeserialize(*graph, _handle);
|
|
|
|
// All tensors should have been auto-assigned unique UIDs
|
|
ASSERT_EQ(graphT.tensors.size(), 7u);
|
|
std::unordered_set<int64_t> uids;
|
|
for(const auto& t : graphT.tensors)
|
|
{
|
|
uids.insert(t->uid);
|
|
}
|
|
EXPECT_EQ(uids.size(), 7u)
|
|
<< "Tensor UIDs are not unique"; // NOLINT(readability-implicit-bool-conversion)
|
|
|
|
// The batchnorm inference operation should reference the auto-assigned UIDs
|
|
ASSERT_EQ(graphT.nodes.size(), 1u);
|
|
auto* rmsBwd = graphT.nodes[0]->attributes.AsRMSNormBackwardAttributes();
|
|
ASSERT_NE(rmsBwd, nullptr);
|
|
|
|
// Tensor UIDs in the node should match tensors in the graph
|
|
EXPECT_TRUE(uids.count(rmsBwd->dy_tensor_uid) > 0)
|
|
<< "DY tensor UID " << rmsBwd->dy_tensor_uid
|
|
<< " not found in graph tensors"; // NOLINT(readability-implicit-bool-conversion)
|
|
EXPECT_TRUE(uids.count(rmsBwd->x_tensor_uid) > 0)
|
|
<< "X tensor UID " << rmsBwd->x_tensor_uid
|
|
<< " not found in graph tensors"; // NOLINT(readability-implicit-bool-conversion)
|
|
EXPECT_TRUE(uids.count(rmsBwd->scale_tensor_uid) > 0)
|
|
<< "Scale tensor UID " << rmsBwd->scale_tensor_uid
|
|
<< " not found in graph tensors"; // NOLINT(readability-implicit-bool-conversion)
|
|
EXPECT_TRUE(uids.count(rmsBwd->inv_rms_tensor_uid) > 0)
|
|
<< "INV RMS tensor UID " << rmsBwd->inv_rms_tensor_uid
|
|
<< " not found in graph tensors"; // NOLINT(readability-implicit-bool-conversion)
|
|
EXPECT_TRUE(uids.count(rmsBwd->dx_tensor_uid) > 0)
|
|
<< "DX tensor UID " << rmsBwd->dx_tensor_uid // NOLINT(readability-implicit-bool-conversion)
|
|
<< " not found in graph tensors"; // NOLINT(readability-implicit-bool-conversion)
|
|
EXPECT_TRUE(uids.count(rmsBwd->dscale_tensor_uid) > 0)
|
|
<< "DScale tensor UID " << rmsBwd->dscale_tensor_uid
|
|
<< " not found in graph tensors"; // NOLINT(readability-implicit-bool-conversion)
|
|
EXPECT_TRUE(rmsBwd->dbias_tensor_uid.has_value());
|
|
EXPECT_TRUE(uids.count(rmsBwd->dbias_tensor_uid.value()) > 0)
|
|
<< "DBias tensor UID " << rmsBwd->dbias_tensor_uid.value()
|
|
<< " not found in graph tensors"; // NOLINT(readability-implicit-bool-conversion)
|
|
|
|
// All six tensor UIDs referenced by the node should be distinct
|
|
const std::unordered_set<int64_t> nodeUids = {rmsBwd->dy_tensor_uid,
|
|
rmsBwd->x_tensor_uid,
|
|
rmsBwd->scale_tensor_uid,
|
|
rmsBwd->inv_rms_tensor_uid,
|
|
rmsBwd->dx_tensor_uid,
|
|
rmsBwd->dscale_tensor_uid,
|
|
rmsBwd->dbias_tensor_uid.value()};
|
|
EXPECT_EQ(nodeUids.size(), 7u)
|
|
<< "RMSNorm backward node tensor UIDs are not distinct"; // NOLINT(readability-implicit-bool-conversion)
|
|
}
|
|
} // namespace
|