code: def compute_swap_map(matrix, used_stripes):
      do_gpu = use_gpu()
      assert do_gpu
  
      stripe_pairs = build_stripe_pairs(matrix, used_stripes).astype(
          np.uint32
      )
      matrix_view = matrix.astype(np.float32).flatten()
      stripe_pairs_view = stripe_pairs.flatten()
      output = np.zeros(
          (len(stripe_pairs) * 16), dtype=np.float32
      ).flatten()
      result = permutation_search_cuda_kernels.build_swap_map(
          matrix_view,
          matrix.shape[0],
          matrix.shape[1],
          stripe_pairs_view,
          output,
      )
  
      # translate the flat array from the GPU to a map
      pair_improvement_map = {}
      for i, pair in enumerate(stripe_pairs):
          for swap_idx in range(16):
              col0, col1 = stripes_and_swap_idx_to_columns(
                  pair[0], pair[1], swap_idx
              )
              pair_improvement_map[(col0, col1)] = output[
                  i * 16 + swap_idx
              ]
      return pair_improvement_map
  
file_path: apex/apex/contrib/sparsity/permutation_search_kernels/channel_swap.py
imports: from .permutation_utilities import *
label_desc: Use image node_id 5 for calling a global function with example usage: compute_swap_map(matrix, used_stripes) and returns: pair_improvement_map