shap--shap/shap/cext/_cext_gpu.cu

#include <Python.h>

#include "gpu_treeshap.h"
#include "tree_shap.h"

const float inf = std::numeric_limits<tfloat>::infinity();

struct CategoryConstraint {
  CategoryConstraint() = default;

  __host__ __device__ static CategoryConstraint All() {
    return CategoryConstraint{};
  }

  __host__ __device__ static CategoryConstraint Include(int categories) {
    CategoryConstraint constraint;
    constraint.has_include_categories = true;
    constraint.include_categories = categories;
    return constraint;
  }

  __host__ __device__ static CategoryConstraint Exclude(int categories) {
    CategoryConstraint constraint;
    constraint.exclude_categories = categories;
    return constraint;
  }

  __host__ __device__ static bool CategoryInMask(int categories,
                                                 float category) {
    int category_int = static_cast<int>(category);
    if (category_int < 1) {
      return false;
    }
    int category_flag = 1 << (category_int - 1);
    return (categories & category_flag) != 0;
  }

  __host__ __device__ bool Contains(float category) const {
    bool included = !has_include_categories ||
                    CategoryInMask(include_categories, category);
    return included && !CategoryInMask(exclude_categories, category);
  }

  __host__ __device__ void Intersect(const CategoryConstraint &other) {
    if (other.has_include_categories) {
      include_categories = has_include_categories
                               ? include_categories & other.include_categories
                               : other.include_categories;
      has_include_categories = true;
    }
    exclude_categories |= other.exclude_categories;
    if (has_include_categories) {
      include_categories &= ~exclude_categories;
    }
  }

  bool has_include_categories = false;
  int include_categories = 0;
  int exclude_categories = 0;
};

struct ShapSplitCondition {
  ShapSplitCondition() = default;
  ShapSplitCondition(tfloat feature_lower_bound, tfloat feature_upper_bound,
                     bool is_missing_branch)
      : feature_lower_bound(feature_lower_bound),
        feature_upper_bound(feature_upper_bound),
        is_missing_branch(is_missing_branch),
        categories(CategoryConstraint::All()) {
    assert(feature_lower_bound <= feature_upper_bound);
  }
  ShapSplitCondition(tfloat feature_lower_bound, tfloat feature_upper_bound,
                     bool is_missing_branch,
                     CategoryConstraint category_constraint)
      : feature_lower_bound(feature_lower_bound),
        feature_upper_bound(feature_upper_bound),
        is_missing_branch(is_missing_branch),
        categories(category_constraint) {
    assert(feature_lower_bound <= feature_upper_bound);
  }

  /*! Feature values >= lower and < upper flow down this path. */
  tfloat feature_lower_bound;
  tfloat feature_upper_bound;
  /*! Do missing values flow down this path? */
  bool is_missing_branch;
  CategoryConstraint categories;

  // Does this instance flow down this path?
  __host__ __device__ bool EvaluateSplit(float x) const {
    // is nan
    if (isnan(x)) {
      return is_missing_branch;
    }
    if (!categories.Contains(x)) {
      return false;
    }
    return x > feature_lower_bound && x <= feature_upper_bound;
  }

  // Combine two split conditions on the same feature
  __host__ __device__ void
  Merge(const ShapSplitCondition &other) {  // Combine duplicate features
    feature_lower_bound = max(feature_lower_bound, other.feature_lower_bound);
    feature_upper_bound = min(feature_upper_bound, other.feature_upper_bound);
    categories.Intersect(other.categories);
    is_missing_branch = is_missing_branch && other.is_missing_branch;
  }
};


// Inspired by: https://en.cppreference.com/w/cpp/iterator/size
// Limited implementation of std::size fo arrays
template <class T, size_t N>
constexpr size_t array_size(const T (&array)[N]) noexcept
{
    return N;
}

