run_permute_element_example.inc 2.39 KB
Newer Older
Po-Yen, Chen's avatar
Po-Yen, Chen committed
1
2
3
4
5
// SPDX-License-Identifier: MIT
// Copyright (c) 2018-2022, Advanced Micro Devices, Inc. All rights reserved.

#pragma once

Po-Yen, Chen's avatar
Po-Yen, Chen committed
6
bool run_permute_element(const Problem& problem)
Po-Yen, Chen's avatar
Po-Yen, Chen committed
7
{
8
9
    const auto& input_shape = problem.shape;
    const auto& input_axes  = problem.axes;
10

11
    const auto output_shape = transpose(input_shape, input_axes);
Po-Yen, Chen's avatar
Po-Yen, Chen committed
12

13
14
    Tensor<InDataType> input_tensor(input_shape);
    Tensor<OutDataType> output_tensor(output_shape);
Po-Yen, Chen's avatar
Po-Yen, Chen committed
15

16
17
18
19
    ck::utils::FillUniformDistribution<InDataType>{-1.f, 1.f}(input_tensor);

    DeviceMem input_device_buf(input_tensor.GetElementSpaceSizeInBytes());
    DeviceMem output_device_buf(output_tensor.GetElementSpaceSizeInBytes());
Po-Yen, Chen's avatar
Po-Yen, Chen committed
20

21
    using std::data;
22
    input_device_buf.ToDevice(data(input_tensor));
Po-Yen, Chen's avatar
Po-Yen, Chen committed
23

24
25
    static_assert(std::is_default_constructible_v<DevicePermuteInstance>);

26
    auto permute  = DevicePermuteInstance{};
27
28
29
30
    auto argument = permute.MakeArgument(to_array(input_shape),
                                         to_array(input_tensor.GetStrides()),
                                         to_array(output_shape),
                                         to_array(output_tensor.GetStrides()),
Po-Yen, Chen's avatar
Po-Yen, Chen committed
31
32
                                         input_device_buf.GetDeviceBuffer(),
                                         output_device_buf.GetDeviceBuffer(),
33
                                         PassThrough{});
Po-Yen, Chen's avatar
Po-Yen, Chen committed
34

35
    if(!permute.IsSupportedArgument(argument))
Po-Yen, Chen's avatar
Po-Yen, Chen committed
36
    {
37
38
39
        std::cerr << "The runtime parameters seems not supported by the device instance, exiting!"
                  << std::endl;
        return false;
Po-Yen, Chen's avatar
Po-Yen, Chen committed
40
41
    };

42
    auto invoker   = permute.MakeInvoker();
Po-Yen, Chen's avatar
Po-Yen, Chen committed
43
    float ave_time = invoker.Run(argument, StreamConfig{nullptr, true});
Po-Yen, Chen's avatar
Po-Yen, Chen committed
44
45
46

    std::cout << "Perf: " << ave_time << " ms" << std::endl;

47
    output_device_buf.FromDevice(data(output_tensor));
Po-Yen, Chen's avatar
Po-Yen, Chen committed
48

49
50
    Tensor<OutDataType> output_tensor_host(output_shape);
    if(!host_permute(input_tensor, input_axes, PassThrough{}, output_tensor_host))
Po-Yen, Chen's avatar
Po-Yen, Chen committed
51
52
53
54
    {
        return false;
    }

55
56
    return ck::utils::check_err(output_tensor.AsSpan<const OutDataType>(),
                                output_tensor_host.AsSpan<const OutDataType>(),
57
58
59
                                "Error: incorrect results in output tensor",
                                1e-6,
                                1e-6);
Po-Yen, Chen's avatar
Po-Yen, Chen committed
60
61
}

62
bool run_permute_element_example(const Problem::Shape& shape, const Problem::Axes& axes)
Po-Yen, Chen's avatar
Po-Yen, Chen committed
63
{
64
    return run_permute_element(Problem{shape, axes});
Po-Yen, Chen's avatar
Po-Yen, Chen committed
65
}