vsel
Description
Compares the condition elements with True or 1 based on its data type. If the expression is true, the value of x is returned. Otherwise, the value of y is returned.
- If condition is a bool, the condition elements are compared with True. If the expression is true, the value of x is returned. Otherwise, the value of y is returned.
- If condition is a uint8, the condition elements are compared with 1 bitwise. If the expression is true, the value of x is returned. Otherwise, the value of y is returned.
The following uses expressions to explain the comparison operation. i indicates an element in the condition, x' indicates x's element or x, y' indicates y's element or y, z indicates an element of the result tensor, and n (a value ranging from 0 to 7) indicates the bit index of the condition element. The expression is as follows:
- When the data type of condition is of type bool:
z = (i==True)?x':y'
- When the data type of condition is of type uint8:
z = (i[n]==1)? x':y'
Prototype
vsel(condition, lhs, rhs)
Parameters
- condition: a tvm.tensor for the condition tensor of type bool or uint8.
- lhs: a tvm.tensor or scalar for the possible return value
- rhs: a tvm.tensor or scalar for the possible return value
lhs and rhs must have the same data type.
Returns
res_tensor: a tvm.tensor for the result tensor
Restrictions
- When condition is a uint8, the last dimensions of the shapes of lhs and rhs must be divisible by 8.
- When condition is a bool and the tbe.dsl.build API is called for build, bool_storage_as_1bit in the passed configuration parameter must be set to False. Otherwise, unexpected output shape will be obtained.
bool_storage_as_1bit is set to True by default, indicating that it is stored as 1-bit data.
The following gives a build configuration file template.
with tvm.target.cce(): schedule = tbe.dsl.auto_schedule(res) config = {"name": kernel_name, "tensor_list": [data_x, data_y, res], "bool_storage_as_1bit": False} tbe.dsl.build(schedule, config)