MindSpore Hybrid 语法规范

概述

MindSpore Hybrid DSL的语法与Python语法类似，例如函数定义、缩进和注释。把MindSpore Hybrid DSL书写的函数加上kernel装饰器后可以当做普通的numpy函数使用，也可以用于Custom的进行自定义算子。

[1]:

import numpy as np
import mindspore as ms
import mindspore.ops as ops
from mindspore.ops import kernel

@kernel
def outer_product(a, b):
    d = allocate(a.shape, a.dtype)
    c = output_tensor(a.shape, a.dtype)

    for i0 in range(a.shape[0]):
        for i1 in range(b.shape[1]):
            c[i0, i1] = 0.0
            for i2 in range(a.shape[1]):
                d[i0, i2] = 2 * a[i0, i2]
                c[i0, i1] = c[i0, i1] + sin(d[i0, i2] * b[i2, i1])
    return c

np_x = np.random.normal(0, 1, [4, 4]).astype(np.float32)
np_y = np.random.normal(0, 1, [4, 4]).astype(np.float32)

print(outer_product(np_x, np_y))

input_x = ms.Tensor(np_x)
input_y = ms.Tensor(np_y)

test_op_akg = ops.Custom(outer_product)
out = test_op_akg(input_x, input_y)
print(out)

[[-0.7582229   1.9742808  -1.5035899   1.6295254 ]
 [ 0.18717238 -1.1390371  -0.92540735  0.25755903]
 [-0.75234073  0.2182185   0.9805498   0.27473617]
 [ 0.7546873  -0.8488003   0.58964515 -0.23971215]]
[[-0.758223    1.9742805  -1.5035899   1.6295254 ]
 [ 0.18717244 -1.1390371  -0.9254071   0.2575591 ]
 [-0.7523403   0.21821874  0.9805499   0.27473587]
 [ 0.75468683 -0.84879947  0.5896454  -0.23971221]]

语法规则

变量

MindSpore Hybrid DSL中的变量包括Tensor和Scalar两种形式。

对于Tensor类型的变量，除了在输入中提供的变量，其他变量都需要在使用前申明 shape和 dtype。

对于输出Tensor使用 output_tensor，用法为：output_tensor(shape, dtype)。
对于中间结果使用 allocate，用法为：allocate(shape, dtype)。

Tensor分配的示例代码如下：

[2]:

@kernel
def kernel_func(a, b):
    # a和b作为输入tensor，可以直接使用

    # d为一个数据类型为fp16,形状为(2,)的Tensor，在下面的code中作为中间变量使用
    d = allocate((2,), "float16")
    # c为一个数据类型与b相同,形状与a相同的Tensor，在下面的code中作为函数输出使用
    c = output_tensor(a.shape, b.dtype)

    # d作为中间变量，给c赋值
    d[0] = b[0, 0]
    for i in range(4):
        for j in range(4):
            c[i, j] = d[0]

    # c作为输出
    return c

对于Scalar类变量，会将它第一次的赋值运算作为声明。赋值操作可以是一个立即数，也可以是一个计算表达式。Scalar类变量第一次赋值的地方决定了它的定义域（例如，某一个for loop之内），在定义域之外使用Scalar变量会报错。

Scalar变量使用的示例代码如下：

[3]:

@kernel
def kernel_func(a):
    c = output_tensor(a.shape, a.dtype)

    for i in range(10): # i loop
        for j in range(5): # j loop
            # 用一个立即数给Scalar赋值
            d = 2.0
            # 用表达式给Scalar赋值
            e = a[i, j]
            # 正常使用scalar
            c[i, j] = d + e

    # Wrong: c[0, 0] = d
    # 不能在超出Scalar d的定义域（j loop）之外的范围使用

    return c

与原生Python语言不同的是，变量一旦创建，shape和 dtype就不能改变。

计算表达

MindSpore Hybrid DSL支持基本的四则运算表达，包括 +, -, *, /，及赋值运算符，包括 =, +=, -=, *=, /=。用户可以像写Python表达一样书写计算表达式利用变量计算和为变量赋值。

所有的计算需要基于标量计算，如果是Tensor对象那么写清楚所有index，即 C[i, j] = A[i, j] + B[i, j]。当前不支持 C = A + B这种向量化的写法。

在书写计算表达式时，用户需要自行负责类型的合法性。表达式左右两边的类型需要保持一致，否则在算子编译环节会报错。计算式中的整数立即数会被认定为int32，而浮点立即数会被认定为float32。MindSpore Hybrid DSL不提供任何隐式的类型转化，所有类型转化都需要显式的书写出来。类型名即对应类型转换函数的名字，包括：

int32
float16
float32
(仅GPU后端)int8, int16, int64, float64

类型转换代码示例如下：

[4]:

