merge 2nd order derivative of CutlassMLP

include/tiny-cuda-nn/encodings/identity.h

@@ -60,7 +60,7 @@ __global__ void identity(
  const uint32_t j = encoded_index - i * fan_out;
 
  if (j >= num_to_encode) {
- data_out(j, i) = 1;
+ data_out(j, i) = 0; // data_out(j, i) = 0;
  } else {
  data_out(j, i) = data_in(j, i) * scale + offset;
  }
@@ -84,6 +84,25 @@ __global__ void identity_backward(
  dL_dx(j, i) = (T)((float)dL_dy(j, i) * scale);
 }
 
+template <typename T>
+__global__ void identity_backward_backward(
+ const uint32_t num_outputs,
+ const uint32_t num_elements,
+ const uint32_t n_dims_to_encode,
+ const float scale,
+ MatrixView<const float> dL_ddLdy,
+ MatrixView<T> dL_ddLdx)
+{
+ const uint32_t output_index = threadIdx.x + blockIdx.x * blockDim.x;
+ if (output_index >= num_outputs) return;
+
+ const uint32_t i = output_index / n_dims_to_encode;
+ const uint32_t j = output_index - i * n_dims_to_encode;
+
+ // The identity encoding can simply pass through the derivative.
+ dL_ddLdx(j, i) = (T)(dL_ddLdy(j, i) * scale);
+}
+
 template <typename T>
 class IdentityEncoding : public Encoding<T> {
 public:
@@ -139,6 +158,33 @@ class IdentityEncoding : public Encoding<T> {
  );
  }
 
+ void backward_backward_input_impl(
+ cudaStream_t stream,
+ const Context& ctx,
+ const GPUMatrixDynamic<float>& input,
+ const GPUMatrixDynamic<float>& dL_ddLdinput,
+ const GPUMatrixDynamic<T>& dL_doutput,
+ GPUMatrixDynamic<T>* dL_ddLdoutput = nullptr,
+ GPUMatrixDynamic<float>* dL_dinput = nullptr,
+ bool use_inference_params = false,
+ GradientMode param_gradients_mode = GradientMode::Overwrite
+ ) override {
+ if (!dL_dinput || !dL_ddLdoutput || padded_output_width() == 0) {
+ return;
+ }
+
+ linear_kernel(identity_backward_backward<T>, 0, stream,
+ input.n() * m_n_dims_to_encode,
+ input.n(),
+ m_n_dims_to_encode,
+ m_scale,
+ dL_ddLdinput.view(),
+ dL_ddLdoutput->view()
+ );
+
+ // dL_dinput: don't need to calculate this term, it's default set as 0.0
+ }
+
  uint32_t input_width() const override {
  return m_n_dims_to_encode;
  }

include/tiny-cuda-nn/network_with_input_encoding.h

@@ -38,6 +38,33 @@
 
 namespace tcnn {
 
+// element-wise convert float* to T*
+template <typename T>
+__global__ void element_wise_convert(uint32_t n_elements, float* in, T* out) {
+ uint32_t i = blockIdx.x * blockDim.x + threadIdx.x;
+ if (i >= n_elements) return;
+
+ out[i] = (T)in[i];
+}
+
+// element-wise convert T* to float* and then add back to *out
+template <typename T>
+__global__ void element_wise_convert_float(uint32_t n_elements, T* in, float* out) {
+ uint32_t i = blockIdx.x * blockDim.x + threadIdx.x;
+ if (i >= n_elements) return;
+
+ out[i] += (float)in[i];
+}
+
+// element-wise add
+template <typename T>
+__global__ void element_wise_add(uint32_t n_elements, T* in, T* out) {
+ uint32_t i = blockIdx.x * blockDim.x + threadIdx.x;
+ if (i >= n_elements) return;
+
+ out[i] += in[i];
+}
+
 template <typename T>
 class NetworkWithInputEncoding : public Network<float, T> {
 public:
@@ -90,8 +117,8 @@ class NetworkWithInputEncoding : public Network<float, T> {
  bool use_inference_params = false,
  GradientMode param_gradients_mode = GradientMode::Overwrite
  ) override {
- GPUMatrixDynamic<T> dL_dnetwork_input;
  if (m_encoding->n_params() > 0 || dL_dinput) {
+ // dL_dnetwork_input becomes a member of the class instance
  dL_dnetwork_input = {m_encoding->padded_output_width(), input.n(), stream, m_encoding->preferred_output_layout()};
  }
 
@@ -112,6 +139,89 @@ class NetworkWithInputEncoding : public Network<float, T> {
  }
  }
 
+ void backward_backward_input_impl(
+ cudaStream_t stream,
+ const Context& ctx,
+ const GPUMatrixDynamic<float>& input,
+ const GPUMatrixDynamic<float>& dL_ddLdinput,
+ const GPUMatrixDynamic<T>& dL_doutput,
+ GPUMatrixDynamic<T>* dL_ddLdoutput = nullptr,
+ GPUMatrixDynamic<float>* dL_dinput = nullptr,
+ bool use_inference_params = false,
+ GradientMode param_gradients_mode = GradientMode::Overwrite
+ ) override {
+ const auto& forward = dynamic_cast<const ForwardContext&>(ctx);
+
+ // dL_ddLdinput of m_network->backward_baward_input equals to dL_dLdencoding_output (different names)
+ GPUMatrixDynamic<T> dL_dLdnetwork_input;
+
+ if (m_encoding->n_params() > 0) {
+ dL_dLdnetwork_input = {m_encoding->padded_output_width(), input.n(), stream, dL_ddLdinput.layout()};
+ // cudaMemsetAsync: set dL_dLdnetwork_input.data() with 0.0 to avoid NaN initialization
+ CUDA_CHECK_THROW(cudaMemsetAsync(dL_dLdnetwork_input.data(), 0, dL_dLdnetwork_input.n() * dL_dLdnetwork_input.m() * sizeof(T), stream));
+
+ // encoding backward backward
+ m_encoding->backward_backward_input(
+ stream, 
+ *forward.encoding_ctx, 
+ input,
+ dL_ddLdinput,
+ dL_dnetwork_input, // dL1_denc_output
+ &dL_dLdnetwork_input, // dL2_ddL1_denc_output
+ dL_dinput,
+ use_inference_params,
+ param_gradients_mode
+ );
+ } else { // copy dL_ddLdinput (float) to dL_dLdnetwork_input (T)
+ dL_dLdnetwork_input = {m_encoding->padded_output_width(), input.n(), stream, dL_ddLdinput.layout()};
+ linear_kernel(element_wise_convert<T>, 0, stream, dL_dLdnetwork_input.n() * dL_dLdnetwork_input.m(), dL_ddLdinput.data(), dL_dLdnetwork_input.data());
+ }
+
+ // dL2_dinput of m_network->backward_backward_input
+ GPUMatrixDynamic<T> dL2_dnetwork_input;
+ if (m_encoding->n_params() > 0 || dL_dinput) {
+ dL2_dnetwork_input = {m_encoding->padded_output_width(), input.n(), stream, m_encoding->preferred_output_layout()};
+ }
+
+ // network backward backward
+ m_network->backward_backward_input(
+ stream,
+ *forward.network_ctx,
+ forward.network_input, // enc_output i.e. network_input
+ dL_dLdnetwork_input, // dL2_dL1dnetwork_input
+ dL_doutput,
+ dL_ddLdoutput ? dL_ddLdoutput : nullptr,
+ dL2_dnetwork_input.data() ? &dL2_dnetwork_input : nullptr, // dL2_dinput of network
+ use_inference_params,
+ param_gradients_mode
+ );
+
+ // dL2dnetwork_input backward to dL2dinput, first order backward
+ GPUMatrixDynamic<float> dL2_dinput;
+ if (m_encoding->n_params() > 0 || dL2_dnetwork_input.data()) {
+ dL2_dinput = {m_encoding->input_width(), input.n(), stream, input.layout()};
+ }
+
+ if (m_encoding->n_params() > 0) {
+ // backward dL2dnetwork_input to dL2dinput
+ m_encoding->backward(
+ stream,
+ *forward.encoding_ctx,
+ input,
+ forward.network_input, // enc_output
+ dL2_dnetwork_input, // dL2_dencoding_output
+ &dL2_dinput,
+ use_inference_params,
+ GradientMode::Accumulate // dL2denc_w : add up 1st order term
+ );
+
+ linear_kernel(element_wise_add<float>, 0, stream, dL_dinput->n() * dL_dinput->m(), dL2_dinput.data(), dL_dinput->data());
+
+ } else if (dL2_dnetwork_input.data()) {
+ linear_kernel(element_wise_convert_float<T>, 0, stream, dL_dinput->n() * dL_dinput->m(), dL2_dnetwork_input.data(), dL_dinput->data());
+ }
+ }
+
  void set_params_impl(T* params, T* inference_params, T* gradients) override {
  size_t offset = 0;
  m_network->set_params(params + offset, inference_params + offset, gradients + offset);
@@ -181,6 +291,7 @@ class NetworkWithInputEncoding : public Network<float, T> {
 private:
  std::shared_ptr<Encoding<T>> m_encoding;
  std::shared_ptr<Network<T>> m_network;
+ GPUMatrixDynamic<T> dL_dnetwork_input;
 
  struct ForwardContext : public Context {
  GPUMatrixDynamic<T> network_input;

include/tiny-cuda-nn/networks/cutlass_mlp.h

@@ -66,6 +66,28 @@ class CutlassMLP : public Network<T> {
  GradientMode param_gradients_mode = GradientMode::Overwrite
  ) override;
 
+ void backward_backward_input_impl(
+ cudaStream_t stream,
+ const Context& ctx,
+ const GPUMatrixDynamic<T>& input,
+ const GPUMatrixDynamic<T>& dL_ddLdinput,
+ const GPUMatrixDynamic<T>& dL_doutput,
+ GPUMatrixDynamic<T>* dL_ddLdoutput = nullptr,
+ GPUMatrixDynamic<T>* dL_dinput = nullptr,
+ bool use_inference_params = false,
+ GradientMode param_gradients_mode = GradientMode::Overwrite
+ ) override;
+
+ bool prepare_backward_variables(
+ cudaStream_t stream,
+ const std::vector<GPUMatrix<T>>& output,
+ const GPUMatrixDynamic<T>& dL_doutput,
+ GPUMatrixDynamic<T>& backward_output_tmp,
+ std::vector<GPUMatrix<T>>& dL1dp,
+ std::vector<GPUMatrix<T>>& dL1doutput,
+ bool use_inference_params
+ );
+
  void set_params_impl(T* params, T* inference_params, T* gradients) override;
  void initialize_params(pcg32& rnd, float* params_full_precision, float scale = 1) override;