Add files using upload-large-folder tool

.gitattributes

@@ -331,3 +331,8 @@ my_container_sandbox/workspace/anaconda3/lib/python3.8/site-packages/opencv_pyth
 my_container_sandbox/workspace/anaconda3/lib/python3.8/site-packages/pyzmq.libs/libsodium-bcf9f097.so.23.3.0 filter=lfs diff=lfs merge=lfs -text
 my_container_sandbox/workspace/anaconda3/lib/python3.8/site-packages/yarl/_quoting_c.cpython-38-x86_64-linux-gnu.so filter=lfs diff=lfs merge=lfs -text
 my_container_sandbox/workspace/anaconda3/lib/python3.8/site-packages/h5py.libs/libaec-9c9e97eb.so.0.0.10 filter=lfs diff=lfs merge=lfs -text
+my_container_sandbox/workspace/anaconda3/lib/python3.8/site-packages/opencv_python.libs/libQt5Test-c38a5234.so.5.15.0 filter=lfs diff=lfs merge=lfs -text
+my_container_sandbox/workspace/anaconda3/lib/python3.8/site-packages/scipy.libs/libgfortran-040039e1.so.5.0.0 filter=lfs diff=lfs merge=lfs -text
+my_container_sandbox/workspace/anaconda3/lib/python3.8/site-packages/brotli/_brotli.abi3.so filter=lfs diff=lfs merge=lfs -text
+my_container_sandbox/workspace/anaconda3/lib/python3.8/site-packages/opencv_python.libs/libxcb-xkb-9ba31ab3.so.1.0.0 filter=lfs diff=lfs merge=lfs -text
+my_container_sandbox/workspace/anaconda3/lib/python3.8/site-packages/opencv_python.libs/libssl-28bef1ac.so.1.1 filter=lfs diff=lfs merge=lfs -text

my_container_sandbox/workspace/anaconda3/lib/python3.8/site-packages/brotli/_brotli.abi3.so

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1c9ea7e74f258c0527249f553bdd1a136e016017a222f38186df92e85361a4f0
+size 746208

my_container_sandbox/workspace/anaconda3/lib/python3.8/site-packages/monai/_extensions/gmm/gmm.cpp

