Module ml::gpu_miopen
MIOpen convolution, pooling and U-Net on ROCm.
The AMD counterpart of gpu_cudnn.h (same API with rocm_ in place of cuda_), layered on gpu_rocm.h: tensors are RocmBufs in NCHW order (cnn.h's CHW FeatureMap with a leading batch axis), and every op is queued on gpu_rocm's stream, so it orders with the RocmBuf ops there and can be recorded into a RocmGraph.
Convolutions run through MIOpen. The first call for each shape asks MIOpen to time its algorithms on that call's buffers (miopenFindConvolutionForwardAlgorithm) and keeps the fastest, with its workspace, in a plan cache — so make that first call outside graph capture. The bias add and ReLU after a convolution are one fused HIP kernel (gpu_rocm.h's rocm_bias_channels). Link with: -lMIOpen plus gpu_rocm's -lamdhip64 -lhiprtc -lrocblas.
Structs
| RocmConv2D | A Conv2D (cnn.h) whose weight [outC, inC, kH, kW] and bias [outC] live on the GPU. weight.dev == 0 if the upload failed. |
| RocmHalfConv2D | The same convolution on 16-bit NCHW tensors (gpu_rocm.h's RocmHalfBuf): MIOpen's half / bfloat16 kernels, which accumulate in fp32 and round once per output; the bias add and ReLU are one fused 16-bit kernel (rocm_half_bias). Weights are rounded from a Conv2D's fp32 ones on upload. x, y and the layer must share one dtype. |
| RocmUNet | Weights plus every intermediate activation, sized once for a fixed batch x height x width, so a forward pass allocates nothing. Each decoder step upsamples straight into its concat buffer (rocm_upsample2x_concat). |
Functions
| rocm_conv_out_dim | Output side length of a convolution or pooling window: (dim + 2*padding - kernel)/stride + 1. |
| rocm_maxpool2d | kernel x kernel window, given stride, no padding (FeatureMap::maxpool2d / FeatureMap::avgpool2d). x: [batch, channels, height, width]. |
| rocm_avgpool2d | |
| rocm_concat_channels | out[b] = (a[b], c[b]) along channels, for each of 'batch' images whose planes are 'plane' = height*width elements (concat_channels). |
| conv2d_forward_rocm | Same signatures and results as the CPU functions; each falls back to the CPU implementation if there is no usable GPU. Every call uploads the weights — inside a loop, upload once and use the queued ops above. |
| unet_forward_rocm |