@kernel
def kernel_func(a):
    c = output_tensor((2,), "float16")

    # Wrong: c[0] = 0.1 此处c的类型为fp16, 而0.1的类型为fp32
    c[0] = float16(0.1) # float16(0.1)把表达式的类型转化为fp16
    c[1] = float16(a[0, 0]) # float16(a[0, 0])把表达式的类型转化为fp16

    return c

循环

当前只支持 for loop，不支持 while、 break、 continue关键词。

基本循环的写法和Python一样，循环维度的表达可以使用 range和 grid关键词。range表示一维的循环维度，接受一个参数表示循环的上限，例如：

[5]:

@kernel
def kernel_func(a, b):
    c = output_tensor((3, 4, 5), "float16")

    for i in range(3):
        for j in range(4):
            for k in range(5):
                out[i, j, k] = a[i, j, k] + b[i, j, k]
    return  c

则循环表达的计算空间为 0 <= i < 3, 0 <= j < 4, 0 <= k < 5。

grid表示多维网格，接受的输入为 tuple ，例如上面的代码用 grid表达后如下：

[6]:

@kernel
def kernel_func(a, b):
    c = output_tensor((3, 4, 5), "float16")

    for arg in grid((4, 5, 6)):
        out[arg] = a[arg] + b[arg]
    return  c

此时，参数 arg等价于一个三维index (i,j,k)，其上限分别为4，5，6。对参数 arg我们可以取其中的某个分量，例如

[7]:

@kernel
def kernel_func(a, b):
    c = output_tensor((3, 4, 5), "float16")

    for arg in grid((4, 5, 6)):
        out[arg] = a[arg] + b[arg[0]]
    return  c

那么循环内的表达式等价于 out[i, j, k] = a[i, j, k] + b[i]。

调度原语

从1.8版本开始，MindSpore Hybrid DSL 提供调度原语以描述不同类型的循环。在 Ascend 后端，调度原语将协助新 DSA 多面体调度器生成代码。此类调度原语包括：serial， vectorize， parallel，和 reduce。

serial 会提示调度器该循环在调度生成时应保持前后顺序，不要做会改变顺序的调度变换，例如：

[8]:

@kernel
def serial_test(a, b):
    row = a.shape[0]
    for i in serial(row):
        for j in serial(i):
            b[i] = b[i] - a[i, j] * b[j]
    return b

这里 serial 提示 i 和 j 的计算有依赖关系，调度时应保持 i 和 j 从小的大的顺序。

vectorize 一般用于最内层循环，会提示调度器该循环有生成向量化指令的机会，例如：

[9]:

@kernel
def vector_test(a, b):
    out = output_tensor(a.shape, a.dtype)
    row = a.shape[0]
    col = a.shape[1]
    for i in range(row):
        for j in vectorize(col):
            out[i, j] = a[i, j] + b[0, i]
    return out

这里 vectorize 提示最内层 j 轴循环包含同质化计算，调度时可以生成向量化指令加速内层循环。

parallel 一般用于最外层循环，会提示调度器该循环有并行执行机会，例如：

[10]:

@kernel
def parallel_test(a, b):
    out = output_tensor(a.shape, a.dtype)
    row = a.shape[0]
    col = a.shape[1]
    for i in parallel(row):
        for j in range(col):
            out[i, j] = a[i, j] + b[0, j]
    return out

这里 parallel 提示最外层 i 轴循环无依赖关系，调度时可以并行加速。

reduce 会提示调度器该循环为运算中的一个 Reduction 轴，例如：

[11]:

def reduce_test(a):
    out = output_tensor((a.shape[0], ), a.dtype)
    row = a.shape[0]
    col = a.shape[1]
    for i in range(row):
        out[i] = 0.0
        for k in reduce(col):
            out[i] = out[i] + a[i, k]
    return out

这里 reduce 对应的 k 轴为累加轴。

用户在使用调度原语的时候需要注意：

上述调度原语只会在 Ascend 后端影响调度。在CPU和GPU后端，上述调度原语将被处理成普通的 for 循环关键词。
调度原语对于调度器只是提示作用，当调度原语的提示和调度器自身的分析验证相矛盾时，调度器将把上述调度原语将被处理成普通的 for 循环关键词。

属性

当前只支持对Tensor对象属性shape和dtype，例如 a.shape，c.dtype。

一个Tensor的shape属性会表达为一个 tuple，我们可以对它进行固定下标的取分量操作，例如 a.shape[0]。

同时，在 grid关键词中我们接受某个Tensor对象的 shape属性，那么循环的维度由Tensor的维度决定。例如：

[12]:

@kernel
def kernel_func(a, b):
    c = output_tensor(a.shape, "float16")

    for arg in grid(a.shape):
        out[arg] = a[arg] + b[arg[0]]
    return  c

