"tests/kernels/cache.py" did not exist on "cbf8779afafdaba2ddc6e2212d67c40f1b6e11fd"
run_permute_bundle_example.inc 3.25 KB
Newer Older
1
2
3
4
5
6
7
// SPDX-License-Identifier: MIT
// Copyright (c) 2018-2022, Advanced Micro Devices, Inc. All rights reserved.

#pragma once

bool run_permute_bundle(const Problem& problem)
{
Po-Yen, Chen's avatar
Po-Yen, Chen committed
8
    constexpr std::size_t NumElemsInBundle = sizeof(BundleType) / sizeof(DataType);
9
10
11
12
13
14
15

    using std::begin, std::end;

    const auto& shape = problem.shape;
    ck::remove_cvref_t<decltype(shape)> transposed_shape;
    transpose_shape(problem.shape, problem.axes, begin(transposed_shape));

Po-Yen, Chen's avatar
Po-Yen, Chen committed
16
17
    Tensor<BundleType> a(shape);
    Tensor<BundleType> b(transposed_shape);
18

Po-Yen, Chen's avatar
Po-Yen, Chen committed
19
    // initialize tensor by assigning DataType values
20
    using std::data, std::size;
21
22
    ck::utils::FillUniformDistribution<DataType>{-1.f, 1.f}(ck::span<DataType>{
        reinterpret_cast<DataType*>(data(a)), a.GetElementSpaceSize() * NumElemsInBundle});
23

24
25
    DeviceMem a_device_buf(a.GetElementSpaceSizeInBytes());
    DeviceMem b_device_buf(b.GetElementSpaceSizeInBytes());
26

27
    a_device_buf.ToDevice(data(a));
28

29
30
    std::array<ck::index_t, Problem::NumDim> a_lengths, b_lengths;
    std::array<ck::index_t, Problem::NumDim> a_strides, b_strides;
31
32
33
34
35

    const void* input = a_device_buf.GetDeviceBuffer();
    void* output      = b_device_buf.GetDeviceBuffer();

    std::copy(begin(shape), end(shape), begin(a_lengths));
36
    std::copy(begin(a.GetStrides()), end(a.GetStrides()), begin(a_strides));
37
    std::copy(begin(transposed_shape), end(transposed_shape), begin(b_lengths));
38
    std::copy(begin(b.GetStrides()), end(b.GetStrides()), begin(b_strides));
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57

    static_assert(std::is_default_constructible_v<DevicePermuteInstance>);

    auto permute  = DevicePermuteInstance{};
    auto argument = permute.MakeArgument(
        a_lengths, a_strides, b_lengths, b_strides, input, output, PassThrough{});

    if(!permute.IsSupportedArgument(argument))
    {
        std::cerr << "The runtime parameters seems not supported by the device instance, exiting!"
                  << std::endl;
        return false;
    };

    auto invoker   = permute.MakeInvoker();
    float ave_time = invoker.Run(argument, StreamConfig{nullptr, true});

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

58
    b_device_buf.FromDevice(data(b));
59

Po-Yen, Chen's avatar
Po-Yen, Chen committed
60
61
    // extend tensor shape from [N, H, W] to [N, H, W, NumElemsInBundle]
    const auto extended_shape = extend_shape(shape, NumElemsInBundle);
62
63
64
65
66
67
    const auto extended_axes  = extend_axes(problem.axes);

    ck::remove_cvref_t<decltype(extended_shape)> transposed_extended_shape;
    transpose_shape(extended_shape, extended_axes, begin(transposed_extended_shape));

    Tensor<DataType> extended_a(extended_shape);
68
    std::memcpy(data(extended_a), data(a), a.GetElementSpaceSizeInBytes());
69
70
71
72
73
74
75
76

    Tensor<DataType> extended_host_b(transposed_extended_shape);
    if(!host_permute(extended_a, extended_axes, PassThrough{}, extended_host_b))
    {
        return false;
    }

    return ck::utils::check_err(
77
78
        ck::span<const DataType>{reinterpret_cast<DataType*>(data(b)),
                                 b.GetElementSpaceSize() * NumElemsInBundle},
79
80
81
82
83
84
        ck::span<const DataType>{extended_host_b.mData},
        "Error: incorrect results in output tensor",
        1e-6,
        1e-6);
}

85
bool run_permute_bundle_example(const Problem::Shape& shape, const Problem::Axes& axes)
86
{
87
    return run_permute_bundle(Problem{shape, axes});
88
}