void RecurseTree(
    unsigned pos, const TreeEnsemble &tree,
    std::vector<gpu_treeshap::PathElement<ShapSplitCondition>> *tmp_path,
    std::vector<gpu_treeshap::PathElement<ShapSplitCondition>> *paths,
    size_t *path_idx, int num_outputs) {
  if (tree.is_leaf(pos)) {
    for (auto j = 0ull; j < num_outputs; j++) {
      auto v = tree.values[pos * num_outputs + j];
      if (v == 0.0) {
        // The tree has no output for this class, don't bother adding the path
        continue;
      }
      // Go back over path, setting v, path_idx
      for (auto &e : *tmp_path) {
        e.v = v;
        e.group = j;
        e.path_idx = *path_idx;
      }

      paths->insert(paths->end(), tmp_path->begin(), tmp_path->end());
      // Increment path index
      (*path_idx)++;
    }
    return;
  }

  unsigned left_child = tree.children_left[pos];
  bool is_left_default = tree.children_default[pos] == left_child;
  double left_zero_fraction =
      tree.node_sample_weights[left_child] / tree.node_sample_weights[pos];
  auto threshold_type = tree.threshold_types[pos];
  auto threshold = tree.thresholds[pos];

  ShapSplitCondition left_condition{-inf, threshold, is_left_default};
  ShapSplitCondition right_condition{threshold, inf, !is_left_default};
  if (threshold_type == 1) {
    int categories = static_cast<int>(threshold);
    left_condition =
        ShapSplitCondition{-inf, inf, is_left_default,
                           CategoryConstraint::Include(categories)};
    right_condition =
        ShapSplitCondition{-inf, inf, !is_left_default,
                           CategoryConstraint::Exclude(categories)};
  }

  // Add left split to the path
  tmp_path->emplace_back(0, tree.features[pos], 0, left_condition,
                         left_zero_fraction, 0.0f);

  RecurseTree(left_child, tree, tmp_path, paths, path_idx, num_outputs);

  // Add right split to the path
  tmp_path->back() = gpu_treeshap::PathElement<ShapSplitCondition>(
      0, tree.features[pos], 0, right_condition,
      1.0 - left_zero_fraction, 0.0f);

  RecurseTree(tree.children_right[pos], tree, tmp_path, paths, path_idx,
              num_outputs);

  tmp_path->pop_back();
}

std::vector<gpu_treeshap::PathElement<ShapSplitCondition>>
ExtractPaths(const TreeEnsemble &trees) {
  std::vector<gpu_treeshap::PathElement<ShapSplitCondition>> paths;
  size_t path_idx = 0;
  for (auto i = 0; i < trees.tree_limit; i++) {
    TreeEnsemble tree;
    trees.get_tree(tree, i);
    std::vector<gpu_treeshap::PathElement<ShapSplitCondition>> tmp_path;
    tmp_path.reserve(tree.max_depth);
    tmp_path.emplace_back(0, -1, 0, ShapSplitCondition{-inf, inf, false}, 1.0,
                          0.0f);
    RecurseTree(0, tree, &tmp_path, &paths, &path_idx, tree.num_outputs);
  }
  return paths;
}

class DeviceExplanationDataset {
  thrust::device_vector<tfloat> data;
  thrust::device_vector<bool> missing;
  size_t num_features;
  size_t num_rows;

 public:
  DeviceExplanationDataset(const ExplanationDataset &host_data,
                           bool background_dataset = false) {
    num_features = host_data.M;
    if (background_dataset) {
      num_rows = host_data.num_R;
      data = thrust::device_vector<tfloat>(
          host_data.R, host_data.R + host_data.num_R * host_data.M);
      missing = thrust::device_vector<bool>(host_data.R_missing,
                                            host_data.R_missing +
                                                host_data.num_R * host_data.M);

    } else {
      num_rows = host_data.num_X;
      data = thrust::device_vector<tfloat>(
          host_data.X, host_data.X + host_data.num_X * host_data.M);
      missing = thrust::device_vector<bool>(host_data.X_missing,
                                            host_data.X_missing +
                                                host_data.num_X * host_data.M);
    }
  }

  class DenseDatasetWrapper {
    const tfloat *data;
    const bool *missing;
    int num_rows;
    int num_cols;

   public:
    DenseDatasetWrapper() = default;
    DenseDatasetWrapper(const tfloat *data, const bool *missing, int num_rows,
                        int num_cols)
        : data(data), missing(missing), num_rows(num_rows), num_cols(num_cols) {
    }
    __device__ tfloat GetElement(size_t row_idx, size_t col_idx) const {
      auto idx = row_idx * num_cols + col_idx;
      if (missing[idx]) {
        return std::numeric_limits<tfloat>::quiet_NaN();
      }
      return data[idx];
    }
    __host__ __device__ size_t NumRows() const { return num_rows; }
    __host__ __device__ size_t NumCols() const { return num_cols; }
  };

  DenseDatasetWrapper GetDeviceAccessor() {
    return DenseDatasetWrapper(data.data().get(), missing.data().get(),
                               num_rows, num_features);
  }
};

