[HostIr] Refactor hir optimization pass

[samnordmann]

What

This PR is only about cleaning and refactoring, no change in the behavior.

• We create a new base class for HIR lowering optimization pass hir_pass::OptimizationPass in host_ir/pass/optimization_pass.h.
• We create an option hir_lowering_logging to control debug logging through NVFUSER_DUMP
• We create a guard to enable/disable a pass
• We make the existing pass StreamLowering and InsertDeallocations inherit from this pass
• We factor out converting a resharding op to a communication into a separate pass

Why

preparation for #4387

[samnordmann]

!test

[github-actions]

Review updated until commit f76fc10

Description

• Created a new base class OptimizationPass for HIR lowering optimization passes.

• Introduced ConvertOpToCommunication pass to handle converting operations to communications.

• Moved communication lowering logic to lower_to_communication.cpp.

• Refactored HostIrLower to use convertSingleOpToCommunication.

---

Changes walkthrough 📝

Relevant files

Enhancement

19 files executor.cpp Updated to use `convertSingleOpToCommunication`.                  +4/-3      lower.cpp Removed communication lowering logic, updated to use convertSingleOpToCommunication. +4/-556  lower_to_communication.cpp Added new file with communication lowering logic.                +550/-0  convert_op_to_communication.cpp Added new pass for converting operations to communications. +75/-0    insert_deallocations.cpp Updated to use new `InsertDeallocations` class.                    +9/-5      stream_parallel_type.cpp Updated to use new `OptimizationPass` base class.                +3/-3      options.cpp Added new option `host_ir_lowering_logging`.                          +1/-0      fusion_kernel_runtime.cpp Updated to use new convertSingleOpToCommunication and InsertDeallocations class. +7/-4      fusion_definition.cpp Updated to use new `OptimizationPassGuard` with `StreamParallelType`. +1/-1      test_host_ir_stream_lowering.cpp Updated to use new `StreamParallelType` class.                      +15/-22  test_multidevice_host_ir.cpp Updated to use new OptimizationPassGuard with StreamParallelType and ReorderShardedAxisPass. +11/-5    lower.h Added declaration for `convertSingleOpToCommunication`.    +9/-1      lower_to_communication.h Added new header for communication lowering logic.              +20/-0    convert_op_to_communication.h Added new header for `ConvertOpToCommunication` pass.        +36/-0    insert_deallocations.h Updated to use new `InsertDeallocations` class.                    +12/-3    optimization_pass.h Added new base class `OptimizationPass`.                                  +88/-0    stream_parallel_type.h Updated to use new `OptimizationPass` base class.                +6/-7      options.h Added new option `host_ir_lowering_logging`.                          +1/-0      CMakeLists.txt Added new source files for communication lowering and passes. +2/-0

---

PR Reviewer Guide 🔍

Here are some key observations to aid the review process:

🧪 PR contains tests

