Anndrey24 commented on code in PR #17003:
URL: https://github.com/apache/tvm/pull/17003#discussion_r1608503953
##########
tests/python/topi/test_topi_conv2d_nhwc.py:
##########
@@ -117,21 +140,34 @@ def test_conv2d_nhwc_gemm_fp32(device, ref_data, dtype,
stride, padding, dilatio
A = te.placeholder(a_np.shape, name="A", dtype=dtype)
W = te.placeholder(w_np.shape, name="W", dtype=dtype)
- target, compute, schedule = device
- dev = tvm.device(target, 0)
+ target_string, compute, schedule, use_tir_schedule = device
+ dev = tvm.device(target_string, 0)
+ target = tvm.target.Target(target_string)
- with tvm.target.Target(target) as target:
- B = compute(A, W, stride, padding, dilation, dtype)
- s = schedule([B])
+ if (target.features.has_sve and llvm_version_major() < 15) or (
+ target.features.has_sme and llvm_version_major() < 16
+ ):
+ return
+
+ with target:
a = tvm.nd.array(a_np, dev)
w = tvm.nd.array(w_np, dev)
+ B = compute(A, W, stride, padding, dilation, dtype)
b = tvm.nd.array(np.zeros(get_const_tuple(B.shape), dtype=B.dtype),
dev)
- func = tvm.build(s, [A, W, B], target)
+ if use_tir_schedule:
+ primfunc = te.create_prim_func([A, W, B],
index_dtype_override="int64")
Review Comment:
I think it had to do with me initially using `tile_N = T.int64(tile_N)`
inside the schedule since that seems to be the index dtype when lowering from
Relay (e.g. in `test_conv2d.py`), so I also made the TOPI tests use int64.
However, after changing the line to `tile_N = T.cast(tile_N, N_padded.dtype)`
this is not an issue anymore, so I will remove it!
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