inline void dense_tree_path_dependent_gpu(
    const TreeEnsemble &trees, const ExplanationDataset &data,
    tfloat *out_contribs, tfloat transform(const tfloat, const tfloat)) {
  auto paths = ExtractPaths(trees);
  DeviceExplanationDataset device_data(data);
  DeviceExplanationDataset::DenseDatasetWrapper X =
      device_data.GetDeviceAccessor();

  thrust::device_vector<float> phis((X.NumCols() + 1) * X.NumRows() *
                                    trees.num_outputs);
  gpu_treeshap::GPUTreeShap(X, paths.begin(), paths.end(), trees.num_outputs,
                            phis.begin(), phis.end());
  // Add the base offset term to bias
  thrust::device_vector<double> base_offset(
      trees.base_offset, trees.base_offset + trees.num_outputs);
  auto counting = thrust::make_counting_iterator(size_t(0));
  auto d_phis = phis.data().get();
  auto d_base_offset = base_offset.data().get();
  size_t num_groups = trees.num_outputs;
  thrust::for_each(counting, counting + X.NumRows() * trees.num_outputs,
                   [=] __device__(size_t idx) {
                     size_t row_idx = idx / num_groups;
                     size_t group = idx % num_groups;
                     auto phi_idx = gpu_treeshap::IndexPhi(
                         row_idx, num_groups, group, X.NumCols(), X.NumCols());
                     d_phis[phi_idx] += d_base_offset[group];
                   });

  // Shap uses a slightly different layout for multiclass
  thrust::device_vector<float> transposed_phis(phis.size());
  auto d_transposed_phis = transposed_phis.data();
  thrust::for_each(
      counting, counting + phis.size(), [=] __device__(size_t idx) {
        size_t old_shape[] = {X.NumRows(), num_groups, (X.NumCols() + 1)};
        size_t old_idx[array_size(old_shape)];
        gpu_treeshap::FlatIdxToTensorIdx(idx, old_shape, old_idx);
        // Define new tensor format, switch num_groups axis to end
        size_t new_shape[] = {X.NumRows(), (X.NumCols() + 1), num_groups};
        size_t new_idx[] = {old_idx[0], old_idx[2], old_idx[1]};
        size_t transposed_idx =
            gpu_treeshap::TensorIdxToFlatIdx(new_shape, new_idx);
        d_transposed_phis[transposed_idx] = d_phis[idx];
      });
  thrust::copy(transposed_phis.begin(), transposed_phis.end(), out_contribs);
}

inline void
dense_tree_independent_gpu(const TreeEnsemble &trees,
                           const ExplanationDataset &data, tfloat *out_contribs,
                           tfloat transform(const tfloat, const tfloat)) {
  auto paths = ExtractPaths(trees);
  DeviceExplanationDataset device_data(data);
  DeviceExplanationDataset::DenseDatasetWrapper X =
      device_data.GetDeviceAccessor();
  DeviceExplanationDataset background_device_data(data, true);
  DeviceExplanationDataset::DenseDatasetWrapper R =
      background_device_data.GetDeviceAccessor();

  thrust::device_vector<float> phis((X.NumCols() + 1) * X.NumRows() *
                                    trees.num_outputs);
  gpu_treeshap::GPUTreeShapInterventional(X, R, paths.begin(), paths.end(),
                                          trees.num_outputs, phis.begin(),
                                          phis.end());
  // Add the base offset term to bias
  thrust::device_vector<double> base_offset(
      trees.base_offset, trees.base_offset + trees.num_outputs);
  auto counting = thrust::make_counting_iterator(size_t(0));
  auto d_phis = phis.data().get();
  auto d_base_offset = base_offset.data().get();
  size_t num_groups = trees.num_outputs;
  thrust::for_each(counting, counting + X.NumRows() * trees.num_outputs,
                   [=] __device__(size_t idx) {
                     size_t row_idx = idx / num_groups;
                     size_t group = idx % num_groups;
                     auto phi_idx = gpu_treeshap::IndexPhi(
                         row_idx, num_groups, group, X.NumCols(), X.NumCols());
                     d_phis[phi_idx] += d_base_offset[group];
                   });

  // Shap uses a slightly different layout for multiclass
  thrust::device_vector<float> transposed_phis(phis.size());
  auto d_transposed_phis = transposed_phis.data();
  thrust::for_each(
      counting, counting + phis.size(), [=] __device__(size_t idx) {
        size_t old_shape[] = {X.NumRows(), num_groups, (X.NumCols() + 1)};
        size_t old_idx[array_size(old_shape)];
        gpu_treeshap::FlatIdxToTensorIdx(idx, old_shape, old_idx);
        // Define new tensor format, switch num_groups axis to end
        size_t new_shape[] = {X.NumRows(), (X.NumCols() + 1), num_groups};
        size_t new_idx[] = {old_idx[0], old_idx[2], old_idx[1]};
        size_t transposed_idx =
            gpu_treeshap::TensorIdxToFlatIdx(new_shape, new_idx);
        d_transposed_phis[transposed_idx] = d_phis[idx];
      });
  thrust::copy(transposed_phis.begin(), transposed_phis.end(), out_contribs);
}