⚡ Recommended focus areas for review Code Duplication The code for lowering expressions to communications was moved from HostIrLower::lower to convertSingleOpToCommunication in lower_to_communication.cpp. However, the original HostIrLower::lower method still contains a significant amount of duplicate code that should be removed to avoid redundancy. bool HostIrLower::canLower(Expr* expr, bool ignore_inner_resharding) { if (!isResharding(expr)) { return true; } if (!ir_utils::isTvOp(expr)) { return false; } if (auto* reduction = dynamic_cast<ReductionOp*>(expr)) { if (!ignore_inner_resharding && isInnerResharding(expr)) { return false; } auto in = reduction->in()->as<TensorView>(); auto out = reduction->out()->as<TensorView>(); // get the reduced axis std::vector<IterDomain*> reduction_axis; std::copy_if( out->getLogicalDomain().begin(), out->getLogicalDomain().end(), std::back_inserter(reduction_axis), [](IterDomain* id) { return id->isReduction(); }); // check whether the reduction involves only one axis if (reduction_axis.size() != 1) { return false; } // We check whether the reduced axis is sharded on the input const auto c2p_map = PairwiseLogicalDomainMap(in, out).mapConsumerToProducer(); auto c2p_map_it = c2p_map.find(reduction_axis.at(0)); return c2p_map_it != c2p_map.end() && c2p_map_it->second->isDeviceDim(); } else if (auto* ldst = dynamic_cast<LoadStoreOp*>(expr)) { if (!ignore_inner_resharding && isInnerResharding(expr)) { return false; } return ldst->as<LoadStoreOp>()->opType() == LoadStoreOpType::Set; } else if (auto* matmul = dynamic_cast<MatmulOp*>(expr)) { // For now we only support out = matmul(a,b) when b, out are fully // replicated, a is sharded on axis 1, and out i stream-parallelized on axis // 0. return !isSharded(matmul->inB()) && !isSharded(matmul->out()) && matmul->inA()->axis(0)->getParallelType() == ParallelType::Serial && getShardedLogicalAxis(matmul->inA(), ParallelType::DIDx) == 1 && matmul->out()->axis(0)->getParallelType() == ParallelType::Stream; } else if (auto* linear = dynamic_cast<LinearOp*>(expr)) { // For now we only support out = linear(a, b, bias) when b, bias, and out // are fully replicated, a is sharded on axis 1, and out i // stream-parallelized on axis 0. auto* a = linear->inA()->as<TensorView>(); auto* b = linear->inB()->as<TensorView>(); auto* bias = (linear->has_bias() ? linear->bias()->as<TensorView>() : nullptr); auto* out = linear->out()->as<TensorView>(); return !isSharded(b) && !(linear->has_bias() && isSharded(bias)) && !isSharded(out) && a->axis(0)->getParallelType() == ParallelType::Serial && getShardedLogicalAxis(a, ParallelType::DIDx) == 1 && out->axis(0)->getParallelType() == ParallelType::Stream; } return false; } bool HostIrLower::isLowerableAsStandaloneHostOp(Expr* expr) { if (expr->isOneOf< MatmulOp, SliceOp, SelectOp, LinearOp, BinaryOp, ReductionOp, Communication, P2PCommunication>()) { return true; } // Lower as standalone op "set" ops, i.e., LoadStoreOp of "Set" type with no // permute if (expr->isA<LoadStoreOp>()) { auto* load_store = expr->as<LoadStoreOp>(); if (load_store->opType() == LoadStoreOpType::Set && load_store->out()->isA<TensorView>()) { auto* tv = load_store->out()->as<TensorView>(); // If the output tensor has no root, it means it has no permute if (!tv->hasRoot()) { return true; } } } return false; } bool HostIrLower::shouldMergeSegmentedGroups( SegmentedGroup* group1, SegmentedGroup* group2) { for (auto group : {group1, group2}) { for (Expr* expr : group->exprs()) { if (isLowerableAsStandaloneHostOp(expr)) { return false; } } } return true; } std::unique_ptr<hir::HostIrContainer> HostIrLower::lower( std::unique_ptr<Fusion> fusion, DeviceIdxType my_device_index) { // Sharding PreSegmenter passes. // Note: passes run before PreSegmenter optimization passes. preseg_passes::OptimizationPass< preseg_passes::PropagateShardingsPass>::runPass(fusion.get()); preseg_passes::OptimizationPass< preseg_passes::InsertReshardingsPass>::runPass(fusion.get()); preseg_passes::OptimizationPass< preseg_passes::ReorderShardedAxisPass>::runPass(fusion.get()); preseg_passes::OptimizationPass< preseg_passes::MakeReshardingContiguousPass>::runPass(fusion.get()); // Performs segmentation at the inter-device communications // Each SegmentedGroup represents a pipeline's stage, and can be either // 1) a Fusion which doesn't involve inter-device communication // 2) a Fusion comprised of one Expr, representing inter-device communication SegmentCandidateFinderOptions options{ .run_translate_welford = false, .run_combine_reductions = false, .run_herrmann_merge = true, .run_final_merge = true, .custom_should_merge_groups = &shouldMergeSegmentedGroups}; std::unique_ptr<SegmentedFusion> staged_fusion = SegmentCandidateFinder::segment( std::move(fusion), KernelArgumentHolder(), options, true); // Infer a topologically ordered traversal of the segmented fusion to // determine the order for launching the kernels/comms RuntimeWorkSpace workspace; prepareRuntimeOrder(staged_fusion.get(), workspace); // Create the HostIrContainer representing the host program. Each segment of // the segmented fusion will be translated to a HostIR auto hic = std::make_unique<hir::HostIrContainer>(); FusionGuard fg(hic.get()); IrCloner ir_cloner(hic.get()); auto clone = [&ir_cloner](const std::vector<Val*>& vals) -> std::vector<Val*> { std::vector<Val*> cloned_vals(vals.size()); std::transform( vals.begin(), vals.end(), cloned_vals.begin(), [&ir_cloner](Val* val) { return ir_cloner.clone(val); }); return cloned_vals; }; for (auto group : workspace.group_run_order) { NVF_ERROR(!group->exprs().empty(), "invalid segmentation"); if (involvedDevices(group->exprs().at(0)).count(my_device_index) == 0) { continue; } // we decide whether to insert the Expr as a standalone op in the // HostIRContainer, which will result in using ATen Op to evaluate it -- // or, alternatively, to wrap them into a PostOnStream(HostUnit(.)) which // will result in a kernel code generation. if (std::all_of( group->exprs().begin(), group->exprs().end(), isLowerableAsStandaloneHostOp)) { NVF_ERROR( group->exprs().size() == 1, "Expr executed as a standalone op cannot be fused"); hic->pushBackTopLevelExprs(ir_cloner.clone(group->exprs().at(0))); } else { auto host_unit = IrBuilder::create<hir::HostUnit>( staged_fusion->makeFusion(group).second); auto post_on_stream = IrBuilder::create<hir::PostOnStream>( host_unit, clone(group->inputs()), clone(group->outputs())); hic->pushBackTopLevelExprs(post_on_stream); } } for (auto input : staged_fusion->inputs()) { hic->addInput(ir_cloner.clone(input)); } for (auto output : staged_fusion->outputs()) { hic->addOutput(ir_cloner.clone(output)); } for (auto tv : hic->allTvs()) { // set all host tensors to global memory type. This must be the case by // definition of a host tensor, and setting the memory type to global is // also required to avoid Allocate HIR nodes to throw tv->setMemoryType(MemoryType::Global); } hir_pass::StreamParallelType().runPass(hic.get()); hir_pass::ConvertOpToCommunication(params_).runPass(hic.get()); return hic; } } // namespace nvfuser Error Handling The convertSingleOpToCommunication function does not handle errors in the same way as the original HostIrLower::lower method. Ensure that all error conditions are properly checked and handled to maintain consistency. Code Duplication The InsertDeallocations class now includes a check for hir::Deallocate that was previously in the insertDeallocations function. This check should be reviewed to ensure it is not redundant and that the function behaves as expected. !expr->isA<hir::Deallocate>(), "Expected hostir container to not have deallocate, but found one anyways"); }); std::unordered_map<TensorView*, int64_t> last_use; for (auto&& [i, expr] : enumerate(top_level_exprs)) { for (auto* val : expr->inputs()) { if (!val->isA<TensorView>()) { continue; } auto tv = val->as<TensorView>(); last_use[tv] = i; } } std::vector<std::pair<int64_t, TensorView*>> last_use_by_index; last_use_by_index.reserve(last_use.size()); for (auto&& [tv, i] : last_use) { last_use_by_index.emplace_back(i, tv); } std::sort(last_use_by_index.begin(), last_use_by_index.end()); for (auto&& [i, tv] : last_use_by_index | std::views::reverse) { auto* deallocate = IrBuilder::create<hir::Deallocate>(tv); hic->insertExprAfter(i, deallocate); } } } // namespace nvfuser::hir_pass

[samnordmann]

!test

[samnordmann]

!test

[wujingyue]

For the main stack, implemented in fusion_kernel_runtime.cpp, I consider this to be part of the fusion-IR-to-host-IR lowering (aka host IR lowering) instead of a host-IR-to-host-IR transformation pass (aka host IR pass).

Therefore, I prefer leaving this function in lower.h so it can be used in both host IR lowering and host IR passes.

[samnordmann]

For the main stack, implemented in fusion_kernel_runtime.cpp, I consider this to be part of the fusion-IR-to-host-IR lowering (aka host IR lowering) instead of a host-IR-to-host-IR transformation pass (aka host IR pass)

Let me motivate the reason why I moved it:

1. I personally I find it cleaner and more logical to keep this function separated, and close to the pass that uses it.
2. The main stack should eventually use the pass and not the function. The function should eventually not even be exposed
3. it is going to be a host-IR-to-host-IR pass after stream lowering will be integrated. To see it as part of lowering and not host IR pass is just because of the current implementation but is not necessary.

Therefore, I prefer leaving this function in lower.h so it can be used in both host IR lowering and host IR passes.

Would you be ok at least to put it a separate file host_ir/lower_to_communication.h|cpp? (as was named in the past)

Otherwise it would means that HostIrLower::lower from host_ir/lower.h calls the pass host_ir/pass/convert_op_to_communication.h which in turns calls HostIrLower:: ConvertSingleOpToCommunication from host_ir/lower.h, which is kind of cyclical and therefore misleading.

[wujingyue]

Would you be ok at least to put it a separate file host_ir/lower_to_communication.h|cpp? (as was named in the past)

SGTM!

[wujingyue]

This defeats the purpose of "recurring template pattern". I realized

Fuser/csrc/preseg_passes/optimization_pass.h

Line 72 in d0ecedd

virtual ~OptimizationPass() = default;

might have misled you and is probably (I'll give that a quick shot) unnecessary. See

Fuser/csrc/preseg_passes/optimization_pass.h

Line 56 in d0ecedd

DerivedClass::runPass(fusion);

for an example of how to avoid virtual methods.

[samnordmann]

This defeats the purpose of "recurring template pattern"

you're right, thanks. My problem is that I do not know how to deal with that ConvertOpToCommunication pass (and others in the future) has a non-trivial constructor. Maybe I am just missing something.

what about not using recurring template pattern but and keep the runPass in the base class and passImplementation as a private member of the derived class ?

Let me know what you suggest

[wujingyue]

Would the following work for non-trivial constructors?

template <typename Derived>
class Base {
public:
    void do_something() {
        static_cast<Derived*>(this)->implementation();
    }
};

class Derived : public Base<Derived> {
public:
    Derived(int x, std::string name) : x_(x), name_(std::move(name)) {}

    void implementation() {
        std::cout << "x: " << x_ << ", name: " << name_ << "\n";
    }

private:
    int x_;
    std::string name_;
};

(Again, it uses no virtual methods here)

Alternatively, if you don't plan to use CRTP for your host IR passes, please build hir::OptimizationPass as a non-template class and simplify. If you choose to do this, keep insertDeallocations as a helper function and call it from your hir::OptimizationPass. This way, I can wrap it differently if/when I need to.

[samnordmann]

Thanks! I was not so familiar with CRTP, now it is clearer. I have removed the virtual method. Let me know.

[wujingyue]

LGTM otherwise!

[samnordmann]

!test

[samnordmann]

!test

[samnordmann]

!test