SubMulMulAddConvbwdRelubwdBnwrw.cpp

// Copyright © Advanced Micro Devices, Inc., or its affiliates.
// SPDX-License-Identifier:  MIT

#include <iostream>

#include <hipdnn_data_sdk/utilities/Tensor.hpp>
#include <hipdnn_data_sdk/utilities/Workspace.hpp>
#include <hipdnn_frontend.hpp>

#include "../utils.hpp"

int main()
{
    using InputType = float;

    const int64_t n = 4;
    const int64_t c = 64;
    const int64_t h = 16;
    const int64_t w = 16;
    const int64_t k = 32;
    const int64_t r = 3;
    const int64_t s = 3;

    auto buildSubMulMulAddConvbwdRelubwdBnwrwGraph = [=](hipdnnHandle_t handle) {
        auto graph = std::make_shared<hipdnn_frontend::graph::Graph>();
        graph->set_name("sub_mul_mul_add_convbwd_relubwd_bnwrw_graph")
            .set_io_data_type(hipdnn_frontend::getDataTypeEnumFromType<InputType>())
            .set_intermediate_data_type(hipdnn_frontend::getDataTypeEnumFromType<InputType>())
            .set_compute_data_type(hipdnn_frontend::DataType::FLOAT);

        auto xBn = std::make_shared<hipdnn_frontend::graph::TensorAttributes>(
            hipdnn_frontend::graph::Tensor_attributes()
                .set_name("x_bn")
                .set_dim({n, c, h, w})
                .set_stride({c * h * w, 1, c * w, c}));
        auto meanBn = std::make_shared<hipdnn_frontend::graph::TensorAttributes>(
            hipdnn_frontend::graph::Tensor_attributes()
                .set_name("mean_bn")
                .set_dim({1, c, 1, 1})
                .set_stride({c, 1, c, c}));
        auto invStdBn = std::make_shared<hipdnn_frontend::graph::TensorAttributes>(
            hipdnn_frontend::graph::Tensor_attributes()
                .set_name("invstd_bn")
                .set_dim({1, c, 1, 1})
                .set_stride({c, 1, c, c}));
        auto scaleBn = std::make_shared<hipdnn_frontend::graph::TensorAttributes>(
            hipdnn_frontend::graph::Tensor_attributes()
                .set_name("scale_bn")
                .set_dim({1, c, 1, 1})
                .set_stride({c, 1, c, c}));
        auto biasBn = std::make_shared<hipdnn_frontend::graph::TensorAttributes>(
            hipdnn_frontend::graph::Tensor_attributes()
                .set_name("bias_bn")
                .set_dim({1, c, 1, 1})
                .set_stride({c, 1, c, c}));

        auto dy = std::make_shared<hipdnn_frontend::graph::TensorAttributes>(
            hipdnn_frontend::graph::Tensor_attributes()
                .set_name("dy")
                .set_dim({n, k, h, w})
                .set_stride({k * h * w, 1, k * w, k}));
        auto filter = std::make_shared<hipdnn_frontend::graph::TensorAttributes>(
            hipdnn_frontend::graph::Tensor_attributes()
                .set_name("filter")
                .set_dim({k, c, r, s})
                .set_stride({c * r * s, 1, c * s, c}));

        auto subAttrs = hipdnn_frontend::graph::PointwiseAttributes()
                            .set_name("sub_node")
                            .set_mode(hipdnn_frontend::PointwiseMode::SUB);
        auto subOut = graph->pointwise(xBn, meanBn, subAttrs);

        auto mulAttrs0 = hipdnn_frontend::graph::PointwiseAttributes()
                             .set_name("mul0_node")
                             .set_mode(hipdnn_frontend::PointwiseMode::MUL);
        auto mulOut0 = graph->pointwise(subOut, invStdBn, mulAttrs0);

        auto mulAttrs1 = hipdnn_frontend::graph::PointwiseAttributes()
                             .set_name("mul1_node")
                             .set_mode(hipdnn_frontend::PointwiseMode::MUL);
        auto mulOut1 = graph->pointwise(mulOut0, scaleBn, mulAttrs1);

        auto addAttrs = hipdnn_frontend::graph::PointwiseAttributes()
                            .set_name("add_node")
                            .set_mode(hipdnn_frontend::PointwiseMode::ADD);
        auto addOut = graph->pointwise(mulOut1, biasBn, addAttrs);

        auto convAttrs = hipdnn_frontend::graph::ConvDgradAttributes()
                             .set_name("conv_dgrad_node")
                             .set_padding({1, 1})
                             .set_stride({1, 1})
                             .set_dilation({1, 1});
        auto dx = graph->conv_dgrad(dy, filter, convAttrs);

        auto reluBwdAttrs = hipdnn_frontend::graph::PointwiseAttributes()
                                .set_name("relu_bwd_node")
                                .set_mode(hipdnn_frontend::PointwiseMode::RELU_BWD);
        auto dRelu = graph->pointwise(dx, addOut, reluBwdAttrs);
        dRelu->set_output(true);

        auto bnWgradAttrs = hipdnn_frontend::graph::BatchnormBackwardWeightAttributes().set_name(
            "bn_backward_weight_node");
        auto [dscale, dbias, eqScaleDy, eqScaleX, eqBias]
            = graph->dbn_weight(dRelu, xBn, meanBn, invStdBn, scaleBn, bnWgradAttrs);
        dscale->set_output(true);
        dbias->set_output(true);
        eqScaleDy->set_output(true);
        eqScaleX->set_output(true);
        eqBias->set_output(true);

        HIPDNN_FE_CHECK(graph->build(handle));

        return std::make_tuple(graph,
                               xBn,
                               meanBn,
                               invStdBn,
                               scaleBn,
                               biasBn,
                               dy,
                               filter,
                               dRelu,
                               dscale,
                               dbias,
                               eqScaleDy,
                               eqScaleX,
                               eqBias);
    };

    auto backend = hipdnn_frontend::detail::hipdnnBackend();
    if(!backend)
    {
        std::cout << "Creat backend failed. \n";
        return 1;
    }

    hipdnnHandle_t handle;
    HIPDNN_CHECK(backend->create(&handle));

    auto [graph,
          xBn,
          meanBn,
          invStdBn,
          scaleBn,
          biasBn,
          dy,
          filter,
          dRelu,
          dscale,
          dbias,
          eqScaleDy,
          eqScaleX,
          eqBias]
        = buildSubMulMulAddConvbwdRelubwdBnwrwGraph(handle);

    hipdnn_data_sdk::utilities::Tensor<InputType> xBnTensor(xBn->get_dim(), xBn->get_stride());
    hipdnn_data_sdk::utilities::Tensor<InputType> meanBnTensor(meanBn->get_dim(),
                                                               meanBn->get_stride());
    hipdnn_data_sdk::utilities::Tensor<InputType> invStdBnTensor(invStdBn->get_dim(),
                                                                 invStdBn->get_stride());
    hipdnn_data_sdk::utilities::Tensor<InputType> scaleBnTensor(scaleBn->get_dim(),
                                                                scaleBn->get_stride());
    hipdnn_data_sdk::utilities::Tensor<InputType> biasBnTensor(biasBn->get_dim(),
                                                               biasBn->get_stride());
    hipdnn_data_sdk::utilities::Tensor<InputType> dyTensor(dy->get_dim(), dy->get_stride());
    hipdnn_data_sdk::utilities::Tensor<InputType> filterTensor(filter->get_dim(),
                                                               filter->get_stride());
    hipdnn_data_sdk::utilities::Tensor<InputType> dReluTensor(dRelu->get_dim(),
                                                              dRelu->get_stride());
    hipdnn_data_sdk::utilities::Tensor<InputType> dscaleTensor(dscale->get_dim(),
                                                               dscale->get_stride());
    hipdnn_data_sdk::utilities::Tensor<InputType> dbiasTensor(dbias->get_dim(),
                                                              dbias->get_stride());
    hipdnn_data_sdk::utilities::Tensor<InputType> eqScaleDyTensor(eqScaleDy->get_dim(),
                                                                  eqScaleDy->get_stride());
    hipdnn_data_sdk::utilities::Tensor<InputType> eqScaleXTensor(eqScaleX->get_dim(),
                                                                 eqScaleX->get_stride());
    hipdnn_data_sdk::utilities::Tensor<InputType> eqBiasTensor(eqBias->get_dim(),
                                                               eqBias->get_stride());

    std::unordered_map<int64_t, void*> variantPack;
    variantPack[xBn->get_uid()] = xBnTensor.memory().deviceData();
    variantPack[meanBn->get_uid()] = meanBnTensor.memory().deviceData();
    variantPack[invStdBn->get_uid()] = invStdBnTensor.memory().deviceData();
    variantPack[scaleBn->get_uid()] = scaleBnTensor.memory().deviceData();
    variantPack[biasBn->get_uid()] = biasBnTensor.memory().deviceData();
    variantPack[dy->get_uid()] = dyTensor.memory().deviceData();
    variantPack[filter->get_uid()] = filterTensor.memory().deviceData();
    variantPack[dRelu->get_uid()] = dReluTensor.memory().deviceData();
    variantPack[dscale->get_uid()] = dscaleTensor.memory().deviceData();
    variantPack[dbias->get_uid()] = dbiasTensor.memory().deviceData();
    variantPack[eqScaleDy->get_uid()] = eqScaleDyTensor.memory().deviceData();
    variantPack[eqScaleX->get_uid()] = eqScaleXTensor.memory().deviceData();
    variantPack[eqBias->get_uid()] = eqBiasTensor.memory().deviceData();

    int64_t workspaceSize = 0;
    HIPDNN_FE_CHECK(graph->get_workspace_size(workspaceSize));
    const hipdnn_data_sdk::utilities::Workspace workspace(static_cast<size_t>(workspaceSize));

    HIPDNN_FE_CHECK(graph->execute(handle, variantPack, workspace.get()));

    std::cout << "SubMulMulAddConvbwdRelubwdBnwrw graph execution complete. \n";

    HIPDNN_CHECK(backend->destroy(handle));
    return 0;
}