Skip to content

Commit

Permalink
Raise an exception when extern function does not return Stmt (#5964)
Browse files Browse the repository at this point in the history
The function for tvm.te.extern should return either PrimExpr or Stmt,
however there is no check if it actually does so. If it does not, the
result may be a segmentation fault later on. Catch this case early on,
so an informative message can be shown.
  • Loading branch information
Krzysztof Parzyszek authored Jun 30, 2020
1 parent 9efe119 commit 44194d0
Showing 1 changed file with 2 additions and 0 deletions.
2 changes: 2 additions & 0 deletions python/tvm/te/operation.py
Original file line number Diff line number Diff line change
Expand Up @@ -313,6 +313,8 @@ def extern(shape,
body = fcompute(input_placeholders, output_placeholders)
if isinstance(body, tvm.tir.PrimExpr):
body = tvm.tir.Evaluate(body)
if not isinstance(body, tvm.tir.Stmt):
raise ValueError("Function '{}' should return PrimExpr or Stmt".format(fcompute.__name__))

op = _ffi_api.ExternOp(name, tag, attrs,
inputs, input_placeholders,
Expand Down

0 comments on commit 44194d0

Please sign in to comment.