“Python 用户友爱却运转功率低”，“C++ 运转功率较高，但完结一个功用代码量会远大于 Python”。平常学习工作中你是否常听到相似的说法？在 Python 大行其道的今天，你是否常常也会面临代码的瓶颈，而为运转加快而烦恼呢？“我的代码刚跑 10 步，隔壁同学的现已跑完第一个 epoch 了。”–这究竟是人性的扭曲还是科学的沦丧？荀子有言“正人性非异也，善假于物也”。本期《源码解读》带你走进 “Pytorch 中 (神秘) 的 C++ / CUDA 扩展“。
- 本期主题：结合 Python 与 C++ 各自的长处，在 PyTorch 中加入 C++ / CUDA的扩展，而让咱们自己更好地运用工具而不为工具所捆绑。
- 代码来源：MMCV,PyTorch。
- 注：C++ / CUDA 扩展一般有”预编译“ 与 ”实时编译“ （just-in-time, JIT）方法。本期首要介绍”预编译“方法。
1. 由扩展的调用方法说起
当你想为自己的代码添加扩展进行加快时，咱们能够先来看看经典的比如中是怎么处理的。对检测或切割稍有了解的同学应该知道，nms 的核算是最常见的用到了 C++ / CUDA 扩展的算子。
from mmcv import _ext as ext_module
from torch.autograd import Function
def nms(boxes, scores, iou_threshold, offset=0):
inds = NMSop.apply(boxes, scores, iou_threshold, offset)
dets = torch.cat((boxes[inds], scores[inds].reshape(-1, 1)), dim=1)
return dets, inds
class NMSop(torch.autograd.Function):
@staticmethod
def forward(ctx, bboxes, scores, iou_threshold, offset):
inds = ext_module.nms(
bboxes, scores, iou_threshold=float(iou_threshold), offset=offset)
return inds
@staticmethod
def symbolic(g, bboxes, scores, iou_threshold, offset):
pass # onnx 转化相关
Function (见往期内容torch.autograd)。NMSop的forward函数内核调用的是mmcv._ext.nms模块，但实践上咱们在 MMCV源码中是看不到_extmodule 的。只要在编译好的mmcv 库 (MMCV_WITH_OPS=True python setup.py build_ext --inplace) 会呈现mmcv/_ext.cpython-xxx.so文件，只要这时在 Python 中运转import mmcv._ext才会成功。看来 C++ 扩展是经过 setup.py 来履行编译的。
2. setup.py — 扩展的编译文件
基于预编译的扩展由于需求编译，而setup.py文件正是基于setuptools的编译脚本。因而一个 Python package 的扩展是能够在setup.py文件中找到其蛛丝马迹的。这里咱们截取一段mmcv的 setup.py文件，
setup(
name='mmcv',
install_requires=install_requires,
# 需求编译的c++/cuda扩展
ext_modules=get_extensions(),
# cmdclass 为python setup.py --build_ext指令指定行为
cmdclass={'build_ext': torch.utils.cpp_extension.BuildExtension})
这里能够看到setup函数中一个首要的参数ext_modules，该参数需求指定为一个Extension列表，代表实践需求编译的扩展。目前该参数由get_extensions函数取得。其间get_extensions函数界说如下（节选）
def get_extensions():
extensions = []
ext_name = 'mmcv._ext'
from torch.utils.cpp_extension import (CUDAExtension, CppExtension)
if torch.cuda.is_available():
# CUDA编译扩展
extra_compile_args['nvcc'] = [cuda_args] if cuda_args else []
# 编译./mmcv/ops/csrc/pytorch文件夹中的一切文件
op_files = glob.glob('./mmcv/ops/csrc/pytorch/*')
extension = CUDAExtension
else:
# C++ 编译扩展
op_files = glob.glob('./mmcv/ops/csrc/pytorch/*.cpp')
extension = CppExtension
include_path = os.path.abspath('./mmcv/ops/csrc')
ext_ops = extension(
name=ext_name, # 扩展模块名
sources=op_files, # 扩展文件
include_dirs=[include_path], # 头文件地址
define_macros=define_macros, # 预界说宏
extra_compile_args=extra_compile_args) # 其他编译选项
extensions.append(ext_ops)
return extensions
在上述代码中咱们总算看到了mmcv._ext，该姓名正是新界说的扩展的姓名。由此咱们便知道上文中提到的mmcv._ext模块实践上是在 setup.py 文件中指定其模块姓名的。 别的咱们发现用于生成扩展的函数会随体系环境不同而有所区别，当体系中没有 CUDA 时会调用CppExtension，且只编译一切 .cpp文件，反之则调用CUDAExtension。其实CppExtension与CUDAExtension都是基于setuptools.Extension的扩展，这两个函数都额定将体系目录中的torch/include加入到 C++ 编译时的include_dirs中，别的CUDAExtension会额定将CUDA相关的库以及头文件加到默许编译查找途径中。 由 setup.py 文件咱们还了解到送给编译的其他信息，如扩展文件的源文件地址，在MMCV中则是存放于./mmcv/ops/csrc/pytorch/中。其他信息如include_dirs,define_macros,extra_compile_args则会在torch/utils/cpp_extension.py:BuildExtension一起构成终究的 gcc /nvcc 的指令。
class BuildExtension(build_ext, object):
# 只显示中心代码
def build_extensions(self):
# 查看二进制接口兼容性
self._check_abi()
# 注册 cuda 代码 (.cu, .cuh)
self.compiler.src_extensions += ['.cu', '.cuh']
def unix_wrap_compile(obj, src, ext, cc_args, extra_postargs, pp_opts):
try:
original_compiler = self.compiler.compiler_so
if _is_cuda_file(src):
# 对 cuda 文件调用 nvcc 指令
nvcc = _join_cuda_home('bin', 'nvcc')
self.compiler.set_executable('compiler_so', nvcc)
if isinstance(cflags, dict):
cflags = cflags['nvcc']
cflags = COMMON_NVCC_FLAGS + ['--compiler-options',
"'-fPIC'"] + cflags + _get_cuda_arch_flags(cflags)
elif isinstance(cflags, dict):
# 默许 c++ 程序的 flags
cflags = cflags['cxx']
# 强制性运用 --std=c++11
if not any(flag.startswith('-std=') for flag in cflags):
cflags.append('-std=c++11')
# c++ / cuda 程序编译进口
original_compile(obj, src, ext, cc_args, cflags, pp_opts)
finally:
# 将之前掩盖的默许编译器复原
self.compiler.set_executable('compiler_so', original_compiler)
以上进程了解清楚之后咱们运转MMCV_WITH_OPS=True python setup.py build_ext --inplace指令。
/usr/local/cuda-10.0/bin/nvcc -DMMCV_WITH_CUDA -I/home-to-/mmcv/mmcv/ops/csrc -I/home-to-//torch/include -I/home-to/torch/include/torch/csrc/api/include -I/home-to/torch/include/TH -I/home-to/torch/include/THC -I/usr/local/cuda-10.0/include -I/home-to/python3.7m -c ./mmcv/ops/csrc/pytorch/nms_cuda.cu -o build/temp.linux-x86_64-3.7/./mmcv/ops/csrc/pytorch/nms_cuda.o -D__CUDA_NO_HALF_OPERATORS__ -D__CUDA_NO_HALF_CONVERSIONS__ -D__CUDA_NO_HALF2_OPERATORS__ --expt-relaxed-constexpr --compiler-options '-fPIC' -DTORCH_API_INCLUDE_EXTENSION_H -DTORCH_EXTENSION_NAME=_ext -D_GLIBCXX_USE_CXX11_ABI=0 -gencode=arch=compute_61,code=sm_61 -std=c++11
gcc -pthread -B compiler_compat -Wl,--sysroot=/ -Wsign-compare -DNDEBUG -g -fwrapv -O3 -Wall -Wstrict-prototypes -fPIC -DMMCV_WITH_CUDA -I/home-to/mmcv/ops/csrc -I/home-to/torch/include -I/home-to/torch/include/torch/csrc/api/include -I/home-to/torch/include/TH -I/home-to/torch/include/THC -I/usr/local/cuda-10.0/include -I/home-to/python3.7m -c ./mmcv/ops/csrc/pytorch/nms.cpp -o build/temp.linux-x86_64-3.7/./mmcv/ops/csrc/pytorch/nms.o -DTORCH_API_INCLUDE_EXTENSION_H -DTORCH_EXTENSION_NAME=_ext -D_GLIBCXX_USE_CXX11_ABI=0 -std=c++11
...
在上述运转成果中咱们能够看到
- 编译器对 CUDA 文件主动调用 nvcc 而对 .cpp 文件则调用 gcc
- 被
CUDAExtension包装往后体系主动加入了 Python，PyTorch， CUDA 等库中的头文件以及库地址，体系架构信息(-gencode)与编译优化信息(-O3等) - 经过
-DTORCH_EXTENSION_NAME=_ext将TORCH_EXTENSION_NAME宏赋值为_ext。这看来也绝非是闲来之笔，欲知后事怎么，咱们看下一节分化
3. PYBIND11_MODULE — Python 与 C++ 的桥梁
上文提到经过 setup.py 咱们编译了扩展文件。可是目前依然有个疑问，为什么编译出来的 C++ / CUDA 二进制文件能够在 Python 中直接被调用呢？再次检测编译的一切文件，发现其间有个文件pybind.cpp非常可疑，其打开后大致方法如下。
#include <torch/extension.h>
// 函数声明，详细完结在其他文件
Tensor nms(Tensor boxes, Tensor scores, float iou_threshold, int offset);
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
m.def("nms", &nms, "nms (CPU/CUDA) ", py::arg("boxes"), py::arg("scores"),
py::arg("iou_threshold"), py::arg("offset"));
}
这里PYBIND11_MODULE是一个宏，界说在 pybind11 库中(见pybind11/include/pybind11/pybind11.h)。而 pybind11 是一个用来在 C++ 代码中创建 Python的连接的库。找到了源头，咱们进一步剖析。
这里PYBIND11_MODULE的作用是为 C++ 代码接入 Python 解释器供给进口。以上述代码为例，TORCH_EXTENSION_NAME正是在上文gcc编译进程中呈现的宏，对应为extension的name变量。因而在这里会被解释成_ext（留意没有双引号） 。m则代表 TORCH_EXTENSION_NAME 所对应的模块实例（实践上能够指定为任何姓名）。{}中的每个m.def都界说了一个_ext的成员函数，其一般方法为m.def("函数名"，详细 C++ 完结的函数指针, "文档", 参数列表)。经过这种方法，nms也就顺畅地成为了mmcv._ext的成员函数，其详细完结为现已界说好的nms函数（对这个函数的剖析咱们会鄙人节讲到）。在 Python 中也就能够运转from mmcv._ext import nms了。如果对这里的界说依然不清楚，咱们能够把该宏用 C++ 编译器打开一下：
Tensor nms(Tensor boxes, Tensor scores, float iou_threshold, int offset);
static void pybind11_init__ext(pybind11::module &);
extern "C" __attribute__ ((visibility("default"))) PyObject *PyInit__ext()
{
// 省略部分代码
auto m = pybind11::module("_ext"); // m 变量的初始化是在宏内部
try { pybind11_init__ext(m); return m.ptr(); }
}
void pybind11_init__ext(pybind11::module &m) {
// 添加成员函数
m.def("nms", &nms, "nms (CPU/CUDA) ", py::arg("boxes"), py::arg("scores"),
py::arg("iou_threshold"), py::arg("offset"));
}
其间PyObject *PyInit_界说在Python.h中，正是 C++ 中声明 Python module 的官方方法（可见官方 Python文档）。这里PyInit_后接的_ext其实就是TORCH_EXTENSION_NAME宏解释得到。 意味着新声明晰一个 名为_ext的 Python module。
4. cpp/cu文件 — 算子的详细完结
经过对PYBIND11_MODULE的剖析后，咱们了解了mmcv._ext.nms详细的完结是一个声明为Tensor nms(Tensor boxes, Tensor scores, float iou_threshold, int offset);的函数。该函数界说在mmcv/ops/csrc/pytorch/nms.cpp中
#include <torch/extension.h>
Tensor nms(Tensor boxes, Tensor scores, float iou_threshold, int offset) {
if (boxes.device().is_cuda()) {
// cuda 完结
return nms_cuda(boxes, scores, iou_threshold, offset);
} else {
// c++ 完结
return nms_cpu(boxes, scores, iou_threshold, offset);
}
}
能够看到这时实践的完结方法针对设备的不同分为了nms_cuda与nms_cpu两种。这里咱们先来看 cpp 的完结。
4.1 CPP 算子完结
#include <torch/extension.h>
using namespace at; // 恰当改写
Tensor nms_cpu(Tensor boxes, Tensor scores, float iou_threshold, int offset) {
// 仅显示中心代码
for (int64_t _i = 0; _i < nboxes; _i++) {
// 遍历一切检测框，称为主检测框
if (select[_i] == false) continue;
for (int64_t _j = _i + 1; _j < nboxes; _j++) {
// 对每个主检测框，遍历其他检测框，称为次检测框
// 这里只用遍历上三角元素即可，节约核算
if (select[_j] == false) continue;
auto ovr = inter / (iarea + areas[j] - inter);
// 如果次检测框与主检测框 iou 过大，则去除该次检测框
if (ovr >= iou_threshold) select[_j] = false;
}
}
return order_t.masked_select(select_t);
}
以上即为nms_cpu的中心代码，对该算法想要有进一步了解的同学能够参看源码。这里呈现了两个for循环，完结上这正是咱们希望完结nms的 C++ / CUDA 扩展的原因。关于有一定 C++ 根底的同学来说代码应该较好理解 （留意这里int64_t可理解为C99规约的为支持不同渠道的int64类型的typedef界说，可直接理解为int64），但这里一起也呈现了一些新的变量类型，最典型的是Tensor数据类型。
其实这里Tensor数据类型由torch/extension.h支持，来源于 pytorch 中 C++ API 中三大 namespace（at,torch与c10)中的at。
小知识点：
at,torch与c10这三个命名空间中at代表 ATen (A Tensor Library)，担任声明和界说Tensor运算相关的逻辑，是pytorch扩展c++接口中最常用到的命名空间，c10(Caffe Tensor Library)其实是 ATen 的根底，包含了PyTorch的中心抽象、Tensor和Storage数据结构的实践完结。torch命名空间下界说的 Tensor 相比于ATen 添加主动求导功用，但 c++ 扩展中一般不常见）\
该类型功用非常强壮，基本支持 PyTorch 中 Tensor 的一切运算方法（如 +, -, *, /, >, < 等运算符，.view,.reshape,.unsqueeze等维度改变功用等）。Tensor 的 API 接口可见官方链接。 当然除了 Tensor 类型外at命名空间也支持几乎一切和 Tensor 有关的函数 （如at::ones,at::zeros，at::where等）， ATen 的 API 接口可见官方链接。基本上只要在程序中加入了#include <torch/extension.h>就能够在 C++ 中调用一切 PyTorch 支持的功用。
4.2 CUDA 算子完结
4.2.1 (番外篇) CUDA 编程根底
该部分内容部分参阅CUDA编程入门极简教程，感兴趣的同学能够看原文。
基本概念
CUDA 是建立在 NVIDIA GPU上的一个通用并行核算渠道和编程模型，CUDA编程能够利用GPUs 的并行核算引擎来愈加高效地处理比较杂乱的核算难题。CUDA 的语法和 C++ 大多部分情况下是共同的，其默许文件名后缀是 .cu，默许头文件名后缀是 .cuh。CUDA 编程是异构的，即CPU担任处理逻辑杂乱的串行程序，而 GPU 要点处理数据密集型的并行核算程序，从而发挥最大成效。其间 CPU 地点方位称为为主机端（host），而 GPU 地点方位称为设备端（device）。
CUDA程序的规划流程
一般来说，CUDA 程序履行会按照如下流程：
- 分配 host 内存，并进行数据初始化
- 分配 device 内存，并从 host 将数据拷贝到 device 上
- 调用 CUDA 的核函数在 device 上完结指定的运算
- 将 device 上的运算成果拷贝到 host 上
- 开释 device 和 host 上分配的内存
而对 PyTorch 的 CUDA 扩展来说， CUDA 扩展传入和传出的 Tensor 都现已在 GPU 上，因而这里的 5 个步骤只要第 3 步了，这会为咱们省下比较宝贵的时刻而将更多留意力放到详细的程序完结上。
CUDA 中指定函数设备关键字
由于 CUDA 编程为异步，因而函数的界说与调用很可能不在同一个 device 上面，因而 CUDA 中骑过添加额定函数类型来规约函数的界说与调用设备。 –__global__：在 device 上履行，从 host 中调用（一些特定的 GPU 也能够从 device 上调用），返回类型有必要是 void，不支持可变参数参数，不能成为类成员函数。留意用__global__界说的 kernel 是异步的，这意味着 host 不会等候 kernel 履行完就履行下一步。 –__device__：在 device 上履行，单仅能够从 device 中调用，不能够和__global__一起用。 –__host__：在 host 上履行，仅能够从 host 上调用，一般省略不写，不能够和__global__一起用，但可和__device__，此刻函数会在 device 和 host 都编译。
CUDA 中线程逻辑架构方法
一旦一个 kernel 在 device 上履行，device 上许多经量级的线程会被发动，一个 kernel 所发动的一切线程分成两级架构。一切线程归为一个网格（grid），同一个网格上的线程共享相同的全局内存空间，而网格又能够分为许多线程块（block），一个线程块里面包含许多线程。线程两层安排结构如下图所示，这是一个 gird 和 block 均为 2-dim 的线程安排。grid 和 block 都是界说为 dim3 类型的变量，dim3 能够看成是包含三个无符号整数（x，y，z）成员的结构体变量，在界说时，可界说为一维或二维，剩下维度短少值为 1。当然这里的 grid, block 层次划分实践上仅仅逻辑层次，线程在 GPU 中的流处理器 (SM) 中是用“线程束”办理的，一个线程束包含 32 个线程。因而一般在规划 block 时要保证其线程个数为 32 的整数倍。
为了更好地理解这里的线程架构，咱们能够直接将一个kernel 拓荒的一切线程理解为一个小区，这个小区就被称为 grid，而该小区 (grid) 是由不同楼栋 (block)组成的，每个楼栋 (block)有其在小区内的三维坐标 (x, y, z)。在每个楼栋中的一切线程按其在该 block 的三维坐标 (x, y, z)来进行索引。
CUDA 中核函数调用
核函数 (kernel) 是在 device上线程中并行履行的函数，核函数用__global__符号声明，在调用时需求用<<<grid, block>>>来指定 kernel 要履行的线程数量。这里grid与block都需求提前界说好，在 CUDA 中，每一个线程都要履行核函数，而且每个线程会分配一个唯一的线程号 thread ID，这个 ID 值能够经过核函数的内置变量 threadIdx 来取得。下面代码即为在上图线程逻辑架构下的核函数调用方法。
dim3 grid(3, 2);
dim3 block(5, 3);
kernel_fun<<< grid, block >>>(prams...);
CUDA 中核函数编写
核函数调用进程中将需求并行履行的部分用不同的的线程进行完结。因而在实践 CUDA 的核函数中，体系界说了两个两个内置的坐标变量blockIdx与threadIdx来唯一标识一个线程，它们都是 dim3 类型变量（包含x,y,z成员），其间blockIdx指明线程地点grid中的方位，而threaIdx指明线程地点block中的方位，这里 grid 与 block 正是在核函数调用进程中界说好的，在核函数的界说中也有 dim3 类型变量gridDim与blockDim来别离指定 grid 与 block 的维度。 如下核函数为矩阵相加的 CUDA 代码。程序履行进程中会按blockIdx与threadIdx的坐标信息将该核函数分配给不同的线程来完结，因而完结高效并行化核算。以下为一个较为典型的矩阵相加的核函数规划。
__global__ void MatAdd(float A[N][N], float B[N][N], float C[N][N])
{
int i = blockIdx.x * blockDim.x + threadIdx.x;
int j = blockIdx.y * blockDim.y + threadIdx.y;
if (i < N && j < N)
C[i][j] = A[i][j] + B[i][j];
}
至此咱们完结了一般 CUDA 算子完结的根底，鄙人一末节中咱们再来剖析 nms CUDA 算子的实例。
4.2.2 CUDA 算子实例
// 以下程序恰当改写，只显示中心代码
#include <torch/extension.h>
using namespace at;
#define DIVUP(m, n) ((m) / (n) + ((m) % (n) > 0))
int const threadsPerBlock = sizeof(unsigned long long int) * 8; // 64
__device__ inline bool devIoU(float const *const a, float const *const b,
const int offset, const float threshold) {
// 界说在 device 上的函数，用于返回iou
// 界说省略
}
__global__ void nms_cuda(const int n_boxes, const float iou_threshold,
const int offset, const float *dev_boxes,
unsigned long long *dev_mask) {
// 核函数
const int row_start = blockIdx.y; // block 在 grid 中的方位
const int col_start = blockIdx.x; // block 在 grid 中的方位
const int tid = threadIdx.x; // thread ID (0-63)
if (row_start > col_start) return;
const int row_size =
fminf(n_boxes - row_start * threadsPerBlock, threadsPerBlock);
const int col_size =
fminf(n_boxes - col_start * threadsPerBlock, threadsPerBlock);
// Shared memory 只会被一个 block 中的一切线程共享
__shared__ float block_boxes[threadsPerBlock * 4];
if (tid < col_size) {
block_boxes[tid * 4 + 0] =
dev_boxes[(threadsPerBlock * col_start + tid) * 4 + 0];
block_boxes[tid * 4 + 1] =
dev_boxes[(threadsPerBlock * col_start + tid) * 4 + 1];
block_boxes[tid * 4 + 2] =
dev_boxes[(threadsPerBlock * col_start + tid) * 4 + 2];
block_boxes[tid * 4 + 3] =
dev_boxes[(threadsPerBlock * col_start + tid) * 4 + 3];
}
__syncthreads();
if (tid < row_size) {
const int cur_box_idx = threadsPerBlock * row_start + tid;
const float *cur_box = dev_boxes + cur_box_idx * 4;
int i = 0;
unsigned long long int t = 0;
int start = 0;
if (row_start == col_start) {
start = tid + 1; // 每个 bbox 只需求和上三角元素核算（节约核算）
}
for (i = start; i < col_size; i++) {
if (devIoU(cur_box, block_boxes + i * 4, offset, iou_threshold)) {
t |= 1ULL << i;
}
}
dev_mask[cur_box_idx * gridDim.y + col_start] = t;
}
}
Tensor NMSCUDAKernelLauncher(Tensor boxes, Tensor scores, float iou_threshold,
int offset) {
// cuda 程序进口
at::cuda::CUDAGuard device_guard(boxes.device()); // 指定默许的显卡
auto order_t = std::get<1>(scores.sort(0, /*descending=*/true));
auto boxes_sorted = boxes.index_select(0, order_t);
int boxes_num = boxes.size(0);
// 这里将
const int col_blocks = DIVUP(boxes_num, threadsPerBlock);
// mask 是用来存储 bboxes 之间两两 iou 是否大于阈值的一个 mask
// 原本长度应该是 (boxes_num, boxes_num)，但这里采用位存储的方法，一个 LongLong 类型能够存取 64
// 个 bool 值，因而存储空间能够缩小64倍，只用拓荒 (boxes_num, boxes_num/64)长度即可。
Tensor mask =
at::empty({boxes_num, col_blocks}, boxes.options().dtype(at::kLong));
dim3 blocks(col_blocks, col_blocks);
dim3 threads(threadsPerBlock); // 每 64 个线程放到一个 block 中，遍历一个LongLong数的一切位
cudaStream_t stream = at::cuda::getCurrentCUDAStream();
// 更完好的核函数调用 <<< blocks, threads, shared_memory, stream>>>
nms_cuda<<<blocks, threads, 0, stream>>>(
boxes_num, iou_threshold, offset, boxes_sorted.data_ptr<float>(),
(unsigned long long*)mask.data_ptr<int64_t>());
at::Tensor mask_cpu = mask.to(at::kCPU);
unsigned long long* mask_host =
(unsigned long long*)mask_cpu.data_ptr<int64_t>();
std::vector<unsigned long long> remv(col_blocks);
memset(&remv[0], 0, sizeof(unsigned long long) * col_blocks);
at::Tensor keep_t =
at::zeros({boxes_num}, boxes.options().dtype(at::kBool).device(at::kCPU));
bool* keep = keep_t.data_ptr<bool>();
for (int i = 0; i < boxes_num; i++) {
int nblock = i / threadsPerBlock;
int inblock = i % threadsPerBlock;
if (!(remv[nblock] & (1ULL << inblock))) {
keep[i] = true;
// set every overlap box with bit 1 in remv
unsigned long long* p = mask_host + i * col_blocks;
for (int j = nblock; j < col_blocks; j++) {
remv[j] |= p[j];
}
}
}
AT_CUDA_CHECK(cudaGetLastError());
return order_t.masked_select(keep_t.to(at::kCUDA));
}
