Skip to content

Conversation

@rolfmorel
Copy link
Contributor

Walks schedule IR, collects options (name -> values mapping), calls an autotune procedure (dummy implementation just selects first value) to select one value for each name, walks schedule IR again to introduce a param constant for the each option's value and replaces the users of the transform.tune.select op with the constant op's result.

Run it as tpp-sched | tpp-tune where tpp-sched produces IR containing transform.tune.select ops.

@rolfmorel
Copy link
Contributor Author

rolfmorel commented Jun 16, 2025

Segment of updated tpp-sched:

    ...
    %13 = transform.apply_registered_pass "pack-conv2DNhwcHwcf" to %12 : (!transform.any_op) -> !transform.any_op                                           
    %14 = transform.apply_registered_pass "rewrite-conv-to-matmul-or-brgemm" to %13 : (!transform.any_op) -> !transform.any_op                              
    %15 = transform.tune.select @m from [2, 4, 8] : !transform.any_param                                                                                    
    %16 = transform.tune.select @n from [4, 8, 16] : !transform.any_param                                                                                   
    %17 = transform.tune.select @k from [2, 4, 8, 16] : !transform.any_param                                                                                
    %18 = transform.apply_registered_pass "pack-matmul" with options = {"block-factors" = [%15, %16, %17]} to %14 : (!transform.any_op, !transform.any_param, !transform.any_param, !transform.any_param) -> !transform.any_op                                                                                       
    %19 = transform.apply_registered_pass "pack-vnni" to %18 : (!transform.any_op) -> !transform.any_op                                                     
    %20 = transform.structured.match ops{["func.func"]} in %10 : (!transform.any_op) -> !transform.any_op
    ...

Corresponding segment of tpp-sched | tpp-tune:

    ...
    %13 = transform.apply_registered_pass "pack-conv2DNhwcHwcf" to %12 : (!transform.any_op) -> !transform.any_op                                           
    %14 = transform.apply_registered_pass "rewrite-conv-to-matmul-or-brgemm" to %13 : (!transform.any_op) -> !transform.any_op                              
    %15 = transform.param.constant 2 : i64 -> !transform.any_param                                                                                          
    %16 = transform.tune.select @m from [2, 4, 8] : !transform.any_param                                                                                    
    %17 = transform.param.constant 4 : i64 -> !transform.any_param                                                                                          
    %18 = transform.tune.select @n from [4, 8, 16] : !transform.any_param                                                                                   
    %19 = transform.param.constant 2 : i64 -> !transform.any_param                                                                                          
    %20 = transform.tune.select @k from [2, 4, 8, 16] : !transform.any_param                                                                                
    %21 = transform.apply_registered_pass "pack-matmul" with options = {"block-factors" = [%15, %17, %19]} to %14 : (!transform.any_op, !transform.any_param, !transform.any_param, !transform.any_param) -> !transform.any_op                                                                                       
    %22 = transform.apply_registered_pass "pack-vnni" to %21 : (!transform.any_op) -> !transform.any_op                                                     
    %23 = transform.structured.match ops{["func.func"]} in %10 : (!transform.any_op) -> !transform.any_op
    ...

And the corresponding segment of tpp-sched | tpp-tune | tpp-opt -canonicalize:

   ...
    %13 = transform.apply_registered_pass "pack-conv2DNhwcHwcf" to %12 : (!transform.any_op) -> !transform.any_op                                           
    %14 = transform.apply_registered_pass "rewrite-conv-to-matmul-or-brgemm" to %13 : (!transform.any_op) -> !transform.any_op                              
    %15 = transform.param.constant 2 : i64 -> !transform.any_param                                                                                          
    %16 = transform.param.constant 4 : i64 -> !transform.any_param                                                                                          
    %17 = transform.param.constant 2 : i64 -> !transform.any_param                                                                                          
    %18 = transform.apply_registered_pass "pack-matmul" with options = {"block-factors" = [%15, %16, %17]} to %14 : (!transform.any_op, !transform.any_param, !transform.any_param, !transform.any_param) -> !transform.any_op                                                                                       
    %19 = transform.apply_registered_pass "pack-vnni" to %18 : (!transform.any_op) -> !transform.any_op
    %20 = transform.structured.match ops{["func.func"]} in %10 : (!transform.any_op) -> !transform.any_op
    ...

Walks schedule IR, collects options (name -> values mapping), calls an
`autotune` procedure (dummy implementation just selects first value) to
select one value for each name, walks schedule IR again to introduce
a param constant for the each option's value and replaces the users of
the `transform.tune.select` op with the constant op's result.

Run it as `tpp-sched | tpp-tune` where `tpp-sched` produces IR
containing `transform.tune.select` ops.
Comment on lines 32 to 34
def autotune(choices: Dict[str, Sequence[ir.Attribute]]) -> Dict[str, ir.Attribute]:
# Aint tuning easy!!
return {key: values[0] for key, values in choices.items()}
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Just need to replace this function's implementation to make it useful.

@rolfmorel rolfmorel marked this pull request as draft June 17, 2025 11:05
@rolfmorel
Copy link
Contributor Author

rolfmorel commented Jun 17, 2025

Latest update allows for specifying options via commandline, e.g. tpp-sched --pack-block-factors "2;4;8,4,8;16;32" | tpp-tune gives:

    %15 = transform.param.constant 8 : i64 -> !transform.any_param                                                                                          
    %16 = transform.tune.select @m from [2, 4, 8] : !transform.any_param                                                                                    
    %17 = transform.param.constant 4 : i64 -> !transform.any_param                                                                                          
    %18 = transform.tune.select @n from [4] : !transform.any_param                                                                                          
    %19 = transform.param.constant 16 : i64 -> !transform.any_param                                                                                         
    %20 = transform.tune.select @k from [8, 16, 32] : !transform.any_param                                                                                  
    %21 = transform.apply_registered_pass "pack-matmul" with options = {"block-factors" = [%15, %17, %19]} to %14 : (!transform.any_op, !transform.any_param, !transform.any_param, !transform.any_param) -> !transform.any_op

@rolfmorel
Copy link
Contributor Author

rolfmorel commented Jun 18, 2025

transform.tune.pick: tune.select but without forgetting the options:

    %15 = transform.param.constant 4 : i64 -> !transform.any_param
    %16 = transform.tune.select @m from [2, 4, 8] : !transform.any_param
    %17 = transform.param.constant 4 : i64 -> !transform.any_param
    %18 = transform.tune.select @n from [4] : !transform.any_param
    %19 = transform.tune.pick @k from [8, 16, 32] {selected = 8 : i64} : !transform.any_param
    %20 = transform.apply_registered_pass "pack-matmul" with options = {"block-factors" = [%15, %17, %19]} to %14 : (!transform.any_op, !transform.any_param, !transform.any_param, !transform.any_param) -> !transform.any_op

@rolfmorel
Copy link
Contributor Author

Closing due to feature getting merged upstream: llvm/llvm-project#146732

@rolfmorel rolfmorel closed this Jul 16, 2025
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants