#include "common/nn_model.h"
#include "system/target_clones.h"
#include <glib/gstdio.h>
#include <json-glib/json-glib.h>
#include <limits.h>
#include <math.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

Include dependency graph for nn_model.c:

Data Structures
struct	nn_conv_t

struct	nn_unet_t

struct	dt_nn_model_t

struct	nn_header_t

struct	dt_nn_cl_t

Macros
#define	NN_MAX_DEPTH 8

#define	NN_MIN(a, b) ((a) < (b) ? (a) : (b))

#define	NN_MAX_DEVICES 16

#define	NN_OC_BLOCK 4

#define	NN_LEDGER(delta)

#define	NN_CL_ALLOC(var, floats)

#define	NN_CL_FREE(var)

Typedefs
typedef struct nn_conv_t	nn_conv_t

typedef struct nn_unet_t	nn_unet_t

typedef struct nn_header_t	nn_header_t

Functions
void	dt_nn_set_allocator (dt_nn_alloc_f alloc_fn, dt_nn_free_f free_fn)

static void *	_nn_alloc (size_t floats, int long_lived)

static void	_nn_free (void *p)

static void	_err (char err, size_t err_len, const char msg)

static int	_wire_conv (const nn_header_t h, const char prefix, int out_ch, int in_ch, int k, nn_conv_t cv, char err, size_t err_len)

static int	_wire_unet (const nn_header_t h, const char stage_prefix, int base, int depth, int in_ch, int out_ch, nn_unet_t u, char err, size_t err_len)

static int	_read_net_cfg (JsonObject cfg, int out_ch_max, int base, int depth, int in_ch, int *out_ch)

dt_nn_model_t *	dt_nn_model_load (const char path, char err, size_t err_len)

static void	dt_nn_model_free_cl (dt_nn_model_t *m)

void	dt_nn_model_free (dt_nn_model_t *m)

int	dt_nn_model_in_channels (const dt_nn_model_t *m)

int	dt_nn_model_out_channels (const dt_nn_model_t *m)

int	dt_nn_model_bin (const dt_nn_model_t *m, const int is_xtrans)

int	dt_nn_model_coarse_in_channels (const dt_nn_model_t *m)

int	dt_nn_model_coarse_out_channels (const dt_nn_model_t *m)

int	dt_nn_model_anchor (const dt_nn_model_t *m)

static int	_lcm (int a, int b)

int	dt_nn_model_alignment (const dt_nn_model_t *m)

static __DT_CLONE_TARGETS__ void	_conv2d (const nn_conv_t cv, const float in, int w, int h, int stride, int pad, float *out)

static __DT_CLONE_TARGETS__ void	_gelu (float *x, size_t n)

static size_t	_unet_peak_floats (const nn_unet_t *u, size_t wh, int cl_variant)

static float	_scratch_per_px (const dt_nn_model_t *m, int cl_variant)

float	dt_nn_unet_scratch_per_px (const dt_nn_model_t *m)

float	dt_nn_unet_scratch_per_px_cl (const dt_nn_model_t *m)

float	dt_nn_unet_scratch_maxblock_per_px (const dt_nn_model_t *m)

size_t	dt_nn_unet_scratch_bytes (const dt_nn_model_t *m, int width, int height)

static __DT_CLONE_TARGETS__ void	_conv2d_cat2 (const nn_conv_t cv, const float a, int in_ch_a, const float b, int w, int h, float out)

static int	_unet_forward (const nn_unet_t u, const float in, float *out, int width, int height, int residual_ch)

int	dt_nn_unet_apply (const dt_nn_model_t m, const float in, float *out, int width, int height)

int	dt_nn_unet_apply_stage (const dt_nn_model_t m, int stage, const float in, float *out, int width, int height, int apply_residual)

__DT_CLONE_TARGETS__ void	dt_nn_bin_planes (const float planes, int pw, int ph, int bin, float out_rgb, float *out_cnt)

__DT_CLONE_TARGETS__ void	dt_nn_upsample_nearest (const float in, int ch, int w, int h, int factor, float out)

dt_nn_cl_t *	dt_nn_cl_create (int program)

void	dt_nn_cl_destroy (dt_nn_cl_t *cl)

static cl_mem	_weights_cl (const dt_nn_model_t *m, int devid)

static int	_conv_cl (dt_nn_cl_t cl, int devid, cl_mem weights, const float blob_base, cl_mem in, cl_mem out, int w, int h, const nn_conv_t *cv, int stride, int pad, int do_gelu)

static int	_upsample_cl (dt_nn_cl_t *cl, int devid, cl_mem in, cl_mem out, int w, int h, int ch)

static int	_unet_forward_cl (const dt_nn_model_t m, const nn_unet_t u, dt_nn_cl_t *cl, int devid, cl_mem dev_in, cl_mem dev_out, int width, int height)

