Add elementwise_affine support for rmsnorm.

Author: liuliu

lib/nnc/ccv_cnnp_model_addons.c

@@ -2555,7 +2555,7 @@ static const ccv_cnnp_model_vtab_t ccv_cnnp_rmsnorm_isa = {
 	.copy = _ccv_cnnp_rmsnorm_copy,
 };
 
-ccv_cnnp_model_t* ccv_cnnp_rmsnorm(const float epsilon, const int axis[CCV_NNC_MAX_DIM_ALLOC], const int axis_count, const int is_trainable, const char* const name)
+ccv_cnnp_model_t* ccv_cnnp_rmsnorm(const float epsilon, const int axis[CCV_NNC_MAX_DIM_ALLOC], const int axis_count, const int elementwise_affine, const int is_trainable, const char* const name)
 {
 	ccv_cnnp_model_rmsnorm_t* const model_rmsnorm = (ccv_cnnp_model_rmsnorm_t*)cccalloc(1, sizeof(ccv_cnnp_model_rmsnorm_t));
 	model_rmsnorm->super.isa = &ccv_cnnp_rmsnorm_isa;
@@ -2568,14 +2568,15 @@ ccv_cnnp_model_t* ccv_cnnp_rmsnorm(const float epsilon, const int axis[CCV_NNC_M
 	model_rmsnorm->scale.graph = 0;
 	model_rmsnorm->params.rmsnorm.epsilon = epsilon;
 	model_rmsnorm->params.rmsnorm.count = axis_count;
+	model_rmsnorm->params.rmsnorm.elementwise_affine = elementwise_affine;
 	memcpy(model_rmsnorm->params.lnorm.axis, axis, sizeof(int) * axis_count);
 	return (ccv_cnnp_model_t*)model_rmsnorm;
 }
 
 static ccv_cnnp_model_t* _ccv_cnnp_rmsnorm_copy(const ccv_cnnp_model_t* const super, void* const context)
 {
 	const ccv_cnnp_model_rmsnorm_t* const self = (const ccv_cnnp_model_rmsnorm_t*)super;
-	return ccv_cnnp_rmsnorm(self->params.rmsnorm.epsilon, self->params.rmsnorm.axis, self->params.rmsnorm.count, self->super.is_trainable, self->super.name);
+	return ccv_cnnp_rmsnorm(self->params.rmsnorm.epsilon, self->params.rmsnorm.axis, self->params.rmsnorm.count, self->params.rmsnorm.elementwise_affine, self->super.is_trainable, self->super.name);
 }
 
 // MARK - Batched Matrix Mul Layer

lib/nnc/ccv_nnc.h

@@ -170,6 +170,7 @@ typedef struct {
 			int axis[CCV_NNC_MAX_DIM_ALLOC]; /**< [rmsnorm.axis[]] The axis selected to compute mean / variance. */
 			int count; /**< [rmsnorm.count] The number of axis selected. */
 			float epsilon; /**< [rmsnorm.epsilon] The epsilon for standard derivation. */
+			int elementwise_affine; /**< [rmsnorm.elementwise_affine] Whether it supports scale. */
 		} rmsnorm;
 		struct {
 			int nesterov; /**< [sgd.nesterov] Nesterov accelerated gradient. */
@@ -4549,11 +4550,12 @@ CCV_WARN_UNUSED(ccv_cnnp_model_t*) ccv_cnnp_group_norm(const int group_axis, con
  * @param epsilon The epsilon in layer norm parameter.
  * @param axis The axis are the feature axis to compute norm.
  * @param axis_count How many axis we count as feature.
+ * @param elementwise_affine Whether it contains scale.
  * @param is_trainable Whether the parameters of this model can be trained.
  * @param name The unique name of the model.
  * @return A rmsnorm model.
  */
-CCV_WARN_UNUSED(ccv_cnnp_model_t*) ccv_cnnp_rmsnorm(const float epsilon, const int axis[CCV_NNC_MAX_DIM_ALLOC], const int axis_count, const int is_trainable, const char* const name);
+CCV_WARN_UNUSED(ccv_cnnp_model_t*) ccv_cnnp_rmsnorm(const float epsilon, const int axis[CCV_NNC_MAX_DIM_ALLOC], const int axis_count, const int elementwise_affine, const int is_trainable, const char* const name);
 /**
  * Add two input tensors together. Different from sum because this support broadcasting.
  * @param p The weight for the first input.

lib/nnc/cmd/ccv_nnc_cmd_easy.h

@@ -183,9 +183,9 @@
 // CCV_NNC_GROUP_NORM_BACKWARD
 #define CMD_GROUP_NORM_BACKWARD(_group_axis, _groups, _epsilon, _elementwise_affine, ...) ccv_nnc_cmd(CCV_NNC_GROUP_NORM_BACKWARD, 0, ((ccv_nnc_cmd_param_t){.size={.dim={1,1,1}},.gnorm={.group_axis=_group_axis,.groups=_groups,.epsilon=_epsilon,.elementwise_affine=_elementwise_affine,.reduce_count=LIST_COUNT(__VA_ARGS__),.reduce_axis={__VA_ARGS__}}}), 0)
 // CCV_NNC_RMSNORM_FORWARD
-#define CMD_RMSNORM_FORWARD(_epsilon, ...) ccv_nnc_cmd(CCV_NNC_RMSNORM_FORWARD, 0, ((ccv_nnc_cmd_param_t){.size={.dim={1,1,1}},.rmsnorm={.epsilon=_epsilon,.count=LIST_COUNT(__VA_ARGS__),.axis={__VA_ARGS__}}}), 0)
+#define CMD_RMSNORM_FORWARD(_epsilon, _elementwise_affine, ...) ccv_nnc_cmd(CCV_NNC_RMSNORM_FORWARD, 0, ((ccv_nnc_cmd_param_t){.size={.dim={1,1,1}},.rmsnorm={.epsilon=_epsilon,.elementwise_affine=_elementwise_affine,.count=LIST_COUNT(__VA_ARGS__),.axis={__VA_ARGS__}}}), 0)
 // CCV_NNC_RMSNORM_BACKWARD
-#define CMD_RMSNORM_BACKWARD(_epsilon, ...) ccv_nnc_cmd(CCV_NNC_RMSNORM_BACKWARD, 0, ((ccv_nnc_cmd_param_t){.size={.dim={1,1,1}},.rmsnorm={.epsilon=_epsilon,.count=LIST_COUNT(__VA_ARGS__),.axis={__VA_ARGS__}}}), 0)
+#define CMD_RMSNORM_BACKWARD(_epsilon, _elementwise_affine, ...) ccv_nnc_cmd(CCV_NNC_RMSNORM_BACKWARD, 0, ((ccv_nnc_cmd_param_t){.size={.dim={1,1,1}},.rmsnorm={.epsilon=_epsilon,.elementwise_affine=_elementwise_affine,.count=LIST_COUNT(__VA_ARGS__),.axis={__VA_ARGS__}}}), 0)
 // CCV_NNC_PAD_FORWARD
 #define CMD_PAD_FORWARD(_type, _begin, _end) ccv_nnc_cmd(CCV_NNC_PAD_FORWARD, 0, ((ccv_nnc_cmd_param_t){.size={.dim={ESCAPE_X _begin}},.pad={.type=_type,.end={ESCAPE_X _end}}}), 0)
 // CCV_NNC_PAD_BACKWARD

lib/nnc/cmd/norm/ccv_nnc_norm.c

@@ -287,28 +287,45 @@ REGISTER_COMMAND(CCV_NNC_GROUP_NORM_BACKWARD)(ccv_nnc_cmd_registry_t* const regi
 
 static int _ccv_nnc_rmsnorm_forw_bitmask(const ccv_nnc_cmd_param_t cmd, const int input_size, const int output_size, const uint64_t* const input_bitmasks, const int input_bitmask_size, const uint64_t* const output_bitmasks, const int output_bitmask_size)
 {
-	// 2 inputs (x, gamma)
-	// 2 outputs (y, saved_inv_std)
-	if (input_bitmasks[0] == 3u && output_bitmasks[0] == 3u)
-		return 1;
+	if (cmd.rmsnorm.elementwise_affine)
+	{
+		// 2 inputs (x, gamma)
+		// 2 outputs (y, saved_inv_std)
+		if (input_bitmasks[0] == 3u && output_bitmasks[0] == 3u)
+			return 1;
+	} else {
+		// 1 inputs (x)
+		// 2 outputs (y, saved_inv_std)
+		if (input_bitmasks[0] == 1u && output_bitmasks[0] == 3u)
+			return 1;
+	}
 	return 0;
 }
 
 static int _ccv_nnc_rmsnorm_back_bitmask(const ccv_nnc_cmd_param_t cmd, const int input_size, const int output_size, const uint64_t* const input_bitmasks, const int input_bitmask_size, const uint64_t* const output_bitmasks, const int output_bitmask_size)
 {
-	// 1 + 4 + 8 + 32
-	// Inputs (gradient, 0, x, gamma, 0, saved_inv_std)
-	// Output the propagated error, dgamma
-	if ((input_bitmasks[0] & 45u) == 45u && (output_bitmasks[0] & 3u) == 3u)
-		return 1;
-	if ((input_bitmasks[0] & 45u) == 45u && (output_bitmasks[0] & 1u) == 1u)
-		return 1;
+	if (cmd.rmsnorm.elementwise_affine)
+	{
+		// 1 + 4 + 8 + 32
+		// Inputs (gradient, 0, x, gamma, 0, saved_inv_std)
+		// Output the propagated error, dgamma
+		if ((input_bitmasks[0] & 45u) == 45u && (output_bitmasks[0] & 3u) == 3u)
+			return 1;
+		if ((input_bitmasks[0] & 45u) == 45u && (output_bitmasks[0] & 1u) == 1u)
+			return 1;
+	} else {
+		// 1 + 4 + 16
+		// Inputs (gradient, 0, x, 0, saved_inv_std)
+		// Output the propagated error
+		if ((input_bitmasks[0] & 21u) == 21u && (output_bitmasks[0] & 1u) == 1u)
+			return 1;
+	}
 	return 0;
 }
 
 static void _ccv_nnc_rmsnorm_tensor_auto_forw(const ccv_nnc_cmd_param_t cmd, const ccv_nnc_tensor_param_t* const inputs, const int input_size, const ccv_nnc_hint_t hint, ccv_nnc_tensor_param_t* const outputs, const int output_size)
 {
-	assert(input_size == 2);
+	assert(input_size == 2 || input_size == 1);
 	assert(output_size == 1 || output_size == 2);
 	outputs[0] = inputs[0];
 	if (output_size == 1)
@@ -324,7 +341,7 @@ static void _ccv_nnc_rmsnorm_tensor_auto_forw(const ccv_nnc_cmd_param_t cmd, con
 
 static void _ccv_nnc_rmsnorm_tensor_auto_back(const ccv_nnc_cmd_param_t cmd, const ccv_nnc_tensor_param_t* const inputs, const int input_size, const ccv_nnc_hint_t hint, ccv_nnc_tensor_param_t* const outputs, const int output_size)
 {
-	assert(input_size == 6);
+	assert(input_size == 6 || input_size == 5);
 	assert(output_size == 1 || output_size == 2);
 	outputs[0] = inputs[0];
 	int i, j;
@@ -351,6 +368,6 @@ REGISTER_COMMAND(CCV_NNC_RMSNORM_BACKWARD)(ccv_nnc_cmd_registry_t* const registr
 }
 
 //@REGISTER_EASY_COMMAND_MACRO(CCV_NNC_RMSNORM_FORWARD)
-#define CMD_RMSNORM_FORWARD(_epsilon, ...) ccv_nnc_cmd(CCV_NNC_RMSNORM_FORWARD, 0, ((ccv_nnc_cmd_param_t){.size={.dim={1,1,1}},.rmsnorm={.epsilon=_epsilon,.count=LIST_COUNT(__VA_ARGS__),.axis={__VA_ARGS__}}}), 0)
+#define CMD_RMSNORM_FORWARD(_epsilon, _elementwise_affine, ...) ccv_nnc_cmd(CCV_NNC_RMSNORM_FORWARD, 0, ((ccv_nnc_cmd_param_t){.size={.dim={1,1,1}},.rmsnorm={.epsilon=_epsilon,.elementwise_affine=_elementwise_affine,.count=LIST_COUNT(__VA_ARGS__),.axis={__VA_ARGS__}}}), 0)
 //@REGISTER_EASY_COMMAND_MACRO(CCV_NNC_RMSNORM_BACKWARD)
-#define CMD_RMSNORM_BACKWARD(_epsilon, ...) ccv_nnc_cmd(CCV_NNC_RMSNORM_BACKWARD, 0, ((ccv_nnc_cmd_param_t){.size={.dim={1,1,1}},.rmsnorm={.epsilon=_epsilon,.count=LIST_COUNT(__VA_ARGS__),.axis={__VA_ARGS__}}}), 0)
+#define CMD_RMSNORM_BACKWARD(_epsilon, _elementwise_affine, ...) ccv_nnc_cmd(CCV_NNC_RMSNORM_BACKWARD, 0, ((ccv_nnc_cmd_param_t){.size={.dim={1,1,1}},.rmsnorm={.epsilon=_epsilon,.elementwise_affine=_elementwise_affine,.count=LIST_COUNT(__VA_ARGS__),.axis={__VA_ARGS__}}}), 0)

lib/nnc/cmd/norm/ccv_nnc_rmsnorm_cpu_ref.c

@@ -15,9 +15,9 @@
 
 static int _ccv_nnc_rmsnorm_forw(const ccv_nnc_cmd_t cmd, const ccv_nnc_hint_t hint, const int flags, ccv_nnc_tensor_t* const* const inputs, const int input_size, ccv_nnc_tensor_t* const* const outputs, const int output_size, ccv_nnc_stream_context_t* const stream_context)
 {
-	assert(input_size == 2);
+	assert(input_size == 2 || input_size == 1);
 	ccv_nnc_tensor_view_t* const a = (ccv_nnc_tensor_view_t*)inputs[0];
-	ccv_nnc_tensor_view_t* const scale = (ccv_nnc_tensor_view_t*)inputs[1];
+	ccv_nnc_tensor_view_t* const scale = input_size >= 2 ? (ccv_nnc_tensor_view_t*)inputs[1] : 0;
 	ccv_nnc_tensor_view_t* const b = (ccv_nnc_tensor_view_t*)outputs[0];
 	ccv_nnc_tensor_view_t* const saved_inv_std = (ccv_nnc_tensor_view_t*)outputs[1];
 	assert(ccv_nnc_tensor_nd(a->info.dim) <= CCV_NNC_MAX_DIM + 2);
@@ -33,7 +33,8 @@ static int _ccv_nnc_rmsnorm_forw(const ccv_nnc_cmd_t cmd, const ccv_nnc_hint_t h
 	int bstride[CCV_NNC_MAX_DIM_ALLOC];
 	int scale_stride[CCV_NNC_MAX_DIM_ALLOC];
 	ccv_nnc_tensor_view_get_stride(a, astride);
-	ccv_nnc_tensor_view_get_stride(scale, scale_stride);
+	if (scale)
+		ccv_nnc_tensor_view_get_stride(scale, scale_stride);
 	ccv_nnc_tensor_view_get_stride(b, bstride);
 	// The epsilon is used a little bit differently from batch norm, it is outside of the sqrt in this case.
 	const float epsilon = cmd.info.rmsnorm.epsilon;
@@ -91,36 +92,66 @@ static int _ccv_nnc_rmsnorm_forw(const ccv_nnc_cmd_t cmd, const ccv_nnc_hint_t h
 			}
 		}
 	}
-	float* const scalep = scale->data.f32;
-	int sdim[CCV_NNC_MAX_DIM_ALLOC];
-	ccv_nnc_tensor_view_get_dim(scale, sdim);
-	// Do the straight-forward one, y = x * inv_std * scale + bias, we cannot allocate extra memory to help.
-	// There is no need for precompute since scale / bias is per element.
-	float* const bp = b->data.f32;
-	for (i[0] = 0; i[0] < adim[0]; i[0]++)
+	if (cmd.info.rmsnorm.elementwise_affine)
 	{
-		float* const ap0 = ap + i[0] * astride[0];
-		float* const bp0 = bp + i[0] * bstride[0];
-		float* const varp0 = rdim[0] == 1 ? varp : varp + i[0] * saved_inv_std_stride[0];
-		float* const scalep0 = sdim[0] == 1 ? scalep : scalep + i[0] * scale_stride[0];
-		for (i[1] = 0; i[1] < adim[1]; i[1]++)
+		float* const scalep = scale->data.f32;
+		int sdim[CCV_NNC_MAX_DIM_ALLOC];
+		ccv_nnc_tensor_view_get_dim(scale, sdim);
+		// Do the straight-forward one, y = x * inv_std * scale + bias, we cannot allocate extra memory to help.
+		// There is no need for precompute since scale / bias is per element.
+		float* const bp = b->data.f32;
+		for (i[0] = 0; i[0] < adim[0]; i[0]++)
 		{
-			float* ap1 = ap0 + i[1] * astride[1];
-			float* bp1 = bp0 + i[1] * bstride[1];
-			float* const varp1 = rdim[1] == 1 ? varp0 : varp0 + i[1] * saved_inv_std_stride[1];
-			float* const scalep1 = sdim[1] == 1 ? scalep0 : scalep0 + i[1] * scale_stride[1];
-			for (i[2] = 0; i[2] < adim[2]; i[2]++)
+			float* const ap0 = ap + i[0] * astride[0];
+			float* const bp0 = bp + i[0] * bstride[0];
+			float* const varp0 = rdim[0] == 1 ? varp : varp + i[0] * saved_inv_std_stride[0];
+			float* const scalep0 = sdim[0] == 1 ? scalep : scalep + i[0] * scale_stride[0];
+			for (i[1] = 0; i[1] < adim[1]; i[1]++)
 			{
-				float* const varp2 = rdim[2] == 1 ? varp1 : varp1 + i[2] * saved_inv_std_stride[2];
-				float* const scalep2 = sdim[2] == 1 ? scalep1 : scalep1 + i[2] * scale_stride[2];
-				if (rdim[3] == 1)
-					for (x = 0; x < adim[3]; x++)
-						bp1[x] = ap1[x * astride[3]] * varp2[0] * scalep2[sdim[3] == 1 ? 0 : x];
-				else
-					for (x = 0; x < adim[3]; x++)
-						bp1[x] = ap1[x * astride[3]] * varp2[x] * scalep2[sdim[3] == 1 ? 0 : x];
-				ap1 += astride[2];
-				bp1 += bstride[2];
+				float* ap1 = ap0 + i[1] * astride[1];
+				float* bp1 = bp0 + i[1] * bstride[1];
+				float* const varp1 = rdim[1] == 1 ? varp0 : varp0 + i[1] * saved_inv_std_stride[1];
+				float* const scalep1 = sdim[1] == 1 ? scalep0 : scalep0 + i[1] * scale_stride[1];
+				for (i[2] = 0; i[2] < adim[2]; i[2]++)
+				{
+					float* const varp2 = rdim[2] == 1 ? varp1 : varp1 + i[2] * saved_inv_std_stride[2];
+					float* const scalep2 = sdim[2] == 1 ? scalep1 : scalep1 + i[2] * scale_stride[2];
+					if (rdim[3] == 1)
+						for (x = 0; x < adim[3]; x++)
+							bp1[x] = ap1[x * astride[3]] * varp2[0] * scalep2[sdim[3] == 1 ? 0 : x];
+					else
+						for (x = 0; x < adim[3]; x++)
+							bp1[x] = ap1[x * astride[3]] * varp2[x] * scalep2[sdim[3] == 1 ? 0 : x];
+					ap1 += astride[2];
+					bp1 += bstride[2];
+				}
+			}
+		}
+	} else {
+		// Do the straight-forward one, y = x * inv_std, we cannot allocate extra memory to help.
+		float* const bp = b->data.f32;
+		for (i[0] = 0; i[0] < adim[0]; i[0]++)
+		{
+			float* const ap0 = ap + i[0] * astride[0];
+			float* const bp0 = bp + i[0] * bstride[0];
+			float* const varp0 = rdim[0] == 1 ? varp : varp + i[0] * saved_inv_std_stride[0];
+			for (i[1] = 0; i[1] < adim[1]; i[1]++)
+			{
+				float* ap1 = ap0 + i[1] * astride[1];
+				float* bp1 = bp0 + i[1] * bstride[1];
+				float* const varp1 = rdim[1] == 1 ? varp0 : varp0 + i[1] * saved_inv_std_stride[1];
+				for (i[2] = 0; i[2] < adim[2]; i[2]++)
+				{
+					float* const varp2 = rdim[2] == 1 ? varp1 : varp1 + i[2] * saved_inv_std_stride[2];
+					if (rdim[3] == 1)
+						for (x = 0; x < adim[3]; x++)
+							bp1[x] = ap1[x * astride[3]] * varp2[0];
+					else
+						for (x = 0; x < adim[3]; x++)
+							bp1[x] = ap1[x * astride[3]] * varp2[x];
+					ap1 += astride[2];
+					bp1 += bstride[2];
+				}
 			}
 		}
 	}
@@ -129,12 +160,13 @@ static int _ccv_nnc_rmsnorm_forw(const ccv_nnc_cmd_t cmd, const ccv_nnc_hint_t h
 
 static int _ccv_nnc_rmsnorm_back(const ccv_nnc_cmd_t cmd, const ccv_nnc_hint_t hint, const int flags, ccv_nnc_tensor_t* const* const inputs, const int input_size, ccv_nnc_tensor_t* const* const outputs, const int output_size, ccv_nnc_stream_context_t* const stream_context)
 {
-	assert(input_size == 6);
+	assert(input_size == 6 || input_size == 5);
 	assert(output_size >= 1);
+	const int elementwise_affine = cmd.info.rmsnorm.elementwise_affine;
 	ccv_nnc_tensor_view_t* const g = (ccv_nnc_tensor_view_t*)inputs[0];
 	ccv_nnc_tensor_view_t* const a = (ccv_nnc_tensor_view_t*)inputs[2];
-	ccv_nnc_tensor_view_t* const scale = (ccv_nnc_tensor_view_t*)inputs[3];
-	ccv_nnc_tensor_view_t* const saved_inv_std = (ccv_nnc_tensor_view_t*)inputs[5];
+	ccv_nnc_tensor_view_t* const scale = elementwise_affine ? (ccv_nnc_tensor_view_t*)inputs[3] : 0;
+	ccv_nnc_tensor_view_t* const saved_inv_std = (ccv_nnc_tensor_view_t*)inputs[elementwise_affine ? 5 : 4];
 	ccv_nnc_tensor_view_t* const h = (ccv_nnc_tensor_view_t*)outputs[0];
 	ccv_nnc_tensor_view_t* const dscale = output_size > 1 ? (ccv_nnc_tensor_view_t*)outputs[1] : 0;
 	assert(ccv_nnc_tensor_nd(g->info.dim) <= CCV_NNC_MAX_DIM + 2);
@@ -146,7 +178,8 @@ static int _ccv_nnc_rmsnorm_back(const ccv_nnc_cmd_t cmd, const ccv_nnc_hint_t h
 	ccv_nnc_tensor_view_get_dim(g, gdim);
 	ccv_nnc_tensor_view_get_dim(saved_inv_std, rdim);
 	int sdim[CCV_NNC_MAX_DIM_ALLOC];
-	ccv_nnc_tensor_view_get_dim(scale, sdim);
+	if (scale)
+		ccv_nnc_tensor_view_get_dim(scale, sdim);
 	if (dscale)
 		{ assert(ccv_nnc_tensor_view_check_dim(dscale, sdim)); }
 	assert(ccv_nnc_tensor_view_check_dim(a, gdim));
@@ -160,7 +193,8 @@ static int _ccv_nnc_rmsnorm_back(const ccv_nnc_cmd_t cmd, const ccv_nnc_hint_t h
 	ccv_nnc_tensor_view_get_stride(a, astride);
 	ccv_nnc_tensor_view_get_stride(g, gstride);
 	ccv_nnc_tensor_view_get_stride(h, hstride);
-	ccv_nnc_tensor_view_get_stride(scale, scale_stride);
+	if (scale)
+		ccv_nnc_tensor_view_get_stride(scale, scale_stride);
 	ccv_nnc_tensor_view_get_stride(saved_inv_std, inv_std_stride);
 	if (dscale)
 		ccv_nnc_tensor_view_get_stride(dscale, dscale_stride);
@@ -252,29 +286,55 @@ static int _ccv_nnc_rmsnorm_back(const ccv_nnc_cmd_t cmd, const ccv_nnc_hint_t h
 	} else {
 		float* gssp = gss;
 		const float* const gp = g->data.f32;
-		const float* const scalep = scale->data.f32;
-		for (i[0] = 0; i[0] < gdim[0]; i[0]++)
+		if (elementwise_affine)
 		{
-			const float* const gp0 = gp + i[0] * gstride[0];
-			const float* const inv_stdp0 = rdim[0] == 1 ? inv_stdp : inv_stdp + i[0] * inv_std_stride[0];
-			const float* const scalep0 = sdim[0] == 1 ? scalep : scalep + i[0] * scale_stride[0];
-			for (i[1] = 0; i[1] < gdim[1]; i[1]++)
+			const float* const scalep = scale->data.f32;
+			for (i[0] = 0; i[0] < gdim[0]; i[0]++)
 			{
-				const float* gp1 = gp0 + i[1] * gstride[1];
-				const float* const inv_stdp1 = rdim[1] == 1 ? inv_stdp0 : inv_stdp0 + i[1] * inv_std_stride[1];
-				const float* const scalep1 = sdim[1] == 1 ? scalep0 : scalep0 + i[1] * scale_stride[1];
-				for (i[2] = 0; i[2] < gdim[2]; i[2]++)
+				const float* const gp0 = gp + i[0] * gstride[0];
+				const float* const inv_stdp0 = rdim[0] == 1 ? inv_stdp : inv_stdp + i[0] * inv_std_stride[0];
+				const float* const scalep0 = sdim[0] == 1 ? scalep : scalep + i[0] * scale_stride[0];
+				for (i[1] = 0; i[1] < gdim[1]; i[1]++)
 				{
-					const float* const inv_stdp2 = rdim[2] == 1 ? inv_stdp1 : inv_stdp1 + i[2] * inv_std_stride[2];
-					const float* const scalep2 = sdim[2] == 1 ? scalep1 : scalep1 + i[2] * scale_stride[2];
-					if (sdim[3] == 1)
-						for (x = 0; x < gdim[3]; x++)
-							gssp[x] = gp1[x] * scalep2[0] * inv_stdp2[rdim[3] == 1 ? 0 : x];
-					else
-						for (x = 0; x < gdim[3]; x++)
-							gssp[x] = gp1[x] * scalep2[x] * inv_stdp2[rdim[3] == 1 ? 0 : x];
-					gp1 += gstride[2];
-					gssp += gdim[3];
+					const float* gp1 = gp0 + i[1] * gstride[1];
+					const float* const inv_stdp1 = rdim[1] == 1 ? inv_stdp0 : inv_stdp0 + i[1] * inv_std_stride[1];
+					const float* const scalep1 = sdim[1] == 1 ? scalep0 : scalep0 + i[1] * scale_stride[1];
+					for (i[2] = 0; i[2] < gdim[2]; i[2]++)
+					{
+						const float* const inv_stdp2 = rdim[2] == 1 ? inv_stdp1 : inv_stdp1 + i[2] * inv_std_stride[2];
+						const float* const scalep2 = sdim[2] == 1 ? scalep1 : scalep1 + i[2] * scale_stride[2];
+						if (sdim[3] == 1)
+							for (x = 0; x < gdim[3]; x++)
+								gssp[x] = gp1[x] * scalep2[0] * inv_stdp2[rdim[3] == 1 ? 0 : x];
+						else
+							for (x = 0; x < gdim[3]; x++)
+								gssp[x] = gp1[x] * scalep2[x] * inv_stdp2[rdim[3] == 1 ? 0 : x];
+						gp1 += gstride[2];
+						gssp += gdim[3];
+					}
+				}
+			}
+		} else {
+			for (i[0] = 0; i[0] < gdim[0]; i[0]++)
+			{
+				const float* const gp0 = gp + i[0] * gstride[0];
+				const float* const inv_stdp0 = rdim[0] == 1 ? inv_stdp : inv_stdp + i[0] * inv_std_stride[0];
+				for (i[1] = 0; i[1] < gdim[1]; i[1]++)
+				{
+					const float* gp1 = gp0 + i[1] * gstride[1];
+					const float* const inv_stdp1 = rdim[1] == 1 ? inv_stdp0 : inv_stdp0 + i[1] * inv_std_stride[1];
+					for (i[2] = 0; i[2] < gdim[2]; i[2]++)
+					{
+						const float* const inv_stdp2 = rdim[2] == 1 ? inv_stdp1 : inv_stdp1 + i[2] * inv_std_stride[2];
+						if (rdim[3] == 1)
+							for (x = 0; x < gdim[3]; x++)
+								gssp[x] = gp1[x] * inv_stdp2[0];
+						else
+							for (x = 0; x < gdim[3]; x++)
+								gssp[x] = gp1[x] * inv_stdp2[x];
+						gp1 += gstride[2];
+						gssp += gdim[3];
+					}
 				}
 			}
 		}