-
Notifications
You must be signed in to change notification settings - Fork 3.5k
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
[CODEGEN][OPENCL] Fix compile error about ternary expression. (#2821)
Code like this can't be built with NV OpenCL, and it needs an explicit type converison for ternary expression if return type is uchar. uchar i = 0, j = 0; uchar t = max((uchar)j, ((i > 0) ? (uchar)1 : (uchar)0));
- Loading branch information
1 parent
0f6989f
commit fa70983
Showing
3 changed files
with
76 additions
and
0 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,55 @@ | ||
import tvm | ||
|
||
target = 'opencl' | ||
|
||
def test_opencl_ternary_expression(): | ||
def check_if_then_else(ctx, n, dtype): | ||
A = tvm.placeholder((n,), name='A', dtype=dtype) | ||
true_value = tvm.const(1, dtype=dtype) | ||
false_value = tvm.const(3, dtype=dtype) | ||
max_lhs = tvm.const(2, dtype=dtype) | ||
max_rhs = tvm.if_then_else(A[0] > 0, true_value, false_value) | ||
C = tvm.compute((n,), lambda i: tvm.max(max_lhs, max_rhs), name='C') | ||
s = tvm.create_schedule(C.op) | ||
s[C].bind(s[C].op.axis[0], tvm.thread_axis("threadIdx.x")) | ||
fun = tvm.build(s, [A, C], target) | ||
|
||
a = tvm.nd.empty((n,), A.dtype, ctx) | ||
c = tvm.nd.empty((n,), A.dtype, ctx) | ||
# Only need to test compiling here | ||
fun(a, c) | ||
|
||
def check_select(ctx, n, dtype): | ||
A = tvm.placeholder((n,), name='A', dtype=dtype) | ||
true_value = tvm.const(1, dtype=dtype) | ||
false_value = tvm.const(3, dtype=dtype) | ||
max_lhs = tvm.const(2, dtype=dtype) | ||
max_rhs = tvm.expr.Select(A[0] > 0, true_value, false_value) | ||
C = tvm.compute((n,), lambda i: tvm.max(max_lhs, max_rhs), name='C') | ||
s = tvm.create_schedule(C.op) | ||
s[C].bind(s[C].op.axis[0], tvm.thread_axis("threadIdx.x")) | ||
fun = tvm.build(s, [A, C], target) | ||
|
||
a = tvm.nd.empty((n,), A.dtype, ctx) | ||
c = tvm.nd.empty((n,), A.dtype, ctx) | ||
# Only need to test compiling here | ||
fun(a, c) | ||
|
||
if not tvm.module.enabled(target): | ||
print("skip because opencl is not enabled..") | ||
return | ||
|
||
ctx = tvm.context(target, 0) | ||
|
||
check_if_then_else(ctx, 1, 'int8') | ||
check_if_then_else(ctx, 1, 'uint8') | ||
check_if_then_else(ctx, 1, 'int16') | ||
check_if_then_else(ctx, 1, 'uint16') | ||
check_select(ctx, 1, 'int8') | ||
check_select(ctx, 1, 'uint8') | ||
check_select(ctx, 1, 'int16') | ||
check_select(ctx, 1, 'uint16') | ||
|
||
|
||
if __name__ == "__main__": | ||
test_opencl_ternary_expression() |