Function main

examples/migraphx/custom_op_hip_kernel/custom_op_hip_kernel.cpp:93–132

};

int main(int argc, const char* argv[])
{
    square_custom_op square_op;
    migraphx::register_experimental_custom_op(square_op);
    migraphx::program p;
    migraphx::shape s{migraphx_shape_float_type, {32, 256}};
    migraphx::module m = p.get_main_module();
    auto x             = m.add_parameter("x", s);
    auto neg_ins       = m.add_instruction(migraphx::operation("neg"), x);
    // add allocation for the custom_kernel's output buffer
    auto alloc = m.add_allocation(s);
    auto custom_kernel =
        m.add_instruction(migraphx::operation("square_custom_op"), {neg_ins, alloc});
    auto relu_ins = m.add_instruction(migraphx::operation("relu"), {custom_kernel});
    m.add_return({relu_ins});
    migraphx::compile_options options;
    // set offload copy to true for GPUs
    options.set_offload_copy();
    p.compile(migraphx::target("gpu"), options);
    migraphx::program_parameters pp;
    std::vector<float> x_data(s.elements());
    std::iota(x_data.begin(), x_data.end(), 0);
    pp.add("x", migraphx::argument(s, x_data.data()));
    auto results                       = p.eval(pp);
    auto result                        = results[0];
    std::vector<float> expected_result = x_data;
    std::transform(expected_result.begin(),
                   expected_result.end(),
                   expected_result.begin(),
                   [](auto i) { return std::pow(i, 2); });
    if(bool{result == migraphx::argument(s, expected_result.data())})
    {
        std::cout << "Successfully executed custom HIP kernel example\n";
    }
    else
    {
        std::cout << "Custom HIP kernel example failed\n";
    }
    return 0;
}