// SafeC Standard Library — MIOpen convolution, pooling and U-Net (see gpu_miopen.h).
//
// Uses MIOpen's descriptor API on gpu_rocm's stream. Each convolution shape
// gets a cached plan: tensor/filter/convolution descriptors, the algorithm
// miopenFindConvolutionForwardAlgorithm timed fastest, and that algorithm's
// own workspace — per plan rather than shared, so a RocmGraph that captured
// a convolution keeps a valid workspace however many shapes come later.
// MIOpen's Find runs the candidates, so it needs real buffers: the plan is
// built on the first call for a shape, with that call's x, weight and y.
// safeguard: link linux -L${ROCM_PATH:-/opt/rocm}/lib -lMIOpen -lamdhip64
#pragma once
#include <std/stddef.h>
#include <std/ml/gpu_miopen.h>
#include <std/ml/gpu_rocm.h>
#include <std/ml/rocm_ffi.h>
#include <std/ml/cnn.h>
#include <std/ml/unet.h>
#include <std/mem.h>
namespace std {
// ── MIOpen C ABI (miopen/miopen.h; every handle/descriptor is a pointer) ─────
extern int miopenCreate(MiopenHandle_** handle);
extern int miopenSetStream(MiopenHandle_* handle, RocmNativeStream* stream);
extern int miopenCreateTensorDescriptor(MiopenTensorDesc_** desc);
extern int miopenSet4dTensorDescriptor(MiopenTensorDesc_* desc, int dataType, int n, int c, int h, int w);
extern int miopenDestroyTensorDescriptor(MiopenTensorDesc_* desc);
extern int miopenCreateConvolutionDescriptor(MiopenConvDesc_** desc);
extern int miopenInitConvolutionDescriptor(MiopenConvDesc_* desc, int mode, int padH, int padW, int strideH, int strideW,
int dilationH, int dilationW);
extern int miopenDestroyConvolutionDescriptor(MiopenConvDesc_* desc);
extern int miopenConvolutionForwardGetWorkSpaceSize(MiopenHandle_* handle, MiopenTensorDesc_* wDesc, MiopenTensorDesc_* xDesc, MiopenConvDesc_* convDesc,
MiopenTensorDesc_* yDesc, size_t* bytes);
extern int miopenFindConvolutionForwardAlgorithm(MiopenHandle_* handle, MiopenTensorDesc_* xDesc, const HipDevPtr_* x, MiopenTensorDesc_* wDesc,
const HipDevPtr_* w, MiopenConvDesc_* convDesc, MiopenTensorDesc_* yDesc, HipDevPtr_* y,
int requested, int* returned, MiopenConvPerf_* perfResults,
HipDevPtr_* workspace, size_t wsBytes, bool exhaustive);
extern int miopenConvolutionForward(MiopenHandle_* handle, const float* alpha, MiopenTensorDesc_* xDesc, const HipDevPtr_* x,
MiopenTensorDesc_* wDesc, const HipDevPtr_* w, MiopenConvDesc_* convDesc, int algo,
const float* beta, MiopenTensorDesc_* yDesc, HipDevPtr_* y, HipDevPtr_* workspace, size_t wsBytes);
extern int miopenCreatePoolingDescriptor(MiopenPoolDesc_** desc);
extern int miopenSet2dPoolingDescriptor(MiopenPoolDesc_* desc, int mode, int windowH, int windowW,
int padH, int padW, int strideH, int strideW);
extern int miopenPoolingForward(MiopenHandle_* handle, MiopenPoolDesc_* poolDesc, const float* alpha, MiopenTensorDesc_* xDesc, const HipDevPtr_* x,
const float* beta, MiopenTensorDesc_* yDesc, HipDevPtr_* y, bool doBackward,
HipDevPtr_* workspace, size_t wsBytes);
// The HIP runtime is declared in rocm_ffi.h.
#define MIO_HALF 0 // miopenHalf
#define MIO_FLOAT 1 // miopenFloat
#define MIO_BFLOAT16 5 // miopenBFloat16
#define MIO_CONVOLUTION 0 // miopenConvolution (cross-correlation, like cnn.sc)
#define MIO_POOL_MAX 0 // miopenPoolingMax
#define MIO_POOL_AVG 1 // miopenPoolingAverage (padding excluded; there is none)
#define MIO_PERF_BYTES 16 // sizeof(miopenConvAlgoPerf_t): algo int, time float, memory size_t
#define MIO_PERF_COUNT 8
#define MIO_PLANS 32
#define MIO_MAX_WORKSPACE 268435456UL // 256 MiB: skip algorithms that want more
// ── State ────────────────────────────────────────────────────────────────────
static int g_mio_state = 0; // 0 = not tried, 1 = ready, -1 = unavailable
static MiopenHandle_* g_mio = null;
static MiopenPoolDesc_* g_mio_pool = null;
static MiopenTensorDesc_* g_mio_pool_x = null;
static MiopenTensorDesc_* g_mio_pool_y = null;
// Plan cache. g_mp_key[s * 8 ..] = batch, inC, H, W, outC, kH, kW,
// stride | padding << 16; g_mp_used[s] = 0 for a free slot.
static unsigned long g_mp_key[MIO_PLANS * 8];
static int g_mp_used[MIO_PLANS];
static MiopenTensorDesc_* g_mp_x[MIO_PLANS];
static MiopenTensorDesc_* g_mp_w[MIO_PLANS];
static MiopenConvDesc_* g_mp_conv[MIO_PLANS];
static MiopenTensorDesc_* g_mp_y[MIO_PLANS];
static int g_mp_algo[MIO_PLANS];
static HipDevPtr_* g_mp_ws[MIO_PLANS];
static unsigned long g_mp_wsBytes[MIO_PLANS];
static int g_mp_evict = 0;
static RocmNativeStream* mio_stream_() {
?&RocmNativeStream s = rocm_stream_handle();
unsafe { return (RocmNativeStream*)s; }
}
static int miopen_ensure_() {
if (g_mio_state == 1) return rocm_available();
if (g_mio_state != 0) return 0;
g_mio_state = -1;
if (!rocm_available()) return 0;
unsafe {
if (miopenCreate(&g_mio) != 0) return 0;
if (miopenSetStream(g_mio, mio_stream_()) != 0) return 0;
if (miopenCreatePoolingDescriptor(&g_mio_pool) != 0 ||
miopenCreateTensorDescriptor(&g_mio_pool_x) != 0 ||
miopenCreateTensorDescriptor(&g_mio_pool_y) != 0) return 0;
int s = 0;
while (s < MIO_PLANS) { g_mp_used[s] = 0; s = s + 1; }
}
g_mio_state = 1;
return 1;
}
unsigned long rocm_conv_out_dim(unsigned long dim, unsigned long kernel,
unsigned long stride, unsigned long padding) {
if (stride == 0UL || dim + 2UL * padding < kernel) return 0UL;
return (dim + 2UL * padding - kernel) / stride + 1UL;
}
static int mio_int_ok_(unsigned long v) { return v > 0UL && v <= 2147483647UL; }
static void mio_plan_release_(int s) {
unsafe {
if (g_mp_used[s] == 0) return;
miopenDestroyTensorDescriptor(g_mp_x[s]);
miopenDestroyTensorDescriptor(g_mp_w[s]);
miopenDestroyConvolutionDescriptor(g_mp_conv[s]);
miopenDestroyTensorDescriptor(g_mp_y[s]);
if (g_mp_ws[s] != null) {
hipStreamSynchronize(mio_stream_()); // queued work may still use it
hipFree(g_mp_ws[s]);
}
g_mp_used[s] = 0;
}
}
// Finds or builds the plan for this convolution shape; -1 on failure.
// x, the weights and y are the buffers MIOpen's Find times its candidates on.
// The layer geometry the plans are built from (fp32 and 16-bit layers alike).
struct MioGeom_ {
unsigned long inC;
unsigned long outC;
unsigned long kH;
unsigned long kW;
unsigned long stride;
unsigned long padding;
unsigned long wdev; // the weights MIOpen's Find times on
};
static struct MioGeom_ mio_geom_(const &RocmConv2D layer) {
struct MioGeom_ g;
g.inC = layer.inChannels; g.outC = layer.outChannels; g.kH = layer.kH; g.kW = layer.kW;
g.stride = layer.stride; g.padding = layer.padding; g.wdev = layer.weight.dev;
return g;
}
static int mio_plan_(const &MioGeom_ layer, int dt, unsigned long xdev, unsigned long ydev,
unsigned long batch, unsigned long height, unsigned long width) {
unsigned long outH = rocm_conv_out_dim(height, layer.kH, layer.stride, layer.padding);
unsigned long outW = rocm_conv_out_dim(width, layer.kW, layer.stride, layer.padding);
if (!mio_int_ok_(batch) || !mio_int_ok_(height) || !mio_int_ok_(width) || !mio_int_ok_(outH) ||
!mio_int_ok_(outW) || !mio_int_ok_(layer.inC) || !mio_int_ok_(layer.outC) ||
!mio_int_ok_(layer.kH) || !mio_int_ok_(layer.kW) || layer.stride > 65535UL || layer.padding > 65535UL)
return -1;
unsigned long key[8];
key[0] = batch; key[1] = layer.inC; key[2] = height; key[3] = width;
key[4] = layer.outC; key[5] = layer.kH; key[6] = layer.kW;
key[7] = layer.stride | (layer.padding << 16) | ((unsigned long)dt << 32);
int slot = -1;
int s = 0;
while (s < MIO_PLANS) {
if (g_mp_used[s] != 0) {
int same = 1;
int i = 0;
while (i < 8) { if (g_mp_key[s * 8 + i] != key[i]) same = 0; i = i + 1; }
if (same) return s;
} else if (slot < 0) {
slot = s;
}
s = s + 1;
}
if (slot < 0) {
slot = g_mp_evict;
g_mp_evict = (g_mp_evict + 1) % MIO_PLANS;
mio_plan_release_(slot);
}
unsafe {
MiopenTensorDesc_* xd = null;
MiopenTensorDesc_* wd = null;
MiopenConvDesc_* cd = null;
MiopenTensorDesc_* yd = null;
int ok = miopenCreateTensorDescriptor(&xd) == 0 && miopenCreateTensorDescriptor(&wd) == 0 &&
miopenCreateConvolutionDescriptor(&cd) == 0 && miopenCreateTensorDescriptor(&yd) == 0;
ok = ok && miopenSet4dTensorDescriptor(xd, dt, (int)batch, (int)layer.inC,
(int)height, (int)width) == 0;
ok = ok && miopenSet4dTensorDescriptor(wd, dt, (int)layer.outC, (int)layer.inC,
(int)layer.kH, (int)layer.kW) == 0;
ok = ok && miopenInitConvolutionDescriptor(cd, MIO_CONVOLUTION, (int)layer.padding, (int)layer.padding,
(int)layer.stride, (int)layer.stride, 1, 1) == 0;
ok = ok && miopenSet4dTensorDescriptor(yd, dt, (int)batch, (int)layer.outC,
(int)outH, (int)outW) == 0;
// Find needs a workspace big enough for every candidate it may try;
// the winner's own (usually smaller) workspace is allocated after.
size_t findBytes = 0UL;
ok = ok && miopenConvolutionForwardGetWorkSpaceSize(g_mio, wd, xd, cd, yd, &findBytes) == 0;
if (ok && (unsigned long)findBytes > MIO_MAX_WORKSPACE) findBytes = (size_t)MIO_MAX_WORKSPACE;
HipDevPtr_* findWs = null;
if (ok && findBytes > 0UL) ok = hipMalloc(&findWs, (unsigned long)findBytes) == 0;
// miopenConvAlgoPerf_t: { int fwd_algo; float time; size_t memory; },
// fastest first.
struct MiopenConvPerf_ perf[MIO_PERF_COUNT];
int returned = 0;
ok = ok && miopenFindConvolutionForwardAlgorithm(g_mio, xd, (const HipDevPtr_*)xdev, wd,
(const HipDevPtr_*)layer.wdev, cd, yd, (HipDevPtr_*)ydev, MIO_PERF_COUNT, &returned,
&perf[0], findWs, findBytes, false) == 0;
if (findWs != null) hipFree(findWs);
int algo = -1;
unsigned long wsBytes = 0UL;
int r = 0;
while (ok && algo < 0 && r < returned) {
unsigned long mem = perf[r].memory;
if (mem <= (unsigned long)findBytes) {
algo = perf[r].algo;
wsBytes = mem;
}
r = r + 1;
}
ok = ok && algo >= 0;
HipDevPtr_* ws = null;
if (ok && wsBytes > 0UL) ok = hipMalloc(&ws, wsBytes) == 0;
if (!ok) {
if (xd != null) miopenDestroyTensorDescriptor(xd);
if (wd != null) miopenDestroyTensorDescriptor(wd);
if (cd != null) miopenDestroyConvolutionDescriptor(cd);
if (yd != null) miopenDestroyTensorDescriptor(yd);
return -1;
}
g_mp_x[slot] = xd;
g_mp_w[slot] = wd;
g_mp_conv[slot] = cd;
g_mp_y[slot] = yd;
g_mp_algo[slot] = algo;
g_mp_ws[slot] = ws;
g_mp_wsBytes[slot] = wsBytes;
int i = 0;
while (i < 8) { g_mp_key[slot * 8 + i] = key[i]; i = i + 1; }
g_mp_used[slot] = 1;
}
return slot;
}
static int mio_conv_(const &RocmConv2D layer, struct RocmBuf x, struct RocmBuf y,
unsigned long batch, unsigned long height, unsigned long width, int relu) {
if (!miopen_ensure_() || layer.weight.dev == 0UL || layer.bias.dev == 0UL) return 0;
unsigned long outH = rocm_conv_out_dim(height, layer.kH, layer.stride, layer.padding);
unsigned long outW = rocm_conv_out_dim(width, layer.kW, layer.stride, layer.padding);
if (x.dev == 0UL || y.dev == 0UL ||
x.len < batch * layer.inChannels * height * width ||
y.len < batch * layer.outChannels * outH * outW) return 0;
struct MioGeom_ geom = mio_geom_(layer);
int p = mio_plan_(&geom, MIO_FLOAT, x.dev, y.dev, batch, height, width);
if (p < 0) return 0;
float one = 1.0f;
float zero = 0.0f;
unsafe {
if (miopenConvolutionForward(g_mio, (const float*)&one, g_mp_x[p], (const HipDevPtr_*)x.dev,
g_mp_w[p], (const HipDevPtr_*)layer.weight.dev, g_mp_conv[p], g_mp_algo[p],
(const float*)&zero, g_mp_y[p], (HipDevPtr_*)y.dev, g_mp_ws[p], (size_t)g_mp_wsBytes[p]) != 0)
return 0;
}
// Bias (and ReLU) in one pass over y.
struct RocmBuf yv;
yv.dev = y.dev;
yv.len = batch * layer.outChannels * outH * outW;
return rocm_bias_channels(yv, layer.bias, batch, layer.outChannels, outH * outW, relu);
}
int RocmConv2D::forward(struct RocmBuf x, struct RocmBuf y, unsigned long batch, unsigned long height,
unsigned long width) const {
return mio_conv_(self, x, y, batch, height, width, 0);
}
int RocmConv2D::forward_relu(struct RocmBuf x, struct RocmBuf y, unsigned long batch, unsigned long height,
unsigned long width) const {
return mio_conv_(self, x, y, batch, height, width, 1);
}
static int mio_pool_dt_(int mode, int dt, unsigned long xdev, unsigned long xlen, unsigned long ydev, unsigned long ylen,
unsigned long batch, unsigned long channels,
unsigned long height, unsigned long width, unsigned long kernel, unsigned long stride) {
if (!miopen_ensure_()) return 0;
unsigned long outH = rocm_conv_out_dim(height, kernel, stride, 0UL);
unsigned long outW = rocm_conv_out_dim(width, kernel, stride, 0UL);
if (!mio_int_ok_(batch) || !mio_int_ok_(channels) || !mio_int_ok_(height) || !mio_int_ok_(width) ||
!mio_int_ok_(kernel) || !mio_int_ok_(stride) || outH == 0UL || outW == 0UL ||
xdev == 0UL || ydev == 0UL ||
xlen < batch * channels * height * width || ylen < batch * channels * outH * outW) return 0;
float one = 1.0f;
float zero = 0.0f;
unsafe {
if (miopenSet2dPoolingDescriptor(g_mio_pool, mode, (int)kernel, (int)kernel, 0, 0,
(int)stride, (int)stride) != 0 ||
miopenSet4dTensorDescriptor(g_mio_pool_x, dt, (int)batch, (int)channels,
(int)height, (int)width) != 0 ||
miopenSet4dTensorDescriptor(g_mio_pool_y, dt, (int)batch, (int)channels,
(int)outH, (int)outW) != 0) return 0;
// Forward only: no index workspace.
return miopenPoolingForward(g_mio, g_mio_pool, (const float*)&one, g_mio_pool_x, (const HipDevPtr_*)xdev,
(const float*)&zero, g_mio_pool_y, (HipDevPtr_*)ydev, false,
null, (size_t)0) == 0;
}
}
static int mio_pool_(int mode, struct RocmBuf x, struct RocmBuf y, unsigned long batch, unsigned long channels,
unsigned long height, unsigned long width, unsigned long kernel, unsigned long stride) {
return mio_pool_dt_(mode, MIO_FLOAT, x.dev, x.len, y.dev, y.len, batch, channels, height, width, kernel, stride);
}
int rocm_maxpool2d(struct RocmBuf x, struct RocmBuf y, unsigned long batch, unsigned long channels,
unsigned long height, unsigned long width, unsigned long kernel, unsigned long stride) {
return mio_pool_(MIO_POOL_MAX, x, y, batch, channels, height, width, kernel, stride);
}
int rocm_avgpool2d(struct RocmBuf x, struct RocmBuf y, unsigned long batch, unsigned long channels,
unsigned long height, unsigned long width, unsigned long kernel, unsigned long stride) {
return mio_pool_(MIO_POOL_AVG, x, y, batch, channels, height, width, kernel, stride);
}
int rocm_concat_channels(struct RocmBuf a, unsigned long aChannels, struct RocmBuf c, unsigned long cChannels,
struct RocmBuf out, unsigned long batch, unsigned long plane) {
// One kernel over every image, instead of two copies per image.
return rocm_concat(a, aChannels * plane, c, cChannels * plane, out, batch);
}
// ── Layers ───────────────────────────────────────────────────────────────────
struct RocmConv2D RocmConv2D::upload(const &Conv2D layer) {
struct RocmConv2D d;
d.inChannels = layer.inChannels;
d.outChannels = layer.outChannels;
d.kH = layer.kH;
d.kW = layer.kW;
d.stride = layer.stride;
d.padding = layer.padding;
unsigned long wn = layer.outChannels * layer.inChannels * layer.kH * layer.kW;
d.weight = RocmBuf::upload_n_(layer.weight, wn);
d.bias = RocmBuf::upload_n_(layer.bias, layer.outChannels);
if (d.weight.dev == 0UL || d.bias.dev == 0UL) {
d.weight.free();
d.bias.free();
}
return d;
}
void RocmConv2D::free() {
self.weight.free();
self.bias.free();
}
// ── fp16 / bf16 ──────────────────────────────────────────────────────────────
static int mio_dt_(int dtype) { return dtype == ROCM_BF16 ? MIO_BFLOAT16 : MIO_HALF; }
struct RocmHalfConv2D RocmHalfConv2D::upload(const &Conv2D layer, int dtype) {
struct RocmHalfConv2D d;
d.inChannels = layer.inChannels;
d.outChannels = layer.outChannels;
d.kH = layer.kH;
d.kW = layer.kW;
d.stride = layer.stride;
d.padding = layer.padding;
unsigned long wn = layer.outChannels * layer.inChannels * layer.kH * layer.kW;
d.weight = RocmHalfBuf::upload_n_(layer.weight, wn, dtype);
d.bias = RocmHalfBuf::upload_n_(layer.bias, layer.outChannels, dtype);
if (d.weight.dev == 0UL || d.bias.dev == 0UL) {
d.weight.free();
d.bias.free();
}
return d;
}
void RocmHalfConv2D::free() {
self.weight.free();
self.bias.free();
}
static int mio_half_conv_(const &RocmHalfConv2D layer, struct RocmHalfBuf x, struct RocmHalfBuf y,
unsigned long batch, unsigned long height, unsigned long width, int relu) {
if (!miopen_ensure_() || layer.weight.dev == 0UL || layer.bias.dev == 0UL) return 0;
if (x.dtype != layer.weight.dtype || y.dtype != layer.weight.dtype) return 0;
unsigned long outH = rocm_conv_out_dim(height, layer.kH, layer.stride, layer.padding);
unsigned long outW = rocm_conv_out_dim(width, layer.kW, layer.stride, layer.padding);
if (x.dev == 0UL || y.dev == 0UL ||
x.len < batch * layer.inChannels * height * width ||
y.len < batch * layer.outChannels * outH * outW) return 0;
struct MioGeom_ geom;
geom.inC = layer.inChannels; geom.outC = layer.outChannels; geom.kH = layer.kH; geom.kW = layer.kW;
geom.stride = layer.stride; geom.padding = layer.padding; geom.wdev = layer.weight.dev;
int p = mio_plan_(&geom, mio_dt_(x.dtype), x.dev, y.dev, batch, height, width);
if (p < 0) return 0;
float one = 1.0f;
float zero = 0.0f;
unsafe {
if (miopenConvolutionForward(g_mio, (const float*)&one, g_mp_x[p], (const HipDevPtr_*)x.dev,
g_mp_w[p], (const HipDevPtr_*)layer.weight.dev, g_mp_conv[p], g_mp_algo[p],
(const float*)&zero, g_mp_y[p], (HipDevPtr_*)y.dev, g_mp_ws[p], (size_t)g_mp_wsBytes[p]) != 0)
return 0;
}
return rocm_half_bias(y, layer.bias, batch, layer.outChannels, outH * outW, relu);
}
int RocmHalfConv2D::forward(struct RocmHalfBuf x, struct RocmHalfBuf y, unsigned long batch, unsigned long height,
unsigned long width) const {
return mio_half_conv_(self, x, y, batch, height, width, 0);
}
int RocmHalfConv2D::forward_relu(struct RocmHalfBuf x, struct RocmHalfBuf y, unsigned long batch,
unsigned long height, unsigned long width) const {
return mio_half_conv_(self, x, y, batch, height, width, 1);
}
// ── U-Net ────────────────────────────────────────────────────────────────────
struct RocmUNet RocmUNet::new(const &UNet net, unsigned long batch, unsigned long height, unsigned long width) {
struct RocmUNet g;
g.batch = batch;
g.height = height;
g.width = width;
g.encConv1 = RocmConv2D::upload(&net.encConv1);
g.encConv2 = RocmConv2D::upload(&net.encConv2);
g.bottleneck = RocmConv2D::upload(&net.bottleneck);
g.decConv1 = RocmConv2D::upload(&net.decConv1);
g.decConv2 = RocmConv2D::upload(&net.decConv2);
unsigned long c1 = net.encConv1.outChannels;
unsigned long c2 = net.encConv2.outChannels;
unsigned long hw = height * width;
unsigned long hw2 = hw / 4UL;
unsigned long hw4 = hw / 16UL;
g.e1 = RocmBuf::alloc(batch * c1 * hw);
g.p1 = RocmBuf::alloc(batch * c1 * hw2);
g.e2 = RocmBuf::alloc(batch * c2 * hw2);
g.p2 = RocmBuf::alloc(batch * c2 * hw4);
g.mid = RocmBuf::alloc(batch * c2 * hw4);
g.cat1 = RocmBuf::alloc(batch * 2UL * c2 * hw2);
g.d1 = RocmBuf::alloc(batch * c1 * hw2);
g.cat2 = RocmBuf::alloc(batch * 2UL * c1 * hw);
g.ok = batch > 0UL && height % 4UL == 0UL && width % 4UL == 0UL && height > 0UL && width > 0UL &&
g.encConv1.weight.dev != 0UL && g.encConv2.weight.dev != 0UL && g.bottleneck.weight.dev != 0UL &&
g.decConv1.weight.dev != 0UL && g.decConv2.weight.dev != 0UL &&
g.e1.dev != 0UL && g.p1.dev != 0UL && g.e2.dev != 0UL && g.p2.dev != 0UL && g.mid.dev != 0UL &&
g.cat1.dev != 0UL && g.d1.dev != 0UL && g.cat2.dev != 0UL;
return g;
}
int RocmUNet::forward(struct RocmBuf x, struct RocmBuf y) const {
if (!self.ok) return 0;
unsigned long B = self.batch;
unsigned long H = self.height;
unsigned long W = self.width;
unsigned long c1 = self.encConv1.outChannels;
unsigned long c2 = self.encConv2.outChannels;
// Encoder: conv+relu -> 2x2 max pool, twice.
return self.encConv1.forward_relu(x, self.e1, B, H, W) &&
rocm_maxpool2d(self.e1, self.p1, B, c1, H, W, 2UL, 2UL) &&
self.encConv2.forward_relu(self.p1, self.e2, B, H / 2UL, W / 2UL) &&
rocm_maxpool2d(self.e2, self.p2, B, c2, H / 2UL, W / 2UL, 2UL, 2UL) &&
self.bottleneck.forward_relu(self.p2, self.mid, B, H / 4UL, W / 4UL) &&
// Decoder: upsample + concat the skip (one kernel), conv(+relu).
rocm_upsample2x_concat(self.mid, c2, self.e2, c2, self.cat1, B, H / 4UL, W / 4UL) &&
self.decConv1.forward_relu(self.cat1, self.d1, B, H / 2UL, W / 2UL) &&
rocm_upsample2x_concat(self.d1, c1, self.e1, c1, self.cat2, B, H / 2UL, W / 2UL) &&
self.decConv2.forward(self.cat2, y, B, H, W);
}
void RocmUNet::free() {
self.encConv1.free();
self.encConv2.free();
self.bottleneck.free();
self.decConv1.free();
self.decConv2.free();
self.e1.free();
self.p1.free();
self.e2.free();
self.p2.free();
self.mid.free();
self.cat1.free();
self.d1.free();
self.cat2.free();
self.ok = 0;
}
// ── Host convenience ─────────────────────────────────────────────────────────
struct FeatureMap conv2d_forward_rocm(const &Conv2D layer, const &FeatureMap input) {
unsigned long inC = input.channels;
unsigned long h = input.height;
unsigned long w = input.width;
unsigned long outH = rocm_conv_out_dim(h, layer.kH, layer.stride, layer.padding);
unsigned long outW = rocm_conv_out_dim(w, layer.kW, layer.stride, layer.padding);
if (inC == layer.inChannels && outH > 0UL && outW > 0UL && miopen_ensure_()) {
struct RocmConv2D d = RocmConv2D::upload(layer);
struct RocmBuf x = RocmBuf::upload_n_(input.data, inC * h * w);
struct RocmBuf y = RocmBuf::alloc(layer.outChannels * outH * outW);
struct FeatureMap out = FeatureMap::new(layer.outChannels, outH, outW);
int ok = d.weight.dev != 0UL && x.dev != 0UL && y.dev != 0UL &&
d.forward(x, y, 1UL, h, w) && rocm_read_n_(y, out.data, y.len);
d.free();
x.free();
y.free();
if (ok) return out;
out.free();
}
return layer.forward(input);
}
struct FeatureMap unet_forward_rocm(&UNet net, const &FeatureMap input) {
unsigned long h = input.height;
unsigned long w = input.width;
if (input.channels == net.encConv1.inChannels && miopen_ensure_()) {
struct RocmUNet g = RocmUNet::new(net, 1UL, h, w);
unsigned long outC = net.decConv2.outChannels;
struct RocmBuf x = RocmBuf::upload_n_(input.data, input.channels * h * w);
struct RocmBuf y = RocmBuf::alloc(outC * h * w);
struct FeatureMap out = FeatureMap::new(outC, h, w);
int ok = g.ok && x.dev != 0UL && y.dev != 0UL &&
g.forward(x, y) && rocm_read_n_(y, out.data, y.len);
g.free();
x.free();
y.free();
if (ok) return out;
out.free();
}
return net.forward(input);
}
} // namespace std