Add ccode for CeilToInt and IntTrueDiv (#151375)

Summary: As titled

Test Plan: Test in D73052653 -- shape calculator generates successfully

Differential Revision: D73073845

Pull Request resolved: https://github.com/pytorch/pytorch/pull/151375
Approved by: https://github.com/kalpit-meta-1, https://github.com/Skylion007
This commit is contained in:
Sidney Tsang
2025-04-16 16:47:55 +00:00
committed by PyTorch MergeBot
parent 6a3a6d22dc
commit cedcdda0ed

View File

@ -535,6 +535,10 @@ class CeilToInt(sympy.Function):
if isinstance(number, sympy.Number):
return sympy.Integer(math.ceil(float(number)))
def _ccode(self, printer):
number = printer.parenthesize(self.args[0], self.args[0].precedence - 0.5)
return f"ceil({number})"
class FloorToInt(sympy.Function):
is_integer = True
@ -1127,6 +1131,11 @@ class IntTrueDiv(sympy.Function):
if isinstance(base, sympy.Integer) and isinstance(divisor, sympy.Integer):
return sympy.Float(int(base) / int(divisor))
def _ccode(self, printer):
base = printer.parenthesize(self.args[0], PRECEDENCE["Atom"] - 0.5)
divisor = printer.parenthesize(self.args[1], PRECEDENCE["Atom"] - 0.5)
return f"((int){base}/(int){divisor})"
# TODO: As an indicator, this != 0 implies == 1 (and vice versa).
# Because we do not have the ability to guard on the stride permutation