int	dt_nn_unet_apply_stage_cl (const dt_nn_model_t m, int stage, dt_nn_cl_t cl, int devid, cl_mem dev_in, cl_mem dev_out, int width, int height)

Variables
static dt_nn_alloc_f	_nn_alloc_fn = NULL

static dt_nn_free_f	_nn_free_fn = NULL

Macro Definition Documentation

◆ NN_CL_ALLOC

#define NN_CL_ALLOC	(	var,
		floats
	)

Value:

  do                                                                                                          \
  {                                                                                                           \
    var = dt_opencl_alloc_device_buffer(devid, (floats) * sizeof(float));                                     \
    if(!var)                                                                                                  \
    {                                                                                                         \
      err = -1;                                                                                               \
      goto cleanup;                                                                                           \
    }                                                                                                         \
  } while(0)

◆ NN_CL_FREE

#define NN_CL_FREE ( var )

Value:

  do                                                                                                          \
  {                                                                                                           \
    if(var) dt_opencl_release_mem_object(var);                                                                \
    var = NULL;                                                                                               \
  } while(0)

◆ NN_LEDGER

#define NN_LEDGER ( delta )

Value:

  do                                                                                                          \
  {                                                                                                           \
    live += (delta);                                                                                          \
    if(live > peak) peak = live;                                                                              \
  } while(0)

◆ NN_MAX_DEPTH

#define NN_MAX_DEPTH 8

Definition at line 33 of file nn_model.c.

◆ NN_MAX_DEVICES

#define NN_MAX_DEVICES 16

Definition at line 44 of file nn_model.c.

◆ NN_MIN

#define NN_MIN	(	a,
		b
	)	((a) < (b) ? (a) : (b))

Definition at line 34 of file nn_model.c.

◆ NN_OC_BLOCK

#define NN_OC_BLOCK 4

Definition at line 491 of file nn_model.c.

Typedef Documentation

◆ nn_conv_t

typedef struct nn_conv_t nn_conv_t

◆ nn_header_t

typedef struct nn_header_t nn_header_t

◆ nn_unet_t

typedef struct nn_unet_t nn_unet_t

Function Documentation

◆ _conv2d()

static __DT_CLONE_TARGETS__ void _conv2d	(	const nn_conv_t *	cv,
		const float *	in,
		int	w,
		int	h,
		int	stride,
		int	pad,
		float *	out
	)

static

Definition at line 512 of file nn_model.c.

References k, L, NN_MIN, NN_OC_BLOCK, out, r, w1, w2, and w3.

Referenced by _unet_forward().

◆ _conv2d_cat2()

static __DT_CLONE_TARGETS__ void _conv2d_cat2	(	const nn_conv_t *	cv,
		const float *	a,
		int	in_ch_a,
		const float *	b,
		int	w,
		int	h,
		float *	out
	)

static

Definition at line 759 of file nn_model.c.

References k, L, NN_MIN, NN_OC_BLOCK, out, r, w1, w2, and w3.

Referenced by _unet_forward().

◆ _conv_cl()

static int _conv_cl	(	dt_nn_cl_t *	cl,
		int	devid,
		cl_mem	weights,
		const float *	blob_base,
		cl_mem	in,
		cl_mem	out,
		int	w,
		int	h,
		const nn_conv_t *	cv,
		int	stride,
		int	pad,
		int	do_gelu
	)

static

Definition at line 1129 of file nn_model.c.

References dt_opencl_enqueue_kernel_2d_with_local(), dt_opencl_set_kernel_arg(), dt_nn_cl_t::kernel_conv, dt_nn_cl_t::kernel_conv3x3, L, NN_MIN, and out.

Referenced by _unet_forward_cl().

◆ _err()

static void _err	(	char *	err,
		size_t	err_len,
		const char *	msg
	)

static

Definition at line 108 of file nn_model.c.

References L.

Referenced by _wire_conv(), and dt_nn_model_load().

◆ _gelu()

static __DT_CLONE_TARGETS__ void _gelu	(	float *	x,
		size_t	n
	)

static

Definition at line 612 of file nn_model.c.

References i, L, n, and x.

Referenced by _unet_forward().

◆ _lcm()

static int _lcm	(	int	a,
		int	b
	)

static

Definition at line 447 of file nn_model.c.

References t, and x.

Referenced by dt_nn_model_alignment().

◆ _nn_alloc()

static void * _nn_alloc	(	size_t	floats,
		int	long_lived
	)

static

Definition at line 93 of file nn_model.c.

References _nn_alloc_fn, and L.

Referenced by _unet_forward().

◆ _nn_free()

static void _nn_free ( void * p )

static

Definition at line 99 of file nn_model.c.

References _nn_free_fn, and p.

Referenced by _unet_forward().

◆ _read_net_cfg()

