| 37 | virtual bool runs_on_offload_target() const override { return false; } |
| 38 | |
| 39 | virtual migraphx::argument |
| 40 | compute(migraphx::context ctx, migraphx::shape, migraphx::arguments inputs) const override |
| 41 | { |
| 42 | // This custom op simply sets first half size_bytes of the input to 0, and rest of the half |
| 43 | // bytes are copied. for this custom_op, it does its computation on the host. Therefore, |
| 44 | // `runs_on_offload_target()` is set to false. MIGraphX would inject necessary buffer copies |
| 45 | // to and from GPU to Host based on `runs_on_offload_targe()` flag for input buffers as well |
| 46 | // as the output buffers |
| 47 | auto* input_buffer_ptr = inputs[0].data(); |
| 48 | auto* output_buffer_ptr = inputs[1].data(); |
| 49 | auto input_bytes = inputs[0].get_shape().bytes(); |
| 50 | auto copy_bytes = input_bytes / 2; |
| 51 | MIGRAPHX_HIP_ASSERT(hipSetDevice(0)); |
| 52 | MIGRAPHX_HIP_ASSERT(hipMemcpyAsync(output_buffer_ptr, |
| 53 | input_buffer_ptr, |
| 54 | input_bytes, |
| 55 | hipMemcpyHostToHost, |
| 56 | ctx.get_queue<hipStream_t>())); |
| 57 | MIGRAPHX_HIP_ASSERT(hipDeviceSynchronize()); |
| 58 | MIGRAPHX_HIP_ASSERT( |
| 59 | hipMemsetAsync(output_buffer_ptr, 0, copy_bytes, ctx.get_queue<hipStream_t>())); |
| 60 | MIGRAPHX_HIP_ASSERT(hipDeviceSynchronize()); |
| 61 | return inputs[1]; |
| 62 | } |
| 63 | |
| 64 | virtual migraphx::shape compute_shape(migraphx::shapes inputs) const override |
| 65 | { |