Skip to content

Commit c6c048c

Browse files
Hardcode84Hardcode84
authored
Improve sympy expr conversion (#80)
1 parent 3a2f7b2 commit c6c048c

5 files changed

Lines changed: 161 additions & 36 deletions

File tree

‎mlir/include/hc/Dialect/Typing/IR/TypingOps.hpp‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,8 @@ class SymbolicTypeBase : public mlir::Type {
3636
};
3737
} // namespace hc::typing
3838

39+
MLIR_DECLARE_EXPLICIT_TYPE_ID(hc::typing::SymbolicTypeBase);
40+
3941
#include "hc/Dialect/Typing/IR/TypingOpsDialect.h.inc"
4042
#include "hc/Dialect/Typing/IR/TypingOpsEnums.h.inc"
4143

‎mlir/lib/Dialect/Typing/IR/TypingOps.cpp‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,8 @@
1616

1717
#include <llvm/ADT/TypeSwitch.h>
1818

19+
MLIR_DEFINE_EXPLICIT_TYPE_ID(hc::typing::SymbolicTypeBase)
20+
1921
namespace {
2022
struct TypingAsmDialectInterface : public mlir::OpAsmDialectInterface {
2123
using OpAsmDialectInterface::OpAsmDialectInterface;

‎python/cpp/compiler/src/PyWrappers.cpp‎

Lines changed: 22 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -38,7 +38,7 @@ template <typename T> static bool genAttrIsA(MlirAttribute attr) {
3838
}
3939

4040
template <typename T> static MlirTypeID genGetTypeID() {
41-
return wrap(T::getTypeID());
41+
return wrap(::mlir::TypeID::get<T>());
4242
}
4343

4444
using namespace mlir;
@@ -205,12 +205,25 @@ class PyExprType : public PyConcreteType<PyExprType> {
205205
}
206206
};
207207

208+
class PySymbolicType : public PyConcreteType<PySymbolicType> {
209+
public:
210+
static constexpr IsAFunctionTy isaFunction =
211+
genTypeIsA<hc::typing::SymbolicTypeBase>;
212+
static constexpr GetTypeIDFunctionTy getTypeIdFunction =
213+
genGetTypeID<hc::typing::SymbolicTypeBase>;
214+
static constexpr const char *pyClassName = "SymbolicType";
215+
using PyConcreteType::PyConcreteType;
216+
217+
static void bindDerived(ClassTy &c) { defExprOperators(c); }
218+
};
219+
208220
template <typename Op>
209-
static PyType makeExpr(mlir::python::PyType lhs, mlir::python::PyType rhs) {
221+
static PySymbolicType makeExpr(mlir::python::PyType lhs,
222+
mlir::python::PyType rhs) {
210223
auto ulhs = mlir::cast<hc::typing::SymbolicTypeBase>(unwrap(lhs.get()));
211224
auto urhs = mlir::cast<hc::typing::SymbolicTypeBase>(unwrap(rhs.get()));
212225
auto res = Op()(ulhs, urhs);
213-
return PyType(lhs.getContext(), wrap(res));
226+
return PySymbolicType(lhs.getContext(), wrap(res));
214227
}
215228

216229
struct FloorDiv {
@@ -225,12 +238,12 @@ struct CeilDiv {
225238
};
226239

227240
template <typename ClassTy> static void defExprOperators(ClassTy &c) {
228-
c.def("__add__", &makeExpr<std::plus<void>>, py::arg("rhs"), "mul op");
229-
c.def("__sub__", &makeExpr<std::minus<void>>, py::arg("rhs"), "mul op");
241+
c.def("__add__", &makeExpr<std::plus<void>>, py::arg("rhs"), "add op");
242+
c.def("__sub__", &makeExpr<std::minus<void>>, py::arg("rhs"), "sub op");
230243
c.def("__mul__", &makeExpr<std::multiplies<void>>, py::arg("rhs"), "mul op");
231-
c.def("__mod__", &makeExpr<std::modulus<void>>, py::arg("rhs"), "mul op");
232-
c.def("__floordiv__", &makeExpr<FloorDiv>, py::arg("rhs"), "mul op");
233-
c.def("__truediv__", &makeExpr<CeilDiv>, py::arg("rhs"), "mul op");
244+
c.def("__mod__", &makeExpr<std::modulus<void>>, py::arg("rhs"), "mod op");
245+
c.def("__floordiv__", &makeExpr<FloorDiv>, py::arg("rhs"), "floordiv op");
246+
c.def("__truediv__", &makeExpr<CeilDiv>, py::arg("rhs"), "ceildiv op");
234247
}
235248

236249
class PyTypeAttr : public PyConcreteAttribute<PyTypeAttr> {
@@ -263,6 +276,7 @@ static void populateTypingTypes(py::module_ &m) {
263276
PyLiteralType::bind(m);
264277
PySymbolType::bind(m);
265278
PyExprType::bind(m);
279+
PySymbolicType::bind(m);
266280

267281
PyTypeAttr::bind(m);
268282
}

‎python/hckernel/kernel.py‎

Lines changed: 133 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@
55

66
import sympy
77
import inspect
8+
from collections import namedtuple
89
from .kernel_api import _verify_kernel_params
910
from .kernel_api import *
1011
from .dispatcher import create_dispatcher
@@ -88,33 +89,139 @@ def _get_typing_module():
8889
return get_typing_module()
8990

9091

91-
def _get_symbol_type(term, ceiling=False):
92-
match term:
93-
case sympy.Symbol():
94-
return typing.SymbolType.get(term.name)
95-
96-
case sympy.Mul():
97-
lhs, rhs = term.args
98-
lhs = _get_symbol_type(lhs)
99-
if isinstance(rhs, sympy.Pow):
100-
arg, p = rhs.args
101-
assert p == -1, f"Only -1 power is supported, got {p}"
102-
arg = _get_symbol_type(arg)
103-
if ceiling:
104-
return lhs / arg
105-
else:
106-
return lhs // arg
92+
_Rational = namedtuple("_Rational", ["numerator", "denominator"])
93+
94+
95+
def _get_symbol_type(expr):
96+
stack = []
97+
98+
def addi(lhs, rhs):
99+
return lhs + rhs
100+
101+
def muli(lhs, rhs):
102+
return lhs * rhs
103+
104+
def _add(lhs, rhs):
105+
is_rational_lhs = isinstance(lhs, _Rational)
106+
is_rational_rhs = isinstance(rhs, _Rational)
107+
if is_rational_lhs and not is_rational_rhs:
108+
numerator = muli(lhs.denominator, rhs)
109+
numerator = addi(numerator, lhs.numerator)
110+
return _Rational(numerator, lhs.denominator)
111+
elif not is_rational_lhs and is_rational_rhs:
112+
numerator = muli(lhs, rhs.denominator)
113+
numerator = addi(numerator, rhs.numerator)
114+
return _Rational(numerator, rhs.denominator)
115+
elif is_rational_lhs and is_rational_rhs:
116+
lhs_numerator = muli(lhs.numerator, rhs.denominator)
117+
rhs_numerator = muli(rhs.numerator, lhs.denominator)
118+
numerator = addi(lhs_numerator, rhs_numerator)
119+
denominator = muli(lhs.denominator, rhs.denominator)
120+
return _Rational(numerator, denominator)
121+
else:
122+
return addi(lhs, rhs)
123+
124+
# `x * (a/b)` transformed into `(x * a) / b`
125+
def _mul(lhs, rhs):
126+
is_rational_lhs = isinstance(lhs, _Rational)
127+
is_rational_rhs = isinstance(rhs, _Rational)
128+
if is_rational_lhs and not is_rational_rhs:
129+
numerator = muli(lhs.numerator, rhs)
130+
return _Rational(numerator, lhs.denominator)
131+
elif not is_rational_lhs and is_rational_rhs:
132+
numerator = muli(lhs, rhs.numerator)
133+
return _Rational(numerator, rhs.denominator)
134+
elif is_rational_lhs and is_rational_rhs:
135+
numerator = muli(lhs.numerator, rhs.numerator)
136+
denominator = muli(lhs.denominator, rhs.denominator)
137+
return _Rational(numerator, denominator)
138+
else:
139+
return muli(lhs, rhs)
140+
141+
def _floor(value):
142+
if isinstance(value, _Rational):
143+
value = value.numerator // value.denominator
144+
145+
return value
146+
147+
def _ceiling(value):
148+
if isinstance(value, _Rational):
149+
value = value.numerator / value.denominator
150+
151+
return value
152+
153+
def _group_rationals(stack, count):
154+
"""Group rationals and non-rationals args into 2 contiguous sets.
155+
156+
This allows to mul/add all non-rationals first, reducing total number of ops.
157+
"""
158+
rationals = []
159+
non_rationals = []
160+
for _ in range(count):
161+
val = stack.pop()
162+
if isinstance(val, _Rational):
163+
rationals.append(val)
107164
else:
108-
rhs = _get_symbol_type(rhs)
109-
return lhs * rhs
110-
111-
case sympy.floor():
112-
return _get_symbol_type(term.args[0])
113-
114-
case sympy.ceiling():
115-
return _get_symbol_type(term.args[0], ceiling=True)
116-
117-
assert False, f"Can't convert value {term} : {type(term)}"
165+
non_rationals.append(val)
166+
167+
return non_rationals + rationals
168+
169+
def _apply(args, func):
170+
assert len(args) > 0
171+
value = args[0]
172+
for val in args[1:]:
173+
value = func(value, val)
174+
175+
return value
176+
177+
def _enforce_non_rational(val, term):
178+
if isinstance(val, _Rational):
179+
raise ValueError(f"Rational is not supported yet in '{type(term)}'")
180+
181+
if not isinstance(expr, sympy.Expr):
182+
expr = sympy.sympify(expr)
183+
184+
for term in sympy.postorder_traversal(expr):
185+
match term:
186+
case sympy.Symbol():
187+
stack.append(typing.SymbolType.get(term.name))
188+
case sympy.Integer():
189+
stack.append(_get_literal(int(term)))
190+
case sympy.Mul():
191+
args = _group_rationals(stack, len(term.args))
192+
stack.append(_apply(args, _mul))
193+
case sympy.Add():
194+
args = _group_rationals(stack, len(term.args))
195+
stack.append(_apply(args, _add))
196+
case sympy.Mod():
197+
rhs = stack.pop()
198+
lhs = stack.pop()
199+
_enforce_non_rational(rhs, term)
200+
_enforce_non_rational(lhs, term)
201+
mod = lhs % rhs
202+
stack.append(mod)
203+
case sympy.Pow():
204+
assert term.args[1] == -1, f"Only -1 power is supported, got {p}"
205+
p = stack.pop()
206+
val = stack.pop()
207+
stack.append(_Rational(_get_literal(1), val))
208+
case sympy.floor():
209+
stack.append(_floor(stack.pop()))
210+
case sympy.ceiling():
211+
stack.append(_ceiling(stack.pop()))
212+
case sympy.Rational():
213+
numerator = _get_literal(term.p)
214+
denominator = _get_literal(term.q)
215+
stack.append(_Rational(numerator, denominator))
216+
case sympy.UnevaluatedExpr():
217+
continue
218+
case _:
219+
raise ValueError(f"Can not handle {type(term)} : {term}")
220+
221+
if len(stack) != 1 or isinstance(stack[0], _Rational):
222+
raise ValueError(f"Expected single result, got {len(stack)}: {stack}")
223+
224+
return stack[0]
118225

119226

120227
def _get_shape_attr(src):

‎python/hckernel/tests/test_e2e.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -44,10 +44,10 @@ def get_device(backend):
4444

4545
_copy_shapes = [
4646
(1, 27),
47-
(10, 20),
4847
(1, 72),
49-
(111, 813),
5048
(1, 128),
49+
(10, 20),
50+
(111, 813),
5151
(256, 64),
5252
(256, 128),
5353
(256, 256),

0 commit comments

Comments
 (0)