Optimize ttnn.round with a direct implementation

[jdh8]

Rounding is only supported by Wormhole, and Wormhole, and Wormhole has the exact function float_to_int16 if the value is in range.
https://github.com/tenstorrent/tt-metal/blob/main/docs/source/tt-metalium/tt_metal/apis/kernel_apis/sfpu/llk.rst#wormhole-only

However, ttnn.round is implemented as a combination of ttnn.floor, ttnn.add, etc.

tt-metal/ttnn/cpp/ttnn/operations/eltwise/unary/device/unary_composite_op.cpp

Lines 640 to 653 in 679b8d5

	} else { // Bankers' Rounding
	Tensor rounded_non_half = ttnn::floor(
	ttnn::add(
	input,
	ttnn::where(ttnn::logical_and(ttnn::ge(input, 0.4), ttnn::le(input, 0.5)), 0.4f, 0.5f, output_mem_config.value()),
	std::nullopt,
	output_mem_config),
	output_mem_config.value());
	Tensor fractional_part = ttnn::subtract(input, floor_res, std::nullopt, output_mem_config);
	Tensor is_half = ttnn::eq(fractional_part, 0.5, std::nullopt, output_mem_config);
	Tensor rounded_half =
	ttnn::add(floor_res, is_odd(floor_res, output_mem_config), std::nullopt, output_mem_config);
	return ttnn::where(is_half, rounded_half, rounded_non_half, output_mem_config.value());
	}

Then in turn, ttnn.floor calls functions that effectively computes ttnn.round.

tt-metal/tt_metal/hw/ckernels/wormhole_b0/metal/llk_api/llk_sfpu/ckernel_sfpu_floor.h

Lines 25 to 26 in 679b8d5

	vInt tmp = float_to_int16(result, 0); //TODO: Replace float_to_int16 to float_to_int32 once it is available
	result = int32_to_float(tmp, 0);

Rounding to a nearest integer is extremely useful for argument reduction. We can reuse a direct implementation in other mathematical functions (mostly elementary functions) such as:

• Exponential functions
• Trigonometric functions
• ttnn.pow

[mouliraj-mcw]

Hi @jdh8 ,
I examined your approach and found that it doesn't address rounding to a specific number of decimal places (i.e., 2 or 3 decimal places).
Could you please share your thoughts on how this could be managed?

[jdh8]

Thanks for pointing it out! I missed the parameter decimals.

It can be managed with multiplication by 10ⁿ. To be specific,

round(x, n) = 10**-n * round(10**n * x)

[jdh8]

I have two proposals:

1. Implement a native roundeven(x) as conceptually round(x, 0), and then make round(x, n) on top of roundeven. (Named after C23 roundeven)
2. Make a direct, native round(x, n).

Which approach looks better?

[mouliraj-mcw]

I think approach two would be more suitable, as it has a straightforward structure.

[umadevimcw]

@jdh8 Tested the rounding in the jdh8/direct-rounding branch with reference to this comment #13851 (review) and observed that

torch round uses banker's rounding algorithm where it follows round nearest even approach, for example

for the input 94.5

Torch result is 94 whereas in
TT the result is 95 which results in test case failure

94.5 is halfway between 94 and 95 so hence the value is rounded to the nearest even and the result becomes 94 which needs to handled in our TT implementation

Please find the image below (Red is TT's output and Green is Torch output)