@@ -0,0 +1,85 @@
+/*
+Copyright (c) MONAI Consortium
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+    http://www.apache.org/licenses/LICENSE-2.0
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+#include <torch/extension.h>
+
+#include "gmm.h"
+
+py::tuple init() {
+  torch::Tensor gmm_tensor =
+      torch::zeros({GMM_COUNT, GMM_COMPONENT_COUNT}, torch::dtype(torch::kFloat32).device(torch::kCUDA));
+  torch::Tensor scratch_tensor = torch::empty({1}, torch::dtype(torch::kFloat32).device(torch::kCUDA));
+  return py::make_tuple(gmm_tensor, scratch_tensor);
+}
+
+void learn(
+    torch::Tensor gmm_tensor,
+    torch::Tensor scratch_tensor,
+    torch::Tensor input_tensor,
+    torch::Tensor label_tensor) {
+  c10::DeviceType device_type = input_tensor.device().type();
+
+  unsigned int batch_count = input_tensor.size(0);
+  unsigned int element_count = input_tensor.stride(1);
+
+  unsigned int scratch_size =
+      batch_count * (element_count + GMM_COMPONENT_COUNT * GMM_COUNT * (element_count / (32 * 32)));
+
+  if (scratch_tensor.size(0) < scratch_size) {
+    scratch_tensor.resize_({scratch_size});
+  }
+
+  float* gmm = gmm_tensor.data_ptr<float>();
+  float* scratch = scratch_tensor.data_ptr<float>();
+  float* input = input_tensor.data_ptr<float>();
+  int* labels = label_tensor.data_ptr<int>();
+
+  if (device_type == torch::kCUDA) {
+    learn_cuda(input, labels, gmm, scratch, batch_count, element_count);
+  } else {
+    learn_cpu(input, labels, gmm, scratch, batch_count, element_count);
+  }
+}
+
+torch::Tensor apply(torch::Tensor gmm_tensor, torch::Tensor input_tensor) {
+  c10::DeviceType device_type = input_tensor.device().type();
+
+  unsigned int dim = input_tensor.dim();
+  unsigned int batch_count = input_tensor.size(0);
+  unsigned int element_count = input_tensor.stride(1);
+
+  long int* output_size = new long int[dim];
+  memcpy(output_size, input_tensor.sizes().data(), dim * sizeof(long int));
+  output_size[1] = MIXTURE_COUNT;
+  torch::Tensor output_tensor =
+      torch::empty(c10::IntArrayRef(output_size, dim), torch::dtype(torch::kFloat32).device(device_type));
+  delete output_size;
+
+  const float* gmm = gmm_tensor.data_ptr<float>();
+  const float* input = input_tensor.data_ptr<float>();
+  float* output = output_tensor.data_ptr<float>();
+
+  if (device_type == torch::kCUDA) {
+    apply_cuda(gmm, input, output, batch_count, element_count);
+  } else {
+    apply_cpu(gmm, input, output, batch_count, element_count);
+  }
+
+  return output_tensor;
+}
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+  m.def("init", torch::wrap_pybind_function(init));
+  m.def("learn", torch::wrap_pybind_function(learn));
+  m.def("apply", torch::wrap_pybind_function(apply));
+}

my_container_sandbox/workspace/anaconda3/lib/python3.8/site-packages/monai/_extensions/gmm/gmm.h

@@ -0,0 +1,53 @@
+/*
+Copyright (c) MONAI Consortium
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+    http://www.apache.org/licenses/LICENSE-2.0
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+#if !defined(CHANNEL_COUNT) || !defined(MIXTURE_COUNT) || !defined(MIXTURE_SIZE)
+#error Definition of CHANNEL_COUNT, MIXTURE_COUNT, and MIXTURE_SIZE required
+#endif
+
+#if CHANNEL_COUNT < 1 || MIXTURE_COUNT < 1 || MIXTURE_SIZE < 1
+#error CHANNEL_COUNT, MIXTURE_COUNT, and MIXTURE_SIZE must be positive
+#endif
+
+#define MATRIX_COMPONENT_COUNT ((CHANNEL_COUNT + 1) * (CHANNEL_COUNT + 2) / 2)
+#define SUB_MATRIX_COMPONENT_COUNT (CHANNEL_COUNT * (CHANNEL_COUNT + 1) / 2)
+#define GMM_COMPONENT_COUNT (MATRIX_COMPONENT_COUNT + 1)
+#define GMM_COUNT (MIXTURE_COUNT * MIXTURE_SIZE)
+
+void learn_cpu(
+    const float* input,
+    const int* labels,
+    float* gmm,
+    float* scratch_memory,
+    unsigned int batch_count,
+    unsigned int element_count);
+void apply_cpu(
+    const float* gmm,
+    const float* input,
+    float* output,
+    unsigned int batch_count,
+    unsigned int element_count);
+
+void learn_cuda(
+    const float* input,
+    const int* labels,
+    float* gmm,
+    float* scratch_memory,
+    unsigned int batch_count,
+    unsigned int element_count);
+void apply_cuda(
+    const float* gmm,
+    const float* input,
+    float* output,
+    unsigned int batch_count,
+    unsigned int element_count);

my_container_sandbox/workspace/anaconda3/lib/python3.8/site-packages/monai/_extensions/gmm/gmm_cpu.cpp

@@ -0,0 +1,35 @@
+/*
+Copyright (c) MONAI Consortium
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+    http://www.apache.org/licenses/LICENSE-2.0
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+#include <stdexcept>
+
+#include "gmm.h"
+
+void learn_cpu(
+    const float* input,
+    const int* labels,
+    float* gmm,
+    float* scratch_memory,
+    unsigned int batch_count,
+    unsigned int element_count) {
+  throw std::invalid_argument("GMM received a cpu tensor but is not yet implemented for the cpu");
+}
+
+void apply_cpu(
+    const float* gmm,
+    const float* input,
+    float* output,
+    unsigned int batch_count,
+    unsigned int element_count) {
+  throw std::invalid_argument("GMM received a cpu tensor but is not yet implemented for the cpu");
+}

my_container_sandbox/workspace/anaconda3/lib/python3.8/site-packages/monai/_extensions/gmm/gmm_cuda.cu

@@ -0,0 +1,518 @@
+/*
+Copyright (c) MONAI Consortium
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+    http://www.apache.org/licenses/LICENSE-2.0
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+#include <cuda.h>
+#include <cuda_runtime.h>
+
+#include "gmm.h"
+
+#include "gmm_cuda_linalg.cuh"
+
+#define EPSILON 1e-5
+#define BLOCK_SIZE 32
+#define TILE(SIZE, STRIDE) ((((SIZE)-1) / (STRIDE)) + 1)
+
+template <int warp_count, int load_count>
+__global__ void CovarianceReductionKernel(
+    int gaussian_index,
+    const float* g_image,
+    const int* g_alpha,
+    float* g_matrices,
+    int element_count) {
+  constexpr int block_size = warp_count * 32;
+
+  __shared__ float s_matrix_component[warp_count];
+
+  int batch_index = blockIdx.z;
+
+  const float* g_batch_image = g_image + batch_index * element_count * CHANNEL_COUNT;
+  const int* g_batch_alpha = g_alpha + batch_index * element_count;
+  float* g_batch_matrices = g_matrices + batch_index * GMM_COUNT * GMM_COMPONENT_COUNT * gridDim.x;
+
+  int local_index = threadIdx.x;
+  int block_index = blockIdx.x;
+  int warp_index = local_index >> 5;
+  int lane_index = local_index & 31;
+  int global_index = local_index + block_index * block_size * load_count;
+  int matrix_offset = (gaussian_index * gridDim.x + block_index) * GMM_COMPONENT_COUNT;
+
+  float matrix[MATRIX_COMPONENT_COUNT];
+
+  for (int i = 0; i < MATRIX_COMPONENT_COUNT; i++) {
+    matrix[i] = 0;
+  }
+
+  for (int load = 0; load < load_count; load++) {
+    global_index += load * block_size;
+
+    if (global_index < element_count) {
+      int my_alpha = g_batch_alpha[global_index];
+
+      if (my_alpha != -1) {
+        if (gaussian_index == (my_alpha & 15) + (my_alpha >> 4) * MIXTURE_COUNT) {
+          float feature[CHANNEL_COUNT + 1];
+
+          feature[0] = 1;
+
+          for (int i = 0; i < CHANNEL_COUNT; i++) {
+            feature[i + 1] = g_batch_image[global_index + i * element_count];
+          }
+
+          for (int index = 0, i = 0; i < CHANNEL_COUNT + 1; i++) {
+            for (int j = i; j < CHANNEL_COUNT + 1; j++, index++) {
+              matrix[index] += feature[i] * feature[j];
+            }
+          }
+        }
+      }
+    }
+  }
+
+  __syncthreads();
+
+  for (int i = 0; i < MATRIX_COMPONENT_COUNT; i++) {
+    float matrix_component = matrix[i];
+
+    matrix_component += __shfl_down_sync(0xffffffff, matrix_component, 16);
+    matrix_component += __shfl_down_sync(0xffffffff, matrix_component, 8);
+    matrix_component += __shfl_down_sync(0xffffffff, matrix_component, 4);
+    matrix_component += __shfl_down_sync(0xffffffff, matrix_component, 2);
+    matrix_component += __shfl_down_sync(0xffffffff, matrix_component, 1);
+
+    if (lane_index == 0) {
+      s_matrix_component[warp_index] = matrix_component;
+    }
+
+    __syncthreads();
+
+    if (warp_index == 0) {
+      matrix_component = s_matrix_component[lane_index];
+
+      if (warp_count >= 32) {
+        matrix_component += __shfl_down_sync(0xffffffff, matrix_component, 16);
+      }
+      if (warp_count >= 16) {
+        matrix_component += __shfl_down_sync(0xffffffff, matrix_component, 8);
+      }
+      if (warp_count >= 8) {
+        matrix_component += __shfl_down_sync(0xffffffff, matrix_component, 4);
+      }
+      if (warp_count >= 4) {
+        matrix_component += __shfl_down_sync(0xffffffff, matrix_component, 2);
+      }
+      if (warp_count >= 2) {
+        matrix_component += __shfl_down_sync(0xffffffff, matrix_component, 1);
+      }
+
+      if (lane_index == 0) {
+        g_batch_matrices[matrix_offset + i] = matrix_component;
+      }
+    }
+
+    __syncthreads();
+  }
+}
+
+template <int warp_count, bool invert_matrix>
+__global__ void CovarianceFinalizationKernel(const float* g_matrices, float* g_gmm, int matrix_count) {
+  constexpr int block_size = warp_count * 32;
+
+  __shared__ float s_matrix_component[warp_count];
+  __shared__ float s_gmm[GMM_COMPONENT_COUNT];
+
+  int batch_index = blockIdx.z;
+
+  const float* g_batch_matrices = g_matrices + batch_index * GMM_COUNT * GMM_COMPONENT_COUNT * matrix_count;
+  float* g_batch_gmm = g_gmm + batch_index * GMM_COUNT * GMM_COMPONENT_COUNT;
+
+  int local_index = threadIdx.x;
+  int warp_index = local_index >> 5;
+  int lane_index = local_index & 31;
+  int gmm_index = blockIdx.x;
+  int matrix_offset = gmm_index * matrix_count;
+
+  int load_count = TILE(matrix_count, block_size);
+
+  float norm_factor = 1.0f;
+
+  for (int index = 0, i = 0; i < CHANNEL_COUNT + 1; i++) {
+    for (int j = i; j < CHANNEL_COUNT + 1; j++, index++) {
+      float matrix_component = 0.0f;
+
+      for (int load = 0; load < load_count; load++) {
+        int matrix_index = local_index + load * block_size;
+
+        if (matrix_index < matrix_count) {
+          matrix_component += g_batch_matrices[(matrix_offset + matrix_index) * GMM_COMPONENT_COUNT + index];
+        }
+      }
+
+      matrix_component += __shfl_down_sync(0xffffffff, matrix_component, 16);
+      matrix_component += __shfl_down_sync(0xffffffff, matrix_component, 8);
+      matrix_component += __shfl_down_sync(0xffffffff, matrix_component, 4);
+      matrix_component += __shfl_down_sync(0xffffffff, matrix_component, 2);
+      matrix_component += __shfl_down_sync(0xffffffff, matrix_component, 1);
+
+      if (lane_index == 0) {
+        s_matrix_component[warp_index] = matrix_component;
+      }
+
+      __syncthreads();
+
+      if (warp_index == 0) {
+        matrix_component = s_matrix_component[lane_index];
+
+        if (warp_count >= 32) {
+          matrix_component += __shfl_down_sync(0xffffffff, matrix_component, 16);
+        }
+        if (warp_count >= 16) {
+          matrix_component += __shfl_down_sync(0xffffffff, matrix_component, 8);
+        }
+        if (warp_count >= 8) {
+          matrix_component += __shfl_down_sync(0xffffffff, matrix_component, 4);
+        }
+        if (warp_count >= 4) {
+          matrix_component += __shfl_down_sync(0xffffffff, matrix_component, 2);
+        }
+        if (warp_count >= 2) {
+          matrix_component += __shfl_down_sync(0xffffffff, matrix_component, 1);
+        }
+
+        if (lane_index == 0) {
+          float constant = i == 0 ? 0.0f : s_gmm[i] * s_gmm[j];
+
+          if (i != 0 && i == j) {
+            constant -= EPSILON;
+          }
+
+          s_gmm[index] = norm_factor * matrix_component - constant;
+
+          if (index == 0 && matrix_component > 0) {
+            norm_factor = 1.0f / matrix_component;
+          }
+        }
+      }
+
+      __syncthreads();
+    }
+  }
+
+  float* matrix = s_gmm + (CHANNEL_COUNT + 1);
+  float* det_ptr = s_gmm + MATRIX_COMPONENT_COUNT;
+
+  if (local_index == 0) {
+    float square_mat[CHANNEL_COUNT][CHANNEL_COUNT];
+    float cholesky_mat[CHANNEL_COUNT][CHANNEL_COUNT];
+
+    for (int i = 0; i < CHANNEL_COUNT; i++) {
+      for (int j = 0; j < CHANNEL_COUNT; j++) {
+        square_mat[i][j] = 0.0f;
+        cholesky_mat[i][j] = 0.0f;
+      }
+    }
+
+    to_square(matrix, square_mat);
+    cholesky(square_mat, cholesky_mat);
+
+    *det_ptr = chol_det(cholesky_mat);
+
+    if (invert_matrix) {
+      chol_inv(cholesky_mat, square_mat);
+      to_triangle(square_mat, matrix);
+    }
+  }
+
+  if (local_index < GMM_COMPONENT_COUNT) {
+    g_batch_gmm[gmm_index * GMM_COMPONENT_COUNT + local_index] = s_gmm[local_index];
+  }
+}
+
+struct GMMSplit_t {
+  int idx;
+  float threshold;
+  float eigenvector[CHANNEL_COUNT];
+};
+
+// 1 Block, 32xMIXTURE_COUNT
+__global__ void GMMFindSplit(GMMSplit_t* gmmSplit, int gmmK, float* gmm) {
+  int batch_index = blockIdx.z;
+
+  float* g_batch_gmm = gmm + batch_index * GMM_COUNT * GMM_COMPONENT_COUNT;
+  GMMSplit_t* g_batch_gmmSplit = gmmSplit + batch_index * MIXTURE_COUNT;
+
+  int gmm_idx = threadIdx.x * MIXTURE_COUNT + threadIdx.y;
+
+  float eigenvalue = 0;
+  float eigenvector[CHANNEL_COUNT];
+
+  if (threadIdx.x < gmmK) {
+    float* matrix = g_batch_gmm + gmm_idx * GMM_COMPONENT_COUNT + (CHANNEL_COUNT + 1);
+    largest_eigenpair(matrix, eigenvector, &eigenvalue);
+  }
+
+  float max_value = eigenvalue;
+
+  max_value = max(max_value, __shfl_xor_sync(0xffffffff, max_value, 16));
+  max_value = max(max_value, __shfl_xor_sync(0xffffffff, max_value, 8));
+  max_value = max(max_value, __shfl_xor_sync(0xffffffff, max_value, 4));
+  max_value = max(max_value, __shfl_xor_sync(0xffffffff, max_value, 2));
+  max_value = max(max_value, __shfl_xor_sync(0xffffffff, max_value, 1));
+
+  if (max_value == eigenvalue) {
+    GMMSplit_t split;
+
+    float* average_feature = gmm + gmm_idx * GMM_COMPONENT_COUNT + 1;
+
+    split.idx = threadIdx.x;
+    split.threshold = scalar_prod(average_feature, eigenvector);
+
+    for (int i = 0; i < CHANNEL_COUNT; i++) {
+      split.eigenvector[i] = eigenvector[i];
+    }
+
+    g_batch_gmmSplit[threadIdx.y] = split;
+  }
+}
+
+#define DO_SPLIT_DEGENERACY 4
+
+__global__ void GMMDoSplit(const GMMSplit_t* gmmSplit, int k, const float* image, int* alpha, int element_count) {
+  __shared__ GMMSplit_t s_gmmSplit[MIXTURE_COUNT];
+
+  int batch_index = blockIdx.z;
+
+  const GMMSplit_t* g_batch_gmmSplit = gmmSplit + batch_index * MIXTURE_COUNT;
+  const float* g_batch_image = image + batch_index * element_count * CHANNEL_COUNT;
+  int* g_batch_alpha = alpha + batch_index * element_count;
+
+  int* s_linear = (int*)s_gmmSplit;
+  int* g_linear = (int*)g_batch_gmmSplit;
+
+  if (threadIdx.x < MIXTURE_COUNT * sizeof(GMMSplit_t)) {
+    s_linear[threadIdx.x] = g_linear[threadIdx.x];
+  }
+
+  __syncthreads();
+
+  int index = threadIdx.x + blockIdx.x * BLOCK_SIZE * DO_SPLIT_DEGENERACY;
+
+  for (int i = 0; i < DO_SPLIT_DEGENERACY; i++) {
+    index += BLOCK_SIZE;
+
+    if (index < element_count) {
+      int my_alpha = g_batch_alpha[index];
+
+      if (my_alpha != -1) {
+        int select = my_alpha & 15;
+        int gmm_idx = my_alpha >> 4;
+
+        if (gmm_idx == s_gmmSplit[select].idx) {
+          // in the split cluster now
+          float feature[CHANNEL_COUNT];
+
+          for (int i = 0; i < CHANNEL_COUNT; i++) {
+            feature[i] = g_batch_image[index + i * element_count];
+          }
+
+          float value = scalar_prod(s_gmmSplit[select].eigenvector, feature);
+
+          if (value > s_gmmSplit[select].threshold) {
+            // assign pixel to new cluster
+            g_batch_alpha[index] = k + select;
+          }
+        }
+      }
+    }
+  }
+}
+
+// Single block, 32xMIXTURE_COUNT
+__global__ void GMMcommonTerm(float* g_gmm) {
+  int batch_index = blockIdx.z;
+
+  float* g_batch_gmm = g_gmm + batch_index * GMM_COUNT * GMM_COMPONENT_COUNT;
+
+  int gmm_index = (threadIdx.x * MIXTURE_COUNT) + threadIdx.y;
+
+  float gmm_n = threadIdx.x < MIXTURE_SIZE ? g_batch_gmm[gmm_index * GMM_COMPONENT_COUNT] : 0.0f;
+
+  float sum = gmm_n;
+
+  sum += __shfl_xor_sync(0xffffffff, sum, 1);
+  sum += __shfl_xor_sync(0xffffffff, sum, 2);
+  sum += __shfl_xor_sync(0xffffffff, sum, 4);
+  sum += __shfl_xor_sync(0xffffffff, sum, 8);
+  sum += __shfl_xor_sync(0xffffffff, sum, 16);
+
+  if (threadIdx.x < MIXTURE_SIZE) {
+    float det = g_batch_gmm[gmm_index * GMM_COMPONENT_COUNT + MATRIX_COMPONENT_COUNT] + EPSILON;
+    float commonTerm = det > 0.0f ? gmm_n / (sqrtf(det) * sum) : gmm_n / sum;
+
+    g_batch_gmm[gmm_index * GMM_COMPONENT_COUNT + MATRIX_COMPONENT_COUNT] = commonTerm;
+  }
+}
+
+__device__ float GMMTerm(float* feature, const float* gmm) {
+  const float* average_feature = gmm + 1;
+  const float* matrix = gmm + CHANNEL_COUNT + 1;
+
+  float diff[CHANNEL_COUNT];
+
+  for (int i = 0; i < CHANNEL_COUNT; i++) {
+    diff[i] = feature[i] - average_feature[i];
+  }
+
+  float value = 0.0f;
+
+  for (int index = 0, i = 0; i < CHANNEL_COUNT; i++) {
+    for (int j = i; j < CHANNEL_COUNT; j++, index++) {
+      float term = diff[i] * diff[j] * matrix[index];
+
+      value += i == j ? term : 2 * term;
+    }
+  }
+
+  return gmm[MATRIX_COMPONENT_COUNT] * expf(-0.5 * value);
+}
+
+__global__ void GMMDataTermKernel(const float* image, const float* gmm, float* output, int element_count) {
+  int batch_index = blockIdx.z;
+
+  const float* g_batch_image = image + batch_index * element_count * CHANNEL_COUNT;
+  const float* g_batch_gmm = gmm + batch_index * GMM_COUNT * GMM_COMPONENT_COUNT;
+  float* g_batch_output = output + batch_index * element_count * MIXTURE_COUNT;
+
+  int index = blockIdx.x * blockDim.x + threadIdx.x;
+
+  if (index >= element_count)
+    return;
+
+  float feature[CHANNEL_COUNT];
+
+  for (int i = 0; i < CHANNEL_COUNT; i++) {
+    feature[i] = g_batch_image[index + i * element_count];
+  }
+
+  float weights[MIXTURE_COUNT];
+  float weight_total = 0.0f;
+
+  for (int i = 0; i < MIXTURE_COUNT; i++) {
+    float mixture_weight = 0.0f;
+
+    for (int j = 0; j < MIXTURE_SIZE; j++) {
+      mixture_weight += GMMTerm(feature, &g_batch_gmm[(MIXTURE_COUNT * j + i) * GMM_COMPONENT_COUNT]);
+    }
+
+    weights[i] = mixture_weight;
+    weight_total += mixture_weight;
+  }
+
+  for (int i = 0; i < MIXTURE_COUNT; i++) {
+    // protecting against pixels with 0 in all mixtures
+    float final_weight = weight_total > 0.0f ? weights[i] / weight_total : 0.0f;
+    g_batch_output[index + i * element_count] = final_weight;
+  }
+}
+
+#define THREADS 512
+#define WARPS 16
+#define BLOCK (WARPS << 5)
+#define LOAD 4
+
+void GMMInitialize(
+    const float* image,
+    int* alpha,
+    float* gmm,
+    float* scratch_mem,
+    unsigned int batch_count,
+    unsigned int element_count) {
+  unsigned int block_count = TILE(element_count, BLOCK * LOAD);
+
+  float* block_gmm_scratch = scratch_mem;
+  GMMSplit_t* gmm_split_scratch = (GMMSplit_t*)scratch_mem;
+
+  int gmm_N = MIXTURE_COUNT * MIXTURE_SIZE;
+
+  for (unsigned int k = MIXTURE_COUNT; k < gmm_N; k += MIXTURE_COUNT) {
+    for (unsigned int i = 0; i < k; ++i) {
+      CovarianceReductionKernel<WARPS, LOAD>
+          <<<{block_count, 1, batch_count}, BLOCK>>>(i, image, alpha, block_gmm_scratch, element_count);
+    }
+
+    CovarianceFinalizationKernel<WARPS, false><<<{k, 1, batch_count}, BLOCK>>>(block_gmm_scratch, gmm, block_count);
+
+    GMMFindSplit<<<{1, 1, batch_count}, dim3(BLOCK_SIZE, MIXTURE_COUNT)>>>(gmm_split_scratch, k / MIXTURE_COUNT, gmm);
+    GMMDoSplit<<<{TILE(element_count, BLOCK_SIZE * DO_SPLIT_DEGENERACY), 1, batch_count}, BLOCK_SIZE>>>(
+        gmm_split_scratch, (k / MIXTURE_COUNT) << 4, image, alpha, element_count);
+  }
+}
+
+void GMMUpdate(
+    const float* image,
+    int* alpha,
+    float* gmm,
+    float* scratch_mem,
+    unsigned int batch_count,
+    unsigned int element_count) {
+  unsigned int block_count = TILE(element_count, BLOCK * LOAD);
+
+  float* block_gmm_scratch = scratch_mem;
+
+  unsigned int gmm_N = MIXTURE_COUNT * MIXTURE_SIZE;
+
+  for (unsigned int i = 0; i < gmm_N; ++i) {
+    CovarianceReductionKernel<WARPS, LOAD>
+        <<<{block_count, 1, batch_count}, BLOCK>>>(i, image, alpha, block_gmm_scratch, element_count);
+  }
+
+  CovarianceFinalizationKernel<WARPS, true><<<{gmm_N, 1, batch_count}, BLOCK>>>(block_gmm_scratch, gmm, block_count);
+
+  GMMcommonTerm<<<{1, 1, batch_count}, dim3(BLOCK_SIZE, MIXTURE_COUNT)>>>(gmm);
+}
+
+void GMMDataTerm(
+    const float* image,
+    const float* gmm,
+    float* output,
+    unsigned int batch_count,
+    unsigned int element_count) {
+  dim3 block(BLOCK_SIZE, 1);
+  dim3 grid(TILE(element_count, BLOCK_SIZE), 1, batch_count);
+
+  GMMDataTermKernel<<<grid, block>>>(image, gmm, output, element_count);
+}
+
+void learn_cuda(
+    const float* input,
+    const int* labels,
+    float* gmm,
+    float* scratch_memory,
+    unsigned int batch_count,
+    unsigned int element_count) {
+  int* alpha = (int*)scratch_memory;
+  float* scratch_mem = scratch_memory + batch_count * element_count;
+
+  cudaMemcpyAsync(alpha, labels, batch_count * element_count * sizeof(int), cudaMemcpyDeviceToDevice);
+
+  GMMInitialize(input, alpha, gmm, scratch_mem, batch_count, element_count);
+  GMMUpdate(input, alpha, gmm, scratch_mem, batch_count, element_count);
+}
+
+void apply_cuda(
+    const float* gmm,
+    const float* input,
+    float* output,
+    unsigned int batch_count,
+    unsigned int element_count) {
+  GMMDataTerm(input, gmm, output, batch_count, element_count);
+}

my_container_sandbox/workspace/anaconda3/lib/python3.8/site-packages/monai/_extensions/gmm/gmm_cuda_linalg.cuh

@@ -0,0 +1,144 @@
+/*
+Copyright (c) MONAI Consortium
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+    http://www.apache.org/licenses/LICENSE-2.0
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+__device__ void to_square(float in[SUB_MATRIX_COMPONENT_COUNT], float out[CHANNEL_COUNT][CHANNEL_COUNT]) {
+  for (int index = 0, i = 0; i < CHANNEL_COUNT; i++) {
+    for (int j = i; j < CHANNEL_COUNT; j++, index++) {
+      out[i][j] = in[index];
+      out[j][i] = in[index];
+    }
+  }
+}
+
+__device__ void to_triangle(float in[CHANNEL_COUNT][CHANNEL_COUNT], float out[SUB_MATRIX_COMPONENT_COUNT]) {
+  for (int index = 0, i = 0; i < CHANNEL_COUNT; i++) {
+    for (int j = i; j < CHANNEL_COUNT; j++, index++) {
+      out[index] = in[j][i];
+    }
+  }
+}
+
+__device__ void cholesky(float in[CHANNEL_COUNT][CHANNEL_COUNT], float out[CHANNEL_COUNT][CHANNEL_COUNT]) {
+  for (int i = 0; i < CHANNEL_COUNT; i++) {
+    for (int j = 0; j < i + 1; j++) {
+      float sum = 0.0f;
+
+      for (int k = 0; k < j; k++) {
+        sum += out[i][k] * out[j][k];
+      }
+
+      if (i == j) {
+        out[i][j] = sqrtf(in[i][i] - sum);
+      } else {
+        out[i][j] = (in[i][j] - sum) / out[j][j];
+      }
+    }
+  }
+}
+
+__device__ float chol_det(float in[CHANNEL_COUNT][CHANNEL_COUNT]) {
+  float det = 1.0f;
+
+  for (int i = 0; i < CHANNEL_COUNT; i++) {
+    det *= in[i][i];
+  }
+
+  return det * det;
+}
+
+__device__ void chol_inv(float in[CHANNEL_COUNT][CHANNEL_COUNT], float out[CHANNEL_COUNT][CHANNEL_COUNT]) {
+  // Invert cholesky matrix
+  for (int i = 0; i < CHANNEL_COUNT; i++) {
+    in[i][i] = 1.0f / (in[i][i] + 0.0001f);
+
+    for (int j = 0; j < i; j++) {
+      float sum = 0.0f;
+
+      for (int k = j; k < i; k++) {
+        sum += in[i][k] * in[k][j];
+      }
+
+      in[i][j] = -in[i][i] * sum;
+    }
+  }
+
+  // Dot with transpose of self
+  for (int i = 0; i < CHANNEL_COUNT; i++) {
+    for (int j = 0; j < CHANNEL_COUNT; j++) {
+      out[i][j] = 0.0f;
+
+      for (int k = max(i, j); k < CHANNEL_COUNT; k++) {
+        out[i][j] += in[k][i] * in[k][j];
+      }
+    }
+  }
+}
+
+__device__ void normalize(float* v) {
+  float norm = 0.0f;
+
+  for (int i = 0; i < CHANNEL_COUNT; i++) {
+    norm += v[i] * v[i];
+  }
+
+  norm = 1.0f / sqrtf(norm);
+
+  for (int i = 0; i < CHANNEL_COUNT; i++) {
+    v[i] *= norm;
+  }
+}
+
+__device__ float scalar_prod(float* a, float* b) {
+  float product = 0.0f;
+
+  for (int i = 0; i < CHANNEL_COUNT; i++) {
+    product += a[i] * b[i];
+  }
+
+  return product;
+}
+
+__device__ void largest_eigenpair(const float* M, float* evec, float* eval) {
+  float scratch[CHANNEL_COUNT];
+
+  for (int i = 0; i < CHANNEL_COUNT; i++) {
+    scratch[i] = i + 1;
+  }
+
+  for (int itr = 0; itr < 10; itr++) {
+    *eval = 0.0f;
+
+    for (int i = 0; i < CHANNEL_COUNT; i++) {
+      int index = i;
+
+      evec[i] = 0.0f;
+
+      for (int j = 0; j < CHANNEL_COUNT; j++) {
+        evec[i] += M[index] * scratch[j];
+
+        if (j < i) {
+          index += CHANNEL_COUNT - (j + 1);
+        } else {
+          index += 1;
+        }
+      }
+
+      *eval = max(*eval, evec[i]);
+    }
+
+    for (int i = 0; i < CHANNEL_COUNT; i++) {
+      evec[i] /= *eval;
+      scratch[i] = evec[i];
+    }
+  }
+}