Skip to content

Commit

Permalink
[Lang] Migrate irpass::scalarize() after irpass::lower_access() (#8091)
Browse files Browse the repository at this point in the history
Issue: #

### Brief Summary

<!--
copilot:summary
-->
### <samp>🤖 Generated by Copilot at d754141</samp>

Refactored and improved the `scalarize` IR pass and applied it
selectively to block-local storage (BLS) kernels. This enhances the
performance and correctness of the BLS transformation and simplifies the
compilation pipeline.

### Walkthrough

<!--
copilot:walkthrough
-->
### <samp>🤖 Generated by Copilot at d754141</samp>

* Change the `scalarize` function to return a `bool` value indicating if
the IR was modified
([link](https://github.com/taichi-dev/taichi/pull/8091/files?diff=unified&w=0#diff-448ac6e85e192a27e5ec7c54cd8a91545dc7c83f62d030eafb9c190383cfe934L33-R33),
[link](https://github.com/taichi-dev/taichi/pull/8091/files?diff=unified&w=0#diff-97b0d9ab204b703802b3b5d04d036d30f66b34b726128216faf0d8a2a8564528L1184-R1205))
* Move the `scalarize` function to a later stage in the compilation
pipeline, after the offloading is done, and only apply it to offloaded
kernels with BLS enabled
([link](https://github.com/taichi-dev/taichi/pull/8091/files?diff=unified&w=0#diff-8fde186587db97b3bbc8a856e59bc4467b30257335b0fad064b4eebd521a912bL229-L236),
[link](https://github.com/taichi-dev/taichi/pull/8091/files?diff=unified&w=0#diff-8fde186587db97b3bbc8a856e59bc4467b30257335b0fad064b4eebd521a912bR271-R278))
* Check the return value of the `scalarize` function and call
`full_simplify` or `die` only if the IR was modified
([link](https://github.com/taichi-dev/taichi/pull/8091/files?diff=unified&w=0#diff-8fde186587db97b3bbc8a856e59bc4467b30257335b0fad064b4eebd521a912bR271-R278),
[link](https://github.com/taichi-dev/taichi/pull/8091/files?diff=unified&w=0#diff-8fde186587db97b3bbc8a856e59bc4467b30257335b0fad064b4eebd521a912bL359-R366))
* Add a helper function `make_pass_printer` to print the IR after each
pass if the `verbose` flag is true
([link](https://github.com/taichi-dev/taichi/pull/8091/files?diff=unified&w=0#diff-b2368908e6bad68906f40866b6dece421190dfd428d63788f2f5143b14785a45R12-R29))
* Call `scalarize` and `full_simplify` in the `make_block_local_offload`
function to handle the case where the global pointer index is a tensor
type stored in an alloca statement
([link](https://github.com/taichi-dev/taichi/pull/8091/files?diff=unified&w=0#diff-b2368908e6bad68906f40866b6dece421190dfd428d63788f2f5143b14785a45R36-R69))
* Use the `delayed_modifier_` as a return value of the `run` function
for each IR pass class in `scalarize.cpp`, and avoid calling the
modifier twice
([link](https://github.com/taichi-dev/taichi/pull/8091/files?diff=unified&w=0#diff-97b0d9ab204b703802b3b5d04d036d30f66b34b726128216faf0d8a2a8564528L21-L23),
[link](https://github.com/taichi-dev/taichi/pull/8091/files?diff=unified&w=0#diff-97b0d9ab204b703802b3b5d04d036d30f66b34b726128216faf0d8a2a8564528R841-R846),
[link](https://github.com/taichi-dev/taichi/pull/8091/files?diff=unified&w=0#diff-97b0d9ab204b703802b3b5d04d036d30f66b34b726128216faf0d8a2a8564528L901-L903),
[link](https://github.com/taichi-dev/taichi/pull/8091/files?diff=unified&w=0#diff-97b0d9ab204b703802b3b5d04d036d30f66b34b726128216faf0d8a2a8564528R1044-R1050),
[link](https://github.com/taichi-dev/taichi/pull/8091/files?diff=unified&w=0#diff-97b0d9ab204b703802b3b5d04d036d30f66b34b726128216faf0d8a2a8564528L1089-L1090),
[link](https://github.com/taichi-dev/taichi/pull/8091/files?diff=unified&w=0#diff-97b0d9ab204b703802b3b5d04d036d30f66b34b726128216faf0d8a2a8564528R1132-R1137))
* Change the `run` function of the `MergeExternalAndMatrixPtr` class to
return a `bool` value and make it consistent with other IR passes
([link](https://github.com/taichi-dev/taichi/pull/8091/files?diff=unified&w=0#diff-97b0d9ab204b703802b3b5d04d036d30f66b34b726128216faf0d8a2a8564528L1175-R1189))
  • Loading branch information
jim19930609 authored May 30, 2023
1 parent 1b4d152 commit afbc85b
Showing 1 changed file with 9 additions and 8 deletions.
17 changes: 9 additions & 8 deletions taichi/transforms/compile_to_offloads.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -231,14 +231,6 @@ void offload_to_executable(IRNode *ir,
print("Make block local");
}

if (config.real_matrix_scalarize) {
if (irpass::scalarize(ir)) {
// Remove redundant MatrixInitStmt inserted during scalarization
irpass::full_simplify(ir, config, {false, /*autodiff_enabled*/ false});
print("Scalarized");
}
}

if (is_extension_supported(config.arch, Extension::mesh)) {
irpass::demote_mesh_statements(ir, config, {kernel->get_name()});
print("Demote mesh statements");
Expand Down Expand Up @@ -276,6 +268,15 @@ void offload_to_executable(IRNode *ir,
irpass::analysis::verify(ir);
}

if (config.real_matrix_scalarize) {
if (irpass::scalarize(ir)) {
// Remove redundant MatrixInitStmt inserted during scalarization
irpass::full_simplify(ir, config,
{lower_global_access, /*autodiff_enabled*/ false});
print("Scalarized");
}
}

irpass::demote_operations(ir, config);
print("Operations demoted");

Expand Down

0 comments on commit afbc85b

Please sign in to comment.