| 82 | virtual bool runs_on_offload_target() const override { return true; } |
| 83 | |
| 84 | virtual migraphx::argument |
| 85 | compute(migraphx::context ctx, migraphx::shape, migraphx::arguments inputs) const override |
| 86 | { |
| 87 | // This custom op simply sets first half size_bytes of the input to 0, and rest of the half |
| 88 | // bytes are copied. for this custom_op, it does its computation on the "GPU". Therefore, |
| 89 | // `runs_on_offload_target()` is set to "true". |
| 90 | auto* input_buffer_ptr = inputs[0].data(); |
| 91 | auto* output_buffer_ptr = inputs[1].data(); |
| 92 | auto input_bytes = inputs[0].get_shape().bytes(); |
| 93 | auto copy_bytes = input_bytes / 2; |
| 94 | MIGRAPHX_HIP_ASSERT(hipSetDevice(0)); |
| 95 | MIGRAPHX_HIP_ASSERT(hipMemcpyAsync(output_buffer_ptr, |
| 96 | input_buffer_ptr, |
| 97 | input_bytes, |
| 98 | hipMemcpyDeviceToDevice, |
| 99 | ctx.get_queue<hipStream_t>())); |
| 100 | MIGRAPHX_HIP_ASSERT(hipDeviceSynchronize()); |
| 101 | MIGRAPHX_HIP_ASSERT( |
| 102 | hipMemsetAsync(output_buffer_ptr, 0, copy_bytes, ctx.get_queue<hipStream_t>())); |
| 103 | MIGRAPHX_HIP_ASSERT(hipDeviceSynchronize()); |
| 104 | return inputs[1]; |
| 105 | } |
| 106 | |
| 107 | virtual migraphx::shape compute_shape(migraphx::shapes inputs) const override |
| 108 | { |