inline void dense_tree_path_dependent_interactions_gpu(
    const TreeEnsemble &trees, const ExplanationDataset &data,
    tfloat *out_contribs, tfloat transform(const tfloat, const tfloat)) {
  auto paths = ExtractPaths(trees);
  DeviceExplanationDataset device_data(data);
  DeviceExplanationDataset::DenseDatasetWrapper X =
      device_data.GetDeviceAccessor();

  thrust::device_vector<float> phis((X.NumCols() + 1) * (X.NumCols() + 1) *
                                    X.NumRows() * trees.num_outputs);
  gpu_treeshap::GPUTreeShapInteractions(X, paths.begin(), paths.end(),
                                        trees.num_outputs, phis.begin(),
                                        phis.end());
  // Add the base offset term to bias
  thrust::device_vector<double> base_offset(
      trees.base_offset, trees.base_offset + trees.num_outputs);
  auto counting = thrust::make_counting_iterator(size_t(0));
  auto d_phis = phis.data().get();
  auto d_base_offset = base_offset.data().get();
  size_t num_groups = trees.num_outputs;
  thrust::for_each(counting, counting + X.NumRows() * num_groups,
                   [=] __device__(size_t idx) {
                     size_t row_idx = idx / num_groups;
                     size_t group = idx % num_groups;
                     auto phi_idx = gpu_treeshap::IndexPhiInteractions(
                         row_idx, num_groups, group, X.NumCols(), X.NumCols(),
                         X.NumCols());
                     d_phis[phi_idx] += d_base_offset[group];
                   });
  // Shap uses a slightly different layout for multiclass
  thrust::device_vector<float> transposed_phis(phis.size());
  auto d_transposed_phis = transposed_phis.data();
  thrust::for_each(
      counting, counting + phis.size(), [=] __device__(size_t idx) {
        size_t old_shape[] = {X.NumRows(), num_groups, (X.NumCols() + 1),
                              (X.NumCols() + 1)};
        size_t old_idx[array_size(old_shape)];
        gpu_treeshap::FlatIdxToTensorIdx(idx, old_shape, old_idx);
        // Define new tensor format, switch num_groups axis to end
        size_t new_shape[] = {X.NumRows(), (X.NumCols() + 1), (X.NumCols() + 1),
                              num_groups};
        size_t new_idx[] = {old_idx[0], old_idx[2], old_idx[3], old_idx[1]};
        size_t transposed_idx =
            gpu_treeshap::TensorIdxToFlatIdx(new_shape, new_idx);
        d_transposed_phis[transposed_idx] = d_phis[idx];
      });
  thrust::copy(transposed_phis.begin(), transposed_phis.end(), out_contribs);
}

void dense_tree_shap_gpu(const TreeEnsemble &trees,
                         const ExplanationDataset &data, tfloat *out_contribs,
                         const int feature_dependence, unsigned model_transform,
                         bool interactions) {
  // see what transform (if any) we have
  transform_f transform = get_transform(model_transform);

  // dispatch to the correct algorithm handler
  switch (feature_dependence) {
  case FEATURE_DEPENDENCE::independent:
    if (interactions) {
      std::cerr << "FEATURE_DEPENDENCE::independent with interactions not yet "
                   "supported\n";
    } else {
      dense_tree_independent_gpu(trees, data, out_contribs, transform);
    }
    return;

  case FEATURE_DEPENDENCE::tree_path_dependent:
    if (interactions) {
      dense_tree_path_dependent_interactions_gpu(trees, data, out_contribs,
                                                 transform);
    } else {
      dense_tree_path_dependent_gpu(trees, data, out_contribs, transform);
    }
    return;

  case FEATURE_DEPENDENCE::global_path_dependent:
    std::cerr << "FEATURE_DEPENDENCE::global_path_dependent not supported\n";
    return;
  default:
    std::cerr << "Unknown feature dependence option\n";
    return;
  }
}