DSL性能优化
DSL方式进行算子代码实现时,性能优化主要有以下三种方法:
- 避免使用运行时间较长的接口。
当前已知耗时较长的指令有:vrec、vsel、vcmp,对于对性能要求较高的场景,开发者可以通过变换计算公式的方法替换掉耗时较长的接口。例如,计算1/exp(x),可以替换为 exp(-x),可以先求-x,再求指数,从而避免了求倒数的操作。
注意:进行指令替换时,需要同时考虑精度。
示例代码:
# 替换前 res = tbe.dsl.vrec(vsqrt_res) # 替换后 cosh_one = tvm.const(NUM_ONE, "float32") tensor_one = tbe.dsl.broadcast(cosh_one, data_y.shape) res = tbe.dsl.vdiv(tensor_one, vsqrt_res)
- 减少总的计算次数。
注意:变换后的公式要正确、精度要在可接受范围。
例如,在计算(1/vsqrt(x))*data_dy时,可以直接使用data_dy/vsqrt(x):
# 修改前 vsqrt_res = tbe.dsl.vsqrt(num_to_vrsqrt)res = tbe.dsl.vdiv(tvm.const(NUM_ONE, "float32"), vsqrt_res)res = tbe.dsl.vmul(res, data_dy) # 修改后 vsqrt_res = tbe.dsl.vsqrt(num_to_vrsqrt)res = tbe.dsl.vdiv(data_dy, vsqrt_res)
- 减少函数封装。
注意事项:注意单函数变量个数不要超过15个的限制。可以使一些只用一次的中间变量共用一个名字。
#替换前 def _cosh_rec_cloud(data): exp_pos = tbe.dsl.vexp(data) neg_exp = tbe.dsl.vmuls(data, tvm.const(NUM_MINUS_ONE, "float32")) neg_exp_pos = tbe.dsl.vexp(neg_exp) base = tbe.dsl.vadd(exp_pos, neg_exp_pos) base_rec = tbe.dsl.vrec(base) res = tbe.dsl.vmuls(base_rec, tvm.const(NUM_TWO, "float32")) return res cosh_value_rec = _cosh_rec_cloud(data_y) # 修改后 exp_pos = tbe.dsl.vexp(data_y) neg_exp = tbe.dsl.vmuls(data_y, tvm.const(NUM_MINUS_ONE, "float32")) neg_exp_pos = tbe.dsl.vexp(neg_exp) base = tbe.dsl.vadd(exp_pos, neg_exp_pos) base_rec = tbe.dsl.vrec(base)
- 针对DSL实现的算子,tvm.const不单独定义,可提升性能。
省略单独定义tvm.const的步骤,尤其是对于只使用一次的值。在使用的时候直接定义。例如:
# 替换前 cosh_one = tvm.const(NUM_ONE, "float32") tensor_one = tbe.dsl.broadcast(cosh_one, data_y.shape) # 替换后 tensor_one = tbe.dsl.broadcast(tvm.const(NUM_ONE, "float32"),data_y.shape)
父主题: 算子代码实现(TBE DSL)