static int _read_net_cfg	(	JsonObject *	cfg,
		int	out_ch_max,
		int *	base,
		int *	depth,
		int *	in_ch,
		int *	out_ch
	)

static

Definition at line 224 of file nn_model.c.

References L, and NN_MAX_DEPTH.

Referenced by dt_nn_model_load().

◆ _scratch_per_px()

static float _scratch_per_px	(	const dt_nn_model_t *	m,
		int	cl_variant
	)

static

Definition at line 688 of file nn_model.c.

References _unet_peak_floats(), L, m, and NN_MIN.

Referenced by dt_nn_unet_scratch_per_px(), and dt_nn_unet_scratch_per_px_cl().

◆ _unet_forward()

static int _unet_forward	(	const nn_unet_t *	u,
		const float *	in,
		float *	out,
		int	width,
		int	height,
		int	residual_ch
	)

static

Definition at line 881 of file nn_model.c.

References _conv2d(), _conv2d_cat2(), _gelu(), _nn_alloc(), _nn_free(), nn_unet_t::base, nn_unet_t::bot1, nn_unet_t::bot2, nn_unet_t::dec1, nn_unet_t::dec2, nn_unet_t::depth, nn_unet_t::down, nn_unet_t::enc1, nn_unet_t::enc2, nn_unet_t::head, height, i, L, NN_MAX_DEPTH, out, nn_conv_t::out_ch, nn_unet_t::out_ch, nn_unet_t::up, v, and width.

Referenced by dt_nn_unet_apply(), and dt_nn_unet_apply_stage().

◆ _unet_forward_cl()

static int _unet_forward_cl	(	const dt_nn_model_t *	m,
		const nn_unet_t *	u,
		dt_nn_cl_t *	cl,
		int	devid,
		cl_mem	dev_in,
		cl_mem	dev_out,
		int	width,
		int	height
	)

static

Definition at line 1224 of file nn_model.c.

References _conv_cl(), _upsample_cl(), _weights_cl(), nn_unet_t::base, nn_unet_t::bot1, nn_unet_t::bot2, cleanup(), nn_unet_t::dec1, nn_unet_t::dec2, nn_unet_t::depth, nn_unet_t::down, dt_opencl_enqueue_copy_buffer_to_buffer(), dt_opencl_release_mem_object(), nn_unet_t::enc1, nn_unet_t::enc2, nn_unet_t::head, height, i, L, m, NN_CL_ALLOC, NN_CL_FREE, NN_MAX_DEPTH, nn_conv_t::out_ch, nn_unet_t::up, v, and width.

Referenced by dt_nn_unet_apply_stage_cl().

◆ _unet_peak_floats()

static size_t _unet_peak_floats	(	const nn_unet_t *	u,
		size_t	wh,
		int	cl_variant
	)

static

Definition at line 628 of file nn_model.c.

References nn_unet_t::base, nn_unet_t::depth, i, L, NN_LEDGER, and nn_unet_t::out_ch.

Referenced by _scratch_per_px(), and dt_nn_unet_scratch_bytes().

◆ _upsample_cl()

static int _upsample_cl	(	dt_nn_cl_t *	cl,
		int	devid,
		cl_mem	in,
		cl_mem	out,
		int	w,
		int	h,
		int	ch
	)

static

Definition at line 1196 of file nn_model.c.

References ch, dt_opencl_enqueue_kernel_2d(), dt_opencl_set_kernel_arg(), dt_nn_cl_t::kernel_upsample, L, out, ROUNDUPDHT, and ROUNDUPDWD.

Referenced by _unet_forward_cl().

◆ _weights_cl()

static cl_mem _weights_cl	(	const dt_nn_model_t *	m,
		int	devid
	)

static

Definition at line 1109 of file nn_model.c.

References dt_opencl_alloc_device_buffer(), dt_opencl_release_mem_object(), dt_opencl_write_buffer_to_device(), dt_pthread_mutex_lock(), dt_pthread_mutex_unlock(), L, m, and NN_MAX_DEVICES.

Referenced by _unet_forward_cl().

◆ _wire_conv()

static int _wire_conv	(	const nn_header_t *	h,
		const char *	prefix,
		int	out_ch,
		int	in_ch,
		int	k,
		nn_conv_t *	cv,
		char *	err,
		size_t	err_len
	)

static

Definition at line 126 of file nn_model.c.

References _err(), i, k, L, n, name, nn_header_t::payload, nn_header_t::payload_size, size, t, and nn_header_t::tensors.

Referenced by _wire_unet().

◆ _wire_unet()

static int _wire_unet	(	const nn_header_t *	h,
		const char *	stage_prefix,
		int	base,
		int	depth,
		int	in_ch,
		int	out_ch,
		nn_unet_t *	u,
		char *	err,
		size_t	err_len
	)