Conversation
| if opinfo.symbolic_parameter_list is None: | ||
| opinfo.symbolic_parameter_list = [ArgumentType.Symbolic] * len(args) | ||
| num_symbolic_parameters = len(opinfo.symbolic_parameter_list) | ||
| symbolic_parameter_list = ( |
There was a problem hiding this comment.
Avoid updating opinfo.symbolic_parameter_list directly since ops like linear can have variable number of parameters which throws an error for the linear input generator which can have 2 /3 arguments.
|
!build |
| testing::Values(Sizes({k}), Sizes({m, k}), Sizes({b, m, k})), | ||
| testing::Values(Sizes({n, k})), | ||
| testing::Values(Sizes({n})))); |
There was a problem hiding this comment.
Can we add broadcast dims here?
| shapes_input = ((K), (M, K), (B, M, K)) | ||
| shapes_weight = ((K), (N, K)) |
There was a problem hiding this comment.
Can we test with broadcasts? What about more than one batch dim?
jacobhinkle
left a comment
jacobhinkle
left a comment
There was a problem hiding this comment.
Looking good. I had a few minor questions but the only thing I think is missing is a test for IterDomain mapping like in #2246.
csrc/ops/composite.cpp
Outdated
| std::vector<IterDomain*> out_domain = ops::newOutputDomain({mapping_a, mapping_b}); | ||
| for (auto idx : c10::irange(ndims_out)) { | ||
| std::vector<IterDomain*> input_ids; | ||
| input_ids.reserve(2); |
There was a problem hiding this comment.
nit: if input_ids.size() is not always 2 then just skip reserving.
| // Check if the producer is A, B or bias. | ||
| std::optional<MatmulRole> input_role = std::nullopt; | ||
| if (producer->sameAs(op->inA()->as<TensorView>()->domain())) { | ||
| input_role = MatmulRole::INPUT_A; | ||
| } else if (producer->sameAs(op->inB()->as<TensorView>()->domain())) { | ||
| input_role = MatmulRole::INPUT_B; | ||
| } else if (producer->sameAs(op->bias()->as<TensorView>()->domain())){ | ||
| input_role = MatmulRole::INPUT_C; | ||
| } else { | ||
| NVF_ERROR(false, "Producer did not match any LinearOp input.") | ||
| } | ||
|
|
||
| // LinearOp: | ||
| // inputs (INPUT_A) = {*, in_features} | ||
| // weight (INPUT_B) = {out_features, in_features} / {in_features} | ||
| // bias (INPUT_C) = {out_features} / {} | ||
| // output = {*, out_features} / {*} | ||
|
|
||
| const std::vector<IterDomain*>& aligned_producer_ids = | ||
| ops::mapLinearOpIterDomains(producer_root, input_role.value(), out_size); | ||
|
|
||
| for (auto inx : c10::irange(out_size)) { | ||
| IterDomain* producer_id = aligned_producer_ids.at(inx); | ||
| IterDomain* consumer_id = consumer_root.at(inx); | ||
| if (producer_id == nullptr) { | ||
| continue; | ||
| } | ||
| updatePairwiseRootDomainMap(producer_id, consumer_id); | ||
| } | ||
|
|
||
| return dom_map; | ||
| } |
There was a problem hiding this comment.
Is it worth reusing some of this code between matmul and linear? I am not sure if anything but the last loop can be re-used but it seems like there might be a way since the patterns are so similar.
There was a problem hiding this comment.
I'll put the last loop in a separate fn atleast. Will see if combining the rest of the snippet is clean enough or not since we will still need to distinguish between matmul/linear
|
!build |
| Sizes({1, 1}), | ||
| Sizes({b, 1, n})))); | ||
|
|
||
| INSTANTIATE_TEST_SUITE_P( |
There was a problem hiding this comment.
I still need to think about a good naming scheme: Underscore separated shapes looked too verbose and for filtering we will need the exact shapes in that case.
jacobhinkle
left a comment
jacobhinkle
left a comment
There was a problem hiding this comment.
LGTM. Thanks for adding the ID mapping checks.
4 participants
LinearOpnode that has the same functionality asF.linear.linear_op_generatorto cover all cases allowed in thunder /F.linear.Issue #2149, #2132.