Skip to content

Commit 24161f3

Browse files
Hardcode84Hardcode84
authored
Typing work (#14)
1 parent 59a926e commit 24161f3

7 files changed

Lines changed: 130 additions & 20 deletions

File tree

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

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -102,6 +102,15 @@ def Typing_SequenceType : Typing_Type<"Sequence", "sequence", [], "::mlir::Type"
102102
let assemblyFormat = "$params";
103103
}
104104

105+
def Typing_UnpackedSequenceType : Typing_Type<"UnpackedSequence", "unpacked_sequence", [], "::mlir::Type"> {
106+
let parameters = (ins
107+
ArrayRefParameter<"::mlir::Type">:$params
108+
);
109+
110+
let assemblyFormat = "$params";
111+
}
112+
113+
105114
def Typing_SymbolType : Typing_Type<"Symbol", "symbol", [], "::mlir::Type"> {
106115
let parameters = (ins "::mlir::StringAttr":$name);
107116

@@ -409,6 +418,16 @@ def Typing_GetSeqSizeOp : Typing_OpBase<"get_seq_size", [
409418
let assemblyFormat = "attr-dict $seq";
410419
}
411420

421+
def Typing_UnpackSeqOp : Typing_OpBase<"unpack_seq", [
422+
Pure,
423+
DeclareOpInterfaceMethods<TypingInterpreterInterface>
424+
]> {
425+
let arguments = (ins Typing_ValueType:$seq);
426+
let results = (outs Typing_ValueType:$result);
427+
428+
let assemblyFormat = "attr-dict $seq";
429+
}
430+
412431
def Typing_IsSameOp : Typing_OpBase<"is_same", [
413432
Pure, DeclareOpInterfaceMethods<TypingInterpreterInterface>
414433
]> {
@@ -428,6 +447,14 @@ def Typing_CheckOp : Typing_OpBase<"check", [
428447
let assemblyFormat = "attr-dict $condition";
429448
}
430449

450+
def Typing_PrintOp : Typing_OpBase<"print", [
451+
DeclareOpInterfaceMethods<TypingInterpreterInterface>
452+
]> {
453+
let arguments = (ins Typing_ValueType:$value);
454+
455+
let assemblyFormat = "attr-dict $value";
456+
}
457+
431458
def Typing_MakeUnionOp : Typing_OpBase<"make_union", [
432459
Pure, DeclareOpInterfaceMethods<TypingInterpreterInterface>
433460
]> {

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

Lines changed: 21 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -563,11 +563,11 @@ hc::typing::TypeResolverReturnOp::interpret(InterpreterState &state) {
563563
auto &result = *state.result;
564564
mlir::ValueRange args = getArgs();
565565

566-
auto isSeq = [&]() -> typing::SequenceType {
566+
auto isSeq = [&]() -> typing::UnpackedSequenceType {
567567
if (args.size() != 1)
568568
return {};
569569

570-
return mlir::dyn_cast_if_present<typing::SequenceType>(
570+
return mlir::dyn_cast_if_present<typing::UnpackedSequenceType>(
571571
getType(state, args.front()));
572572
};
573573

@@ -773,6 +773,18 @@ hc::typing::GetSeqSizeOp::interpret(InterpreterState &state) {
773773
return InterpreterResult::Advance;
774774
}
775775

776+
mlir::FailureOr<hc::typing::InterpreterResult>
777+
hc::typing::UnpackSeqOp::interpret(InterpreterState &state) {
778+
auto seq = mlir::dyn_cast_if_present<SequenceType>(
779+
::hc::typing::getType(state, getSeq()));
780+
if (!seq)
781+
return emitError("Invalid seq type");
782+
783+
state.state[getResult()] =
784+
UnpackedSequenceType::get(getContext(), seq.getParams());
785+
return InterpreterResult::Advance;
786+
}
787+
776788
mlir::FailureOr<hc::typing::InterpreterResult>
777789
hc::typing::IsSameOp::interpret(InterpreterState &state) {
778790
auto lhs = hc::typing::getType(state, getLhs());
@@ -789,6 +801,13 @@ hc::typing::CheckOp::interpret(InterpreterState &state) {
789801
return *val ? InterpreterResult::Advance : InterpreterResult::MatchFail;
790802
}
791803

804+
mlir::FailureOr<hc::typing::InterpreterResult>
805+
hc::typing::PrintOp::interpret(InterpreterState &state) {
806+
auto type = hc::typing::getType(state, getValue());
807+
llvm::errs() << type << "\n";
808+
return InterpreterResult::Advance;
809+
}
810+
792811
mlir::FailureOr<hc::typing::InterpreterResult>
793812
hc::typing::MakeUnionOp::interpret(InterpreterState &state) {
794813
llvm::SmallSetVector<mlir::Type, 8> types;

‎mlir/lib/Transforms/PyTypeInference.cpp‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -375,7 +375,9 @@ class TypeAnalysis
375375
if (!*res)
376376
return mlir::success();
377377

378-
assert(result.size() == results.size());
378+
if (result.size() != results.size())
379+
return op->emitOpError("Invalid op results count");
380+
379381
for (auto &&[resultLattice, res] : llvm::zip_equal(results, result)) {
380382
auto changed = resultLattice->join(TypeValue(&interpreter, res));
381383
propagateIfChanged(resultLattice, changed);

‎python/cpp/src/TypingPipeline.cpp‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -94,6 +94,7 @@ static mlir::LogicalResult convertCall(mlir::PatternRewriter &builder,
9494
handlers[] = {
9595
{"is_same", i1, {vt, vt}, &convertCall<hc::typing::IsSameOp, 2>},
9696
{"check", none, {i1}, &convertCall<hc::typing::CheckOp, 1, true>},
97+
{"print_val", none, {vt}, &convertCall<hc::typing::PrintOp, 1, true>},
9798
{"make_symbol", vt, {vt}, &convertCall<hc::typing::MakeSymbolOp, 1>},
9899
{"get_num_args",
99100
index,
@@ -113,6 +114,7 @@ static mlir::LogicalResult convertCall(mlir::PatternRewriter &builder,
113114
index,
114115
{vt},
115116
&convertCall<hc::typing::GetSeqSizeOp, 1>},
117+
{"unpack_seq", vt, {vt}, &convertCall<hc::typing::UnpackSeqOp, 1>},
116118
{"get_type_name",
117119
vt,
118120
{vt},

‎python/hckernel/kernel_typing.py‎

Lines changed: 18 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -135,7 +135,7 @@ def resolver(a: ValueType):
135135
@type_resolver(_registry, ["py_ir.tuple_unpack"])
136136
def resolver(a: ValueType):
137137
check_is_tuple(a)
138-
return get_type_param(a, "elements")
138+
return unpack_seq(get_type_param(a, "elements"))
139139

140140

141141
@type_resolver(_registry, ["py_ir.getitem"])
@@ -220,18 +220,27 @@ def getitem_typing(dims: ValueType, index: ValueType):
220220
res_dims = create_seq()
221221
indices_len = get_seq_size(indices)
222222
i = 0
223-
drop_count = 0
223+
dim_idx = 0
224224
while i < indices_len:
225-
if is_same(get_seq_element(indices, i), NoneTyp):
225+
elem = get_seq_element(indices, i)
226+
if is_same(elem, NoneTyp):
226227
res_dims = append_seq(res_dims, 1)
228+
elif is_same(get_type_name(elem), "Slice"):
229+
lower = get_type_param(elem, "lower")
230+
# TODO: check for literals
231+
if is_same(lower, NoneTyp):
232+
lower = 0
233+
234+
upper = get_type_param(elem, "upper")
235+
if is_same(upper, NoneTyp):
236+
upper = get_seq_element(dims, dim_idx)
237+
238+
# size = upper - lower TODO
239+
res_dims = append_seq(res_dims, upper)
240+
dim_idx += 1
227241
else:
228-
drop_count += 1
229-
i += 1
242+
dim_idx += 1
230243

231-
dims_count = get_seq_size(dims)
232-
i = drop_count
233-
while i < dims_count:
234-
res_dims = append_seq(res_dims, get_seq_element(dims, i))
235244
i += 1
236245

237246
return res_dims
@@ -264,14 +273,6 @@ def resolver(target: ValueType):
264273
return make_type("Tuple", elements=dims)
265274

266275

267-
@type_resolver(_registry, ["py_ir.getitem"])
268-
def resolver(target: ValueType, index: ValueType):
269-
check_is_tensor(target)
270-
return make_type(
271-
"Tensor", dims=getitem_typing(get_type_param(target, "dims"), index)
272-
)
273-
274-
275276
@type_resolver(_registry, ["py_ir.binop"])
276277
def resolver(lhs: ValueType, rhs: ValueType):
277278
# TODO: broadcasting

‎python/hckernel/typing.py‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -103,6 +103,11 @@ def check(cond):
103103
_stub_error()
104104

105105

106+
@_register_func
107+
def print_val(value):
108+
_stub_error()
109+
110+
106111
@_register_func
107112
def get_attr(name):
108113
_stub_error()
@@ -148,6 +153,11 @@ def get_seq_size(seq):
148153
_stub_error()
149154

150155

156+
@_register_func
157+
def unpack_seq(seq):
158+
_stub_error()
159+
160+
151161
@_register_func
152162
def make_type(name, **kwargs):
153163
_stub_error()

‎python/ir/typing.mlir‎

Lines changed: 49 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -119,6 +119,31 @@ typing.type_resolver ["py_ir.call"] {
119119
typing.type_resolver_return %3
120120
}
121121

122+
// typing.print
123+
124+
typing.type_resolver ["py_ir.getattr", "print_val"] {
125+
%c0 = arith.constant 0: index
126+
%0 = typing.make_ident "hckernel.typing" []
127+
%1 = typing.get_arg %c0
128+
%2 = typing.is_same %0 %1
129+
typing.check %2
130+
131+
%3 = typing.make_ident "hckernel.typing.print_val" []
132+
typing.type_resolver_return %3
133+
}
134+
135+
typing.type_resolver ["py_ir.call"] {
136+
%c0 = arith.constant 0: index
137+
%0 = typing.make_ident "hckernel.typing.print_val" []
138+
%1 = typing.get_arg %c0
139+
%2 = typing.is_same %0 %1
140+
typing.check %2
141+
142+
%3 = typing.type_constant #typing.type_attr<none> : !typing.value
143+
typing.type_resolver_return %3
144+
}
145+
146+
122147
// typing.get_attr
123148

124149
typing.type_resolver ["py_ir.getattr", "get_attr"] {
@@ -340,6 +365,30 @@ typing.type_resolver ["py_ir.call"] {
340365
typing.type_resolver_return %3
341366
}
342367

368+
// typing.unpack_seq
369+
370+
typing.type_resolver ["py_ir.getattr", "unpack_seq"] {
371+
%c0 = arith.constant 0: index
372+
%0 = typing.make_ident "hckernel.typing" []
373+
%1 = typing.get_arg %c0
374+
%2 = typing.is_same %0 %1
375+
typing.check %2
376+
377+
%3 = typing.make_ident "hckernel.typing.unpack_seq" []
378+
typing.type_resolver_return %3
379+
}
380+
381+
typing.type_resolver ["py_ir.call"] {
382+
%c0 = arith.constant 0: index
383+
%0 = typing.make_ident "hckernel.typing.unpack_seq" []
384+
%1 = typing.get_arg %c0
385+
%2 = typing.is_same %0 %1
386+
typing.check %2
387+
388+
%3 = typing.type_constant #typing.type_attr<!typing.value> : !typing.value
389+
typing.type_resolver_return %3
390+
}
391+
343392
// typing.make_ident
344393

345394
typing.type_resolver ["py_ir.getattr", "make_type"] {

0 commit comments

Comments
 (0)