diff --git a/src/target/intrin_rule.cc b/src/target/intrin_rule.cc index 0ab706704246..ef06f8c34251 100644 --- a/src/target/intrin_rule.cc +++ b/src/target/intrin_rule.cc @@ -33,6 +33,18 @@ namespace codegen { namespace intrin { using tirx::FLowerIntrinsic; +// `tirx.round` is ties-to-even (see include/tvm/tirx/op.h), and constant +// folding implements it with std::nearbyint. The C library's round()/roundf() +// is ties-AWAY-from-zero, so lowering through FloatSuffix would disagree with +// the folder and with every other backend. Rename to nearbyint before the +// float suffix is applied, as the CUDA rule already does. +struct FloatSuffixTiesToEven { + std::string operator()(const PrimType& ty, std::string name) const { + if (name == "round") name = "nearbyint"; + return FloatSuffix()(ty, name); + } +}; + TVM_REGISTER_OP("tirx.exp") .set_attr("default.FLowerIntrinsic", DispatchPureExtern); @@ -115,7 +127,8 @@ TVM_REGISTER_OP("tirx.ceil") .set_attr("default.FLowerIntrinsic", DispatchPureExtern); TVM_REGISTER_OP("tirx.round") - .set_attr("default.FLowerIntrinsic", DispatchPureExtern); + .set_attr("default.FLowerIntrinsic", + DispatchPureExtern); TVM_REGISTER_OP("tirx.nearbyint") .set_attr("default.FLowerIntrinsic", DispatchPureExtern); diff --git a/tests/python/codegen/test_target_codegen_c_host.py b/tests/python/codegen/test_target_codegen_c_host.py index 989dc2129f6d..fbad83845063 100644 --- a/tests/python/codegen/test_target_codegen_c_host.py +++ b/tests/python/codegen/test_target_codegen_c_host.py @@ -193,7 +193,17 @@ def check_c(): fround = m["test_round"] dev = tvm.cpu(0) n = nn - a = tvm.runtime.tensor(np.random.rand(n).astype("float32"), dev) + # Exact midpoints first: this is where ties-to-even (np.round, and the + # semantics every other backend and the constant folder use) differs + # from ties-away-from-zero. np.random.rand never produces them, so the + # random tail alone cannot exercise the tie rule. + midpoints = np.array( + [0.5, 1.5, 2.5, 3.5, -0.5, -1.5, -2.5, -3.5], dtype="float32" + ) + a_np = np.concatenate( + [midpoints, np.random.rand(n - len(midpoints)).astype("float32")] + ).astype("float32") + a = tvm.runtime.tensor(a_np, dev) b = tvm.runtime.tensor(np.zeros(n, dtype="float32"), dev) fround(a, b) tvm.testing.assert_allclose(b.numpy(), (np.round(a.numpy()).view("float32")))