From 70409b7965278d2191921761853334a1b7b8190c Mon Sep 17 00:00:00 2001 From: Krzysztof Parzyszek Date: Tue, 30 Jun 2020 12:58:58 -0500 Subject: [PATCH] Raise an exception when extern function does not return Stmt (#5964) 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. --- python/tvm/te/operation.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/python/tvm/te/operation.py b/python/tvm/te/operation.py index 3ccab5bfd9c3..b03c6f6f0258 100644 --- a/python/tvm/te/operation.py +++ b/python/tvm/te/operation.py @@ -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,