如果a是一个二维Tensor，那么循环内的表达式等价于 out[i, j] = a[i, j] + b[i]。而如果a是一个三维Tensor，那么循环内的表达式等价于 out[i, j, k] = a[i, j, k] + b[i]。

关键词

当前支持的关键词包括

全平台支持数学函数：log、exp、sqrt、tanh、power、floor
内存分配：allocate、 output_tensor
数据类型转化：int32、 float16、 float32、 float64
循环表达：for、 range、 grid
调度源语：serial、 vec、 parallel、 reduce
在当前版本中，我们对CPU/GPU后端提供部分进阶关键词：
- 数学函数：rsqrt、 erf、 isnan、 sin、 cos、 isinf、 isfinite、 atan、 atan2(仅GPU)、 expm1(仅GPU)、 floor、 ceil、 trunc、 round、 ceil_div
- 数据类型转换：int8，int16，int64

常见报错信息及错误归因

为了帮助用户高效地开发和定位bug，MindSpore Hybrid DSL 提供如下报错信息，包括

TypeError: 当使用了while, break 和 continue 等 MindSpore Hybrid DSL 不支持的 Python 关键词。
ValueError:
- 使用了不属于上面的内置函数名；
- 对张量取非 shape 或者 dtype 的属性。
其他常见报错：
- “SyntaxError”: 写的 DSL 不符合基本 Python 语法（非上面的进阶用法中定义的MindSpore Hybrid DSL语法），由 Python 解释器本身报错；
- “ValueError: Compile error”及“The pointer[kernel_mod] is null”: Python DSL符合语法但是编译失败，由 AKG 报错，具体错误原因检查 AKG 相关报错信息；
- “Launch graph failed”: Python DSL符合语法，编译成功但是运行失败。具体原因参考硬件的报错信息。例如在昇腾芯片上遇到运行失败时，MindSpore 端会显示 “Ascend error occurred” 及对应硬件报错信息。

开发用例：利用hybrid类型的自定义算子实现三维张量的加法函数

首先，我们写一个基于MindSpore Hybrid DSL书写一个计算三维张量相加的函数。

注意：

对于输出张量使用 output_tensor，用法为：output_tensor(shape, dtype)；
所有的计算需要基于标量计算，如果是Tensor对象，那么需要写清楚所有index；
基本循环的写法和Python一样，循环维度的表达可以使用 range。

[13]:

import numpy as np
from mindspore import ops
import mindspore as ms
from mindspore.ops import kernel

ms.set_context(device_target="GPU")
@kernel
def tensor_add_3d(x, y):
    result = output_tensor(x.shape, x.dtype)
    #    1. 你需要一个三层循环
    #    2. 第i层循环的上界可以用x.shape[i]获得
    #    3. 你需要基于每个元素表达计算，例如加法为 x[i, j, k] + y[i, j, k]
    for i in range(x.shape[0]):
        for j in range(x.shape[1]):
            for k in range(x.shape[2]):
                result[i, j, k] = x[i, j, k] + y[i, j, k]

    return result

下面我们用上面的函数自定义一个算子。

注意到基于kernel的hybrid函数时，我们可以使用自动的形状和数据类型推导。

因此我们只用给一个func输入（func_type的默认值为"hybrid"）。

[14]:

tensor_add_3d_op = ops.Custom(func=tensor_add_3d)
input_tensor_x = ms.Tensor(np.ones([2, 3, 4]).astype(np.float32))
input_tensor_y = ms.Tensor(np.ones([2, 3, 4]).astype(np.float32) * 2)
result_cus = tensor_add_3d_op(input_tensor_x, input_tensor_y)
print(result_cus)

[[[3. 3. 3. 3.]
  [3. 3. 3. 3.]
  [3. 3. 3. 3.]]

 [[3. 3. 3. 3.]
  [3. 3. 3. 3.]
  [3. 3. 3. 3.]]]

同时我们可以使用pyfunc模式验证上面定义的正确性。

这里我们不需要重新定义算子计算函数tensor_add_3d，直接将func_type改为"pyfunc"即可。

注意pyfunc模式时我们需要手写类型推导函数。

[15]:

def infer_shape_py(x, y):
    return x

def infer_dtype_py(x, y):
    return x

tensor_add_3d_py_func = ops.Custom(func=tensor_add_3d,
                                   out_shape=infer_shape_py,
                                   out_dtype=infer_dtype_py,
                                   func_type="pyfunc")

result_pyfunc = tensor_add_3d_py_func(input_tensor_x, input_tensor_y)
print(result_pyfunc)

[[[3. 3. 3. 3.]
  [3. 3. 3. 3.]
  [3. 3. 3. 3.]]

 [[3. 3. 3. 3.]
  [3. 3. 3. 3.]
  [3. 3. 3. 3.]]]

我们可以得到如下结果，即两个Tensor的和。