check against IntImm instead of Var

apache · Apr 3, 2021 · 0bfc1cc · 0bfc1cc
1 parent 78338c1
commit 0bfc1cc
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/python/tvm/topi/cuda/scan.py b/python/tvm/topi/cuda/scan.py
@@ -105,7 +105,7 @@ def exclusive_scan_ir(data, output, reduction=None, binop=tvm.tir.generic.add, i
         size_cast_dtype = "float64"
         target = tvm.target.Target.current()
 
-        if "vulkan" in str(target) and isinstance(scan_axis_size, tvm.tir.expr.Var):
+        if "vulkan" in str(target) and not isinstance(scan_axis_size, tvm.tir.expr.IntImm):
             # SPIRV seems to have an issue with float64 intrinsic
             # TODO(masahi): Eliminate this concern by adding TIR level CSE
             msg = """