京东|Ubuntu系统下太极图形语言的初始化参数和ti.kernel及函数

京东|Ubuntu系统下太极图形语言的初始化参数和ti.kernel及函数

【京东|Ubuntu系统下太极图形语言的初始化参数和ti.kernel及函数】在上文中 , 我们在Ubuntu系统中安装了太极(taichi)编程语言 , 也就是嵌入Python的第三方库 。 当然 , 由于机智客这台Ubuntu主机运行过深度学习等多种应用 , 所以安装了很多库 , 并不存在依赖问题 , 而有的时候 , 会出现依赖问题 , 这个按照报错提示安装就行了 。 说回来 , 而我们要使用太极的时候除了引用外 , 还要初始化以及用@ti.kernel这个命令 。 有了上述命令 , 才能在太极图形编程语言里使用函数 。 因此 , 谈下初始化和kernel内核的问题 。
Python在用import引入taichi库后 , 就需要在接下来的代码里初始化 。 命令是ti.init(arch=ti.gpu) , 之前我们也说了 , 这个初始化 , 有几种写法 。 简单点 , 不加参数 。 一般加参数可以指定CPU或GPU , 比如上面这个命令就是指定了GPU 。 虽然很多时候这样默认就行 , 不过这个还有个参数就是分配内存 , 具体地说在指定GPU以后可以设置分配显存 。

比如我想让代码使用一部分GPU显存 , 那么就可以另外指定 。 官方英文文档里(说是简体中文 , 机智客看的时候它并没有完全翻译)提到两个参数:device_memory_GB和device_memory_fraction 。 前者是分配显存大小 , 后者是分配显存比例 , 举例device_memory_GB=6.6代表分配6.6GB的显存;device_memory_fraction=0.5则代表分配50%的显存 。 当然 , 其实很多时候 , 都是默认不指定分配显存 , 独占就行了 。
当然 , 初始化代码还不仅只是设置这些 , 还可以指定数据类型 。 毕竟 , 初始化嘛 , 肯定也要涉及到数据 , 只不过我们默认就行 。 数据类型的内容我们放在以后再学习 。 初始化代码指定了 , 之后进入太极语言的代码就要用到@ti.kernel这个命令了 。 具体地说kernel是太极内核 , 有了它的修饰 , 下面执行的代码才是太极内核执行的 。 也就是太极语言的作用域 , 是太极编译器执行的 。 因此kernel之外的代码属于Python 。 对于使用CUDA的用户 , 太极的内核kernel其实就等于__global__函数 。
在之前的文章里 , 我们提到ti.kernel命令并没有说参数 。 而根据官方的文档介绍 , 太极内核可以最多有8个参数 。 这个就方便我们从Python向太极传递了 。 比如我们可以指定参数的类型 。 而太极的类型则包括i32 , f32多个类型 。 而参数类型也有要求 , 只接受上述的标量 , 而不支持矢量或者矩阵类型 。 比如参数中支持ti.i32或ti.f32 , 但不支持ti.Matrix 或者 ti.Vector 。 而返回值(如果有)也需要标注类型 。
另外 , 太极编程语言除了ti.kernel还有一个ti.func 。 顾名思义 , 在其中的函数则是太极函数 , 也是可以被ti.kernel内核调用的函数 。 而对于使用CUDA的用户而言 , 太极的函数Taichi functions其实就等于__device__ functions 。 太极函数可以嵌套使用 , 这里机智客就不多介绍这个了 , 用法类似 。 这次介绍了太极代码的简单结构 , 以后就可以以这样的结构嵌入到Python脚本里撰写代码了 。