[Host irs] Stream lowering of single device fusions

[samnordmann]

This PR belongs to a series of stacked PRs:

1. [Host irs] alias and preallocated output support #4144
2. [Host Ir] refactor and cleanup lowering and segmentation #4145
3. [Host ir] support for set reduce and binary op #4146
4. add HirAliasSelect #4301
5. => You are here: [Host irs] Stream lowering of single device fusions #4147

What

Implement a proper lowering for handling ParallelType::Stream. This PR has the following restrictions:

• Single device fusion
• No split/merge of Stream axis

We add to Hir lowering a new pass that reads the hir container's top level expressions, reads the consumer's stream parallelization and create For Loop with stream management and sync for expressing the stream parallelization. Basic logic for merging For-Loop are written.

Let me explain through some examples that can be found in the PR. We suggest to run those examples as follows:

NVFUSER_DUMP=host_ir test_host_ir --gtest_filter=*

Single expr and for-loop

Look at MultiDeviceExecutorLowerStreamTest.SingleSetOp simple scenario:

  TensorView* tv0 = makeContigTensor(2);
  TensorView* tv1 = set(tv0);
  fusion->addInput(tv0);
  fusion->addOutput(tv1);
  tv1->axis(0)->parallelize(ParallelType::Stream);

the dumped generated Host Ir program is:

%HostIrContainer { (T0_g_float[iS0{i0}, iS1{i2}]) -> (T1_g_float[iStreamIdx2{i0}, iS3{i2}]) :
  T1_g_float[iStreamIdx2{i0}, iS3{i2}] = ALLOCATE(buffer=T1_g_float[iStreamIdx2{i0}, iS3{i2}], mem_type=global, size=( i0 * i2 ), zero_init=false, resets_to_zero=false)
  FOR StreamIdx in iStreamIdx2{i0}:
    GetCurrentStream into Stream 0
    SetCurrentStream to Stream ( StreamIdx % numberOfStreams )
    Synchronize Stream 0
    T2_l_float[iS4{i2}]
       = HirAliasSelect( T0_g_float[iS0{i0}, iS1{i2}], axis = iS0{i0}, index = StreamIdx )
    T3_l_float[iS5{i2}]
       = HirAliasSelect( T1_g_float[iStreamIdx2{i0}, iS3{i2}], axis = iStreamIdx2{i0}, index = StreamIdx )
    T3_l_float[iS5{i2}]
       = Set( T2_l_float[iS4{i2}], cache_op=Streaming )
    SetCurrentStream to Stream 0
    Synchronize Stream ( StreamIdx % numberOfStreams )
} // %HostIrContainer

We can see that the expr, here the "Set", gets embedded into a For Loop. Let us analyze further:

• outside the for loop, we allocate the global output buffer.
• The start of the for loop body does the new stream assignment and sync of that stream to the user stream
• Then, we "Select" (aka slice) through HirAliasSelect into the input and output
• The "Set" operation is executed on the "selected" I/O. Note that the output is an alias to the output's slice.
• At the end of the for loop, we reset to the user's stream (I mean, the currently selected stream before entering the program) and sync the user's stream with the running stream.

Merging for loops

To avoid unnecessary synchronization across streams, it is important to be able to fuse the stream for-loop. This is exercised by the test MultiDeviceExecutorLowerStreamTest.TwoSetOps:

  TensorView* tv0 = makeContigTensor(2);
  TensorView* tv1 = set(tv0);
  TensorView* tv2 = set(tv1);
  fusion->addInput(tv0);
  fusion->addOutput(tv2);
  tv1->axis(0)->parallelize(ParallelType::Stream);
  tv2->axis(0)->parallelize(ParallelType::Stream);

dump:

%HostIrContainer { (T0_g_float[iS0{i0}, iS1{i2}]) -> (T2_g_float[iStreamIdx4{i0}, iS5{i2}]) :
  T1_g_float[iStreamIdx2{i0}, iS3{i2}] = ALLOCATE(buffer=T1_g_float[iStreamIdx2{i0}, iS3{i2}], mem_type=global, size=( i0 * i2 ), zero_init=false, resets_to_zero=false)
  T2_g_float[iStreamIdx4{i0}, iS5{i2}] = ALLOCATE(buffer=T2_g_float[iStreamIdx4{i0}, iS5{i2}], mem_type=global, size=( i0 * i2 ), zero_init=false, resets_to_zero=false)
  FOR StreamIdx in iStreamIdx2{i0}:
    GetCurrentStream into Stream 0
    SetCurrentStream to Stream ( StreamIdx % numberOfStreams )
    Synchronize Stream 0
    T3_l_float[iS6{i2}]
       = HirAliasSelect( T0_g_float[iS0{i0}, iS1{i2}], axis = iS0{i0}, index = StreamIdx )
    T4_l_float[iS7{i2}]
       = HirAliasSelect( T1_g_float[iStreamIdx2{i0}, iS3{i2}], axis = iStreamIdx2{i0}, index = StreamIdx )
    T4_l_float[iS7{i2}]
       = Set( T3_l_float[iS6{i2}], cache_op=Streaming )
    T5_l_float[iS8{i2}]
       = HirAliasSelect( T2_g_float[iStreamIdx4{i0}, iS5{i2}], axis = iStreamIdx4{i0}, index = StreamIdx )
    T5_l_float[iS8{i2}]
       = Set( T4_l_float[iS7{i2}], cache_op=Streaming )
    SetCurrentStream to Stream 0
    Synchronize Stream ( StreamIdx % numberOfStreams )
} // %HostIrContainer

We observe that the For-loop are indeed merged.
Possible future optimization: the allocation of the intermediate buffer could be only of length numberOfStreams

separating for loops

We also need to be able to separate and create new for loops if necessary, as exercised in ThreeSetOpsWithDisjointsForLoops, which considers the Fusion:

  TensorView* tv0 = makeContigTensor(2);
  TensorView* tv1 = set(tv0);
  TensorView* tv2 = set(tv1);
  TensorView* tv3 = set(tv2);
  fusion->addInput(tv0);
  fusion->addOutput(tv3);
  tv1->axis(0)->parallelize(ParallelType::Stream);
  tv3->axis(0)->parallelize(ParallelType::Stream);

Here, tv2 is not stream-parallelized so it should be be produced in a for-loop. Dump:

%HostIrContainer { (T0_g_float[iS0{i0}, iS1{i2}]) -> (T3_g_float[iStreamIdx6{i0}, iS7{i2}]) :
  T1_g_float[iStreamIdx2{i0}, iS3{i2}] = ALLOCATE(buffer=T1_g_float[iStreamIdx2{i0}, iS3{i2}], mem_type=global, size=( i0 * i2 ), zero_init=false, resets_to_zero=false)
  FOR StreamIdx in iStreamIdx2{i0}:
    GetCurrentStream into Stream 0
    SetCurrentStream to Stream ( StreamIdx % numberOfStreams )
    Synchronize Stream 0
    T4_l_float[iS8{i2}]
       = HirAliasSelect( T0_g_float[iS0{i0}, iS1{i2}], axis = iS0{i0}, index = StreamIdx )
    T5_l_float[iS9{i2}]
       = HirAliasSelect( T1_g_float[iStreamIdx2{i0}, iS3{i2}], axis = iStreamIdx2{i0}, index = StreamIdx )
    T5_l_float[iS9{i2}]
       = Set( T4_l_float[iS8{i2}], cache_op=Streaming )
    SetCurrentStream to Stream 0
    Synchronize Stream ( StreamIdx % numberOfStreams )
  T2_g_float[iS4{i0}, iS5{i2}]
     = Set( T1_g_float[iStreamIdx2{i0}, iS3{i2}], cache_op=Streaming )
  T3_g_float[iStreamIdx6{i0}, iS7{i2}] = ALLOCATE(buffer=T3_g_float[iStreamIdx6{i0}, iS7{i2}], mem_type=global, size=( i0 * i2 ), zero_init=false, resets_to_zero=false)
  FOR StreamIdx in iStreamIdx6{i0}:
    GetCurrentStream into Stream 2
    SetCurrentStream to Stream ( StreamIdx % numberOfStreams )
    Synchronize Stream 2
    T6_l_float[iS10{i2}]
       = HirAliasSelect( T2_g_float[iS4{i0}, iS5{i2}], axis = iS4{i0}, index = StreamIdx )
    T7_l_float[iS11{i2}]
       = HirAliasSelect( T3_g_float[iStreamIdx6{i0}, iS7{i2}], axis = iStreamIdx6{i0}, index = StreamIdx )
    T7_l_float[iS11{i2}]
       = Set( T6_l_float[iS10{i2}], cache_op=Streaming )
    SetCurrentStream to Stream 2
    Synchronize Stream ( StreamIdx % numberOfStreams )
} // %HostIrContainer

[github-actions]

Review updated until commit 35ff4da

Description

• Implement StreamParallelType pass for Host IR

• Add tests for Host IR stream lowering

• Fix swizzled_tiles calculation in fillDefaultHopperHeuristic

---

Changes walkthrough 📝

	Relevant files
Refactoring	1 files normalization_inner_outer.cpp Refactor and modularize inner-outer reduction scheduling  +21/-1604
Enhancement	11 files fusion_segmenter.cpp Add edge management and auxiliary group handling                  +226/-246 kernel_ir.cpp Add Continue expression and update utility names                  +23/-2    translation.cpp Add FusionTranslator for translating CPP Fusion to FusionDefinition [link]    stream_parallel_type.cpp Implement StreamParallelType pass for Host IR                        +440/-0  allocation.cpp Update AllocationInserter to handle MmaOp with Blackwell  +7/-5      utils.cpp Add IndexPutAccumulateOp to isTvOp check                                  +1/-15    fusion_cache.cpp Add FusionCache implementation for serialization and deserialization [link]    normalization_inner_outer_multi_wave.cpp Implement inner-outer multi-wave reduction scheduling        +717/-0  normalization_inner_outer_tma_ws.cpp Implement inner-outer TMA warp specialized reduction scheduling +647/-0  vectorize_helper.cpp Add logic to handle vectorization through Resize expressions +85/-64  utils.cpp Update sharding logic to handle reduction IterDomains        +15/-15
Configuration changes	2 files gen_nvfuser_version.py Link to external version generation script                              +0/-75    memory.py Link memory.py to the correct path                                              +0/-28
Tests	1 files test_host_ir_stream_lowering.cpp Add tests for Host IR stream lowering                                        +814/-0
Bug fix	1 files matmul_utils.cpp Fix swizzled_tiles calculation in fillDefaultHopperHeuristic +1/-1
Additional files	101 files build.yml +1/-0      lint.yml +6/-0      .lintrunner.toml +2/-2      CMakeLists.txt +29/-20  README.md +10/-0    core.py +12/-85  test_cross_entropy_loss.py +2/-2      test_matmul.py +6/-0      alias_analysis.cpp +0/-4      codegen.cpp +6/-2      predicate_elimination.cpp +1/-1      lower2device.h +6/-6      circular_buffer.cpp +326/-152 index.cpp +14/-1    index.h +1/-0      inline_ptx.cpp +1/-1      insert_syncs.cpp +59/-37  utils.h +0/-10    dispatch.h +4/-1      fusion_segmenter.h +20/-15  container.cpp +6/-4      executor.cpp +45/-8    executor.h +5/-0      host_ir.cpp +45/-0    host_ir.h +43/-0    lower.cpp +10/-0    stream_parallel_type.h +36/-0    interface_nodes.h +2/-2      internal_nodes.h +60/-0    iostream.cpp +2/-7      nodes.cpp +47/-0    kernel_ir.h +18/-1    logical_domain_map.cpp +33/-16  c10d_mock.h +29/-1    executor.h +4/-0      ipc_handle.cpp +6/-4      utils.h +9/-2      indexing.cpp +38/-0    indexing.h +6/-0      options.cpp +10/-17  options.h +22/-2    parallel_dimension_map.h +2/-2      predicate_compute.cpp +13/-13  insert_reshardings.cpp +55/-4    make_resharding_contiguous.cpp +122/-17 make_resharding_contiguous.h +11/-4    optimization_pass.h +0/-2      pre_segmenter.cpp +10/-6    propagate_shardings.cpp +255/-76 compiled_kernel.cpp +7/-0      executor_utils.h +2/-0      fusion_kernel_runtime.cpp +1/-52    expr_eval_sched.cpp +7/-1      hopper_multi_matmul.cpp +9/-44    hopper_multi_matmul.h +0/-8      multi_matmul.cpp +54/-0    multi_matmul.h +14/-0    normalization_inner_outer_multi_wave.h +31/-0    normalization_inner_outer_tma_ws.h +31/-0    normalization_inner_outer_utils.cpp +301/-0  normalization_inner_outer_utils.h +98/-0    normalization_utils.cpp +23/-5    reduction.cpp +13/-0    reduction_utils.cpp +42/-3    reduction_utils.h +3/-1      registry.cpp +6/-1      domain_map.cpp +12/-9    domain_map.h +1/-1      fusion_cache.fbs +1/-0      fusion_record.cpp +7/-0      type.cpp +3/-3      type.h +3/-3      nvfuser +1/-0      LICENSE +1/-0      README.md [link]    __init__.py [link]    __init__.pyi [link]    benchmark_utils.py +105/-0  __init__.py [link]    __init__.py [link]    normalization.py [link]    nvfuser_version.py [link]    pytorch_utils.py [link]    __init__.py [link]    utils.py [link]    utils.py [link]    pyproject.toml +3/-0      distributed_tensor.cpp [link]    distributed_tensor.h [link]    fusion_cache.h [link]    fusion_definition.cpp +5/-0      fusion_definition.h [link]    fusion_record.h +24/-0    fusion_state.cpp [link]    fusion_state.h [link]    multidevice_bindings.cpp [link]    python_bindings.cpp +42/-0    python_bindings.h [link]    python_bindings_extension.cpp [link]    schedule_bindings.cpp [link]    Additional files not shown

---

PR Reviewer Guide 🔍

Here are some key observations to aid the review process:

🧪 No relevant tests

⚡ Recommended focus areas for review Code Duplication The code for innerOuterWarpSpecializedTmaHeuristic and scheduleTmaWarpSpecializedInnerOuter seems to be very similar to the code for innerOuterPersistentHeuristic and scheduleInnerOuterPersistentKernel. Consider refactoring to reduce duplication. hp.threads_per_block_min, hp.threads_per_block_max); auto rparams = std::make_unique<ReductionParams>( InnerOuterPersistentKernelScheduler::schedulerType()); // Ultimately, we want the heuristic to decide between using the // warp-specialized version or the multi-wave version. The enable option is a // temporary configuration to facilitate testing during development without // disrupting existing behavior. if (isOptionEnabled(EnableOption::WarpSpecializedNormalization)) { inner_outer_tma_warp_specialized::getHeuristics( rparams.get(), properties.total_iteration_numel, properties.total_reduction_numel, buffer_params.regs_buffer_size, Removed Functionality Several functions such as scheduleReductionCombinedOuter, scheduleInnerOuterPersistentKernel, scheduleTmaWarpSpecializedOuter, and scheduleTmaWarpSpecializedInnerOuter have been removed. Ensure that the new code in inner_outer_multi_wave.cpp and inner_outer_tma_warp_specialized.cpp covers all the functionality of the removed functions. New File The new file python/python_frontend/schedule_bindings.cpp introduces a significant amount of new code. Ensure that all new bindings are correctly implemented and tested. // clang-format off /* * SPDX-FileCopyrightText: Copyright (c) 2023-present NVIDIA CORPORATION & AFFILIATES. * All rights reserved. * SPDX-License-Identifier: BSD-3-Clause */ // clang-format on #include <instrumentation.h> #include <ir/interface_nodes.h> #include <multidevice/utils.h> #include <python_frontend/fusion_cache.h> #include <python_frontend/fusion_definition.h> #include <python_frontend/python_bindings.h> #include <scheduler/tools/inlining.h> #include <transform_replay.h> namespace nvfuser::python_frontend { void bindSchedule(py::class_<FusionDefinition>& fusion_def) { //! The SchedOperators class is a nested class of FusionDefinition to allow //! the user to query the class for the list of schedule operators. //! //! Example: //! help(FusionDefinition.SchedOperators) //! //! Additional operators are expected to be defined below as needed. py::class_<FusionDefinition::SchedOperators> nvf_sched( fusion_def, "SchedOperators"); nvf_sched.def(py::init<FusionDefinition*>()); nvf_sched.def( "to_string", [](FusionDefinition::SchedOperators& self, Tensor tensor) { // NOTE: For debugging purposes, print the state of TensorView NVF_CHECK( self.validUse(), "Attempting to use a SchedOperators Op prior to definition!"); // Determine if tensor is a result from a reduction operation. FusionDefinition* fd = self.fusion_definition; TensorView* tv = fd->getFusionState(tensor.index)->template as<TensorView>(); return tv->toString(); }, py::arg("tensor")); nvf_sched.def( "user_schedule_ir", [](FusionDefinition::SchedOperators& self) { return self.fusion_definition->userScheduleIr(); }, py::return_value_policy::reference); //! experimental API for multidevice support nvf_sched.def( "_set_device_mesh", [](FusionDefinition::SchedOperators& self, Tensor tensor, const DeviceMesh& mesh) { NVF_CHECK( self.validUse(), "Attempting to use a SchedOperators Op prior to definition!"); FusionDefinition* fd = self.fusion_definition; auto tv = fd->getFusionState(tensor.index)->template as<TensorView>(); tv->setDeviceMesh(mesh); }, py::arg("tensor"), py::arg("mesh")); nvf_sched.def( "parallelize", [](FusionDefinition::SchedOperators& self, Tensor tensor, int axis, const ParallelType& parallel_type) { NVF_CHECK( self.validUse(), "Attempting to use a SchedOperators Op prior to definition!"); FusionDefinition* fd = self.fusion_definition; auto tv = fd->getFusionState(tensor.index)->template as<TensorView>(); tv->axis(axis)->parallelize(parallel_type); }, py::arg("tensor"), py::arg("axis"), py::arg("parallel_type")); nvf_sched.def( "merge", [](FusionDefinition::SchedOperators& self, Tensor arg, int dim) { FUSER_PERF_SCOPE("SchedOperators.merge"); NVF_CHECK( self.validUse(), "Attempting to use a SchedOperators Op prior to definition!"); FusionDefinition* fd = self.fusion_definition; auto input_tv = fd->getFusionState(arg.index)->template as<TensorView>(); input_tv->merge(dim); }, py::arg("arg"), py::arg("dim")); auto reduction_factor_func = [](FusionDefinition::SchedOperators& self, Tensor arg, const std::vector<int64_t>& dims) -> Tensor { FUSER_PERF_SCOPE("SchedOperators.reduction_factor"); NVF_CHECK( self.validUse(), "Attempting to use a SchedOperators Op prior to definition!"); FusionDefinition* fd = self.fusion_definition; TensorView* input_tv = fd->getFusionState(arg.index)->template as<TensorView>(); TensorView* output_tv = input_tv->rFactor(dims); return fd->addTensor(output_tv); }; nvf_sched.def( "reduction_factor", reduction_factor_func, py::arg("arg"), py::arg("dims")); nvf_sched.def( "rfactor", reduction_factor_func, py::arg("arg"), py::arg("dims")); nvf_sched.def( "reorder", [](FusionDefinition::SchedOperators& self, Tensor arg, const std::unordered_map<int64_t, int64_t>& old2new) { FUSER_PERF_SCOPE("SchedOperators.reorder"); NVF_CHECK( self.validUse(), "Attempting to use a SchedOperators Op prior to definition!"); FusionDefinition* fd = self.fusion_definition; auto input_tv = fd->getFusionState(arg.index)->template as<TensorView>(); input_tv->reorder(old2new); }, py::arg("arg"), py::arg("old2new")); nvf_sched.def( "split", [](FusionDefinition::SchedOperators& self, Tensor arg, int64_t dim, int64_t factor, bool inner_split) { FUSER_PERF_SCOPE("SchedOperators.split"); NVF_CHECK( self.validUse(), "Attempting to use a SchedOperators Op prior to definition!"); FusionDefinition* fd = self.fusion_definition; auto input_tv = fd->getFusionState(arg.index)->template as<TensorView>(); input_tv->split(dim, factor, inner_split); }, py::arg("arg"), py::arg("dim"), py::arg("factor"), py::arg("inner_split") = true); nvf_sched.def( "set_allocation_as_loop", [](FusionDefinition::SchedOperators& self, Tensor arg) { FUSER_PERF_SCOPE("SchedOperators.set_allocation_as_loop"); NVF_CHECK( self.validUse(), "Attempting to use a SchedOperators Op prior to definition!"); FusionDefinition* fd = self.fusion_definition; auto* tv = fd->getFusionState(arg.index)->template as<TensorView>(); tv->setAllocationDomain(tv->getLoopDomain(), true); }, py::arg("arg")); nvf_sched.def( "cache_after", [](FusionDefinition::SchedOperators& self, Tensor tensor, const LoadStoreOpType& op_type, const CacheOp& cache_op) -> Tensor { NVF_CHECK( self.validUse(), "Attempting to use a SchedOperators Op prior to definition!"); FusionDefinition* fd = self.fusion_definition; TensorView* input_tv = fd->getFusionState(tensor.index)->template as<TensorView>(); TensorView* output_tv = input_tv->cacheAfter(op_type, cache_op); return fd->addTensor(output_tv); }, py::arg("tensor"), py::arg("op_type") = LoadStoreOpType::Set, py::arg("cache_op") = CacheOp::Unspecified); nvf_sched.def( "cache_before", [](FusionDefinition::SchedOperators& self, Tensor tensor, const LoadStoreOpType& op_type) -> Tensor { NVF_CHECK( self.validUse(), "Attempting to use a SchedOperators Op prior to definition!"); FusionDefinition* fd = self.fusion_definition; TensorView* input_tv = fd->getFusionState(tensor.index)->template as<TensorView>(); TensorView* output_tv = input_tv->cacheBefore(op_type); return fd->addTensor(output_tv); }, py::arg("tensor"), py::arg("op_type") = LoadStoreOpType::Set); nvf_sched.def( "cache_fork", [](FusionDefinition::SchedOperators& self, Tensor tensor) -> Tensor { NVF_CHECK( self.validUse(), "Attempting to use a SchedOperators Op prior to definition!"); FusionDefinition* fd = self.fusion_definition; TensorView* input_tv = fd->getFusionState(tensor.index)->template as<TensorView>(); TensorView* output_tv = input_tv->cacheFork(); return fd->addTensor(output_tv); }, py::arg("tensor")); nvf_sched.def( "set_memory_type", [](FusionDefinition::SchedOperators& self, Tensor tensor, const MemoryType& memory_type) { NVF_CHECK( self.validUse(), "Attempting to use a SchedOperators Op prior to definition!"); FusionDefinition* fd = self.fusion_definition; TensorView* tv = fd->getFusionState(tensor.index)->template as<TensorView>(); tv->setMemoryType(memory_type); }, py::arg("tensor"), py::arg("memory_type")); nvf_sched.def( "transform_like", [](FusionDefinition::SchedOperators& self, Tensor tensor, const std::vector<Tensor>& selected_tensors) { NVF_CHECK( self.validUse(), "Attempting to use a SchedOperators Op prior to definition!"); FusionDefinition* fd = self.fusion_definition; TensorView* reference_tv = fd->getFusionState(tensor.index)->template as<TensorView>(); TransformPropagator propagator(reference_tv); if (selected_tensors.empty()) { // Propagate scheduler transformations on reference TensorView to the // rest of the fusion. MaxLogicalDomainInfoSpanningTree(reference_tv).traverse(&propagator); } else { // Propagate scheduler transformations on reference TensorView to the // subset of the fusion. std::unordered_set<TensorView*> selected_tv_set; selected_tv_set.reserve(selected_tensors.size()); std::transform( selected_tensors.begin(), selected_tensors.end(), std::inserter(selected_tv_set, selected_tv_set.end()), [&fd](const Tensor& t) { return fd->getFusionState(t.index)->template as<TensorView>(); }); SetSelector selector( {selected_tv_set.begin(), selected_tv_set.end()}); MaxLogicalDomainInfoSpanningTree(reference_tv, &selector) .traverse(&propagator); } }, py::arg("tensor"), py::arg("selected_tensors") = std::vector<Tensor>()); nvf_sched.def( "parallelize_like", [](FusionDefinition::SchedOperators& self, Tensor tensor, int64_t pos, const std::vector<Tensor>& selected_tensors, const std::unordered_set<ParallelType>& selected_parallel_types, bool propagate_padding) { // Propagate the parallelization from the selected dimensions of the // reference tensor to their corresponding dimensions in all selected // tensors in the DAG. // // 1. Position `pos` means selecting all the dimensions // [0, 1, ..., pos - 1]. pos = -1 means selecting all dimensions. // 2. `selected_tvs` are selected tensors in the DAG. Empty // `selected_tvs` means selecting all tensors in the fusion of // `reference_tv`. // 3. `selected_parallel_types` are the selected parallel types. Empty // `selected_parallel_types` means selecting all parallel types. NVF_CHECK( self.validUse(), "Attempting to use a SchedOperators Op prior to definition!"); FusionDefinition* fd = self.fusion_definition; TensorView* reference_tv = fd->getFusionState(tensor.index)->template as<TensorView>(); std::vector<TensorView*> selected_tvs; selected_tvs.reserve(selected_tensors.size()); std::transform( selected_tensors.begin(), selected_tensors.end(), std::back_inserter(selected_tvs), [&fd](const Tensor& t) { return fd->getFusionState(t.index)->template as<TensorView>(); }); nvfuser::scheduler_utils::parallelizeAllLike( reference_tv, pos, selected_tvs, selected_parallel_types, propagate_padding); }, py::arg("tensor"), py::arg("pos") = -1, py::arg("selected_tensors") = std::vector<Tensor>(), py::arg("selected_parallel_types") = std::unordered_set<ParallelType>(), py::arg("propagate_padding") = true); nvf_sched.def( "inline_most", [](FusionDefinition::SchedOperators& self, const std::vector<Tensor>& selected_tensors) { // Inline to the right most allowed position for the selected tensors in // the current fusion. NVF_CHECK( self.validUse(), "Attempting to use a SchedOperators Op prior to definition!"); FusionDefinition* fd = self.fusion_definition; if (selected_tensors.empty()) { nvfuser::inlineMost(); } else { std::vector<TensorView*> selected_tvs; selected_tvs.reserve(selected_tensors.size()); std::transform( selected_tensors.begin(), selected_tensors.end(), std::back_inserter(selected_tvs), [&fd](const Tensor& t) { return fd->getFusionState(t.index)->template as<TensorView>(); }); nvfuser::inlineMost(selected_tvs); } }, py::arg("selected_tensors") = std::vector<Tensor>()); nvf_sched.def( "inline_at", [](FusionDefinition::SchedOperators& self, Tensor tensor, int64_t pos, bool best_effort, const std::vector<Tensor>& selected_tensors) { NVF_CHECK( self.validUse(), "Attempting to use a SchedOperators Op prior to definition!"); FusionDefinition* fd = self.fusion_definition; TensorView* reference_tv = fd->getFusionState(tensor.index)->template as<TensorView>(); if (selected_tensors.empty()) { // Inline to the position corresponding to the reference position in // the reference tensor for all tensors in the current fusion. nvfuser::inlineAllAt(reference_tv, pos, best_effort); } else { // Inline to the position corresponding to the reference position in // the reference tensor for selected tensors in the current fusion. std::unordered_set<TensorView*> selected_tvs; selected_tvs.reserve(selected_tensors.size()); std::transform( selected_tensors.begin(), selected_tensors.end(), std::inserter(selected_tvs, selected_tvs.end()), [&fd](const Tensor& t) { return fd->getFusionState(t.index)->template as<TensorView>(); }); nvfuser::inlineSelectedAt( selected_tvs, reference_tv, pos, best_effort); } }, py::arg("tensor"), py::arg("pos") = -1, py::arg("best_effort") = false, py::arg("selected_tensors") = std::vector<Tensor>()); nvf_sched.def("tensors", [](FusionDefinition::SchedOperators& self) { NVF_CHECK( self.validUse(), "Attempting to use a SchedOperators Op prior to definition!"); // Return all Tensors in FusionDefinition return self.fusion_definition->tensors(); }); nvf_sched.def( "is_reduction", [](FusionDefinition::SchedOperators& self, Tensor tensor) { NVF_CHECK( self.validUse(), "Attempting to use a SchedOperators Op prior to definition!"); // Determine if tensor is a result from a reduction operation. FusionDefinition* fd = self.fusion_definition; TensorView* tv = fd->getFusionState(tensor.index)->template as<TensorView>(); return ( !tv->isFusionInput() && std::any_of( tv->getMaybeRootDomain().begin(), tv->getMaybeRootDomain().end(), [](IterDomain* id) { return id->isReduction(); }) && !isResharding(tv->definition())); }, py::arg("tensor")); nvf_sched.def( "can_schedule", [](FusionDefinition::SchedOperators& self, const SchedulerType& scheduler_type) { NVF_CHECK( self.validUse(), "Attempting to use a SchedOperators Op prior to definition!"); return self.fusion_definition->userSchedule()->canScheduleDebug( scheduler_type); }, py::arg("scheduler_type")); nvf_sched.def( "find_compatible_schedulers", [](FusionDefinition::SchedOperators& self) { NVF_CHECK( self.validUse(), "Attempting to use a SchedOperators Op prior to definition!"); std::vector<SchedulerType> valid_scheduler_types; valid_scheduler_types.reserve(all_heuristics_in_priority_order.size()); std::copy_if( all_heuristics_in_priority_order.begin(), all_heuristics_in_priority_order.end(), std::back_inserter(valid_scheduler_types), [sched = self.fusion_definition->userSchedule()]( SchedulerType scheduler_type) { return sched->canSchedule(scheduler_type); }); return valid_scheduler_types; }); nvf_sched.def( "schedule", [](FusionDefinition::SchedOperators& self, const SchedulerType& scheduler_type) { NVF_CHECK( self.validUse(), "Attempting to use a SchedOperators Op prior to definition!"); UserSchedule* sched = self.fusion_definition->userSchedule(); auto&& [can_schedule, error_msg] = sched->canScheduleDebug(scheduler_type); NVF_CHECK(can_schedule, error_msg); sched->scheduleWithType(scheduler_type); }, py::arg("heuristic")); nvf_sched.def("schedule", [](FusionDefinition::SchedOperators& self) { NVF_CHECK( self.validUse(), "Attempting to use a SchedOperators Op prior to definition!"); UserSchedule* sched = self.fusion_definition->userSchedule(); sched->schedule(); }); nvf_sched.def( "compute_pointwise_heuristics", [](FusionDefinition::SchedOperators& self) -> PointwiseParams& { NVF_CHECK( self.validUse(), "Attempting to use a SchedOperators Op prior to definition!"); UserSchedule* sched = self.fusion_definition->userSchedule(); HeuristicParams* parameters = sched->computeHeuristics(SchedulerType::PointWise); return *parameters->as<PointwiseParams>(); }, py::return_value_policy::reference); nvf_sched.def( "compute_reduction_heuristics", [](FusionDefinition::SchedOperators& self) -> ReductionParams& { NVF_CHECK( self.validUse(), "Attempting to use a SchedOperators Op prior to definition!"); UserSchedule* sched = self.fusion_definition->userSchedule(); HeuristicParams* parameters = sched->computeHeuristics(SchedulerType::Reduction); return *parameters->as<ReductionParams>(); }, py::return_value_policy::reference); nvf_sched.def( "compute_matmul_heuristics", [](FusionDefinition::SchedOperators& self) -> MatmulParams& { NVF_CHECK( self.validUse(), "Attempting to use a SchedOperators Op prior to definition!"); UserSchedule* sched = self.fusion_definition->userSchedule(); HeuristicParams* parameters = sched->computeHeuristics(SchedulerType::Matmul); return *parameters->as<MatmulParams>(); }, py::return_value_policy::reference); nvf_sched.def( "schedule_hyperparameters", [](FusionDefinition::SchedOperators& self) -> scheduler_utils::SchedulerHyperParameters& { NVF_CHECK( self.validUse(), "Attempting to use a SchedOperators Op prior to definition!"); UserSchedule* sched = self.fusion_definition->userSchedule(); auto scheduler_hyperparameters_entry = HeuristicDataCacheEntry< HeuristicCompileTime::SchedulerHyperParameters>( sched->data_cache.get(), []() { return std::make_unique< scheduler_utils::SchedulerHyperParameters>( /*vectorize_factor=*/1, /*unroll_factor=*/1, /*threads_per_block_min=*/1, /*threads_per_block_max=*/1); }); return scheduler_hyperparameters_entry.get(); }, py::return_value_policy::reference); } } // namespace nvfuser::python_frontend

[samnordmann]

I suggest to start the review by reading the PR description and inspecting those tests, running them with the option NVFUSER_DUMP=host_ir.

There are two identical sets of tests: HirLowerStreamTest and MultiDeviceExecutorLowerStreamTest. Both sets exercise semantically identical situations, but the former is "lower level" in the sense that it manually builds a Hir container and apply the pass, while the latter goes through the main level API for defining a fusion, which then gets automatically lowered and executed through MultiDeviceExecutor

[samnordmann]

!test

[samnordmann]

I don't understand the logic behind removing the reduction axis when selecting: since it is an aliasing op, it does not produce the reduction so should keep it. Anyway, this is problematic in our context where we need the reduction axis since the selected tensor will actually become the reduced tensor. So, I decided to add this option to control this behavior

[wujingyue]

IIUC, a reduction dimension always stops at the op that performs the reduction. So the original code looks right and I don't follow the motivation of this change. Do you have an example where we have to keep the reduction dimension in a downstream op?

cc @naoyam in case I missed anything

[samnordmann]

IIUC, a reduction dimension always stops at the op that performs the reduction. So the original code looks right and I don't follow the motivation of this change.

Yes, but indeed, a select op do not performs the reduction, this is why is makes sense to keep the reduction axis.

Do you have an example where we have to keep the reduction dimension in a downstream op?

Consider the case of MultiDeviceExecutorLowerStreamTest.Reduction, that I am reproducing here:

auto tv0 = makeContigTensor(3); // [i0, i1, i2]
auto tv1 = sum(tv0, {2}); // [Stream(i0), i1, r2]

tv1->axis(0)->parallelize(ParallelType::Stream);

The stream lowering pass implements indexing into tv0 and tv1 by creating SelectOps, and letting the reduction operate on the selected tensors. The dumped Host program looks like:

%HostIrContainer { (T0_g_float[iS0{i0}, iS1{i2}, iS2{i3}]) -> (T1_g_float[iStream3{i0}, iS4{i2}, rS5{i3}]) :
  T1_g_float[iStream3{i0}, iS4{i2}, rS5{i3}] = ALLOCATE(buffer=T1_g_float[iStream3{i0}, iS4{i2}, rS5{i3}], mem_type=global, size=( i0 * i2 ), zero_init=false, resets_to_zero=false)
  FOR i10 in iStream3{i0}:
    GetCurrentStream into Stream 0
    SetCurrentStream to Stream ( i10 % numberOfStreams )
    Synchronize Stream 0
    T2_l_float[iS6{i2}, iS7{i3}]
       = select( T0_g_float[iS0{i0}, iS1{i2}, iS2{i3}], axis = iS0{i0}, index = i10 )
    T3_l_float[iS8{i2}, rS9{i3}]
       = select( T1_g_float[iStream3{i0}, iS4{i2}, rS5{i3}], axis = iStream3{i0}, index = i10 )
    T3_l_float[iS8{i2}, rS9{i3}]
       = reduction( T2_l_float[iS6{i2}, iS7{i3}], op = add, initial value = float(0), allreduce = false )
    SetCurrentStream to Stream 0
    Synchronize Stream ( i10 % numberOfStreams )
} // %HostIrContainer

The sum now operates not on tv0 and tv1, but on alias slices of it. This is why we need to keep the reduction axis on tv1's SelectOp's consumer.

I understand I am using SelectOp here to index, while it was originally thought of as a tensor's semantic op given in the user's DAG. But Hir indexing has other needs than kernel indexing: I cannot directly index into C-style arrays, I need to use at::select ops (or other ATen ops), to then feed the indexed tensor into other ATen ops (such as at::matmul, since we do not use matmul generated kernel for now). This is why I need to extend the usage of SelectOp.

[samnordmann]

btw, if this raises too much concern, I can create another HIR op to avoid collisions with SelectOp

[naoyam]

IIUC, a reduction dimension always stops at the op that performs the reduction.

That's right.

[naoyam]

Yes, but indeed, a select op do not performs the reduction, this is why is makes sense to keep the reduction axis.

It isn't about this op, but it's the producer tensor that has reduction IDs. They are there because the op that produces the produce is a reduction. Reduction IDs are not be kept around after the reduction op itself.

[wujingyue]

I can create another HIR op to avoid collisions with SelectOp

SGTM to unblock your experiment. Thanks for the example! IIUC, both the second select and the reduction define T3, causing a mismatch in r. Arguably, select isn't designed for this use case because it belongs to fusion IR that's supposed to be SSA. So creating a new Host IR for this use case makes sense to me.

[samnordmann]

I did as you suggest and created a dedicated HIR in #4301

It actually solves other problems along the way, and allowed me to

1. get rid of aliases in hic: indeed, I do not set the indexed tensor's definition when indexing, thus avoiding cycling dependency. HIC's alias is therefore not necessary anymore and I'll remove it in a separate PR
2. remove the bool keep_reduction_axis option I added in select, newValLike, etc.

So the code is cleaner.

[samnordmann]

!test

[samnordmann]

in the last commit move stream_parallel_type to host_ir/pass folder, I did as you suggested in #4145 (comment) and moved the stream lowering pass to a different folder. However, I let it inherit from OptimizationPass for now since I need the "disable guard" already in the current pr, which is already big. This guard is only needed temporarily in the tests which mix multi device and stream parallel type. Soon, the disable guards won't be needed anymore, and I'll remove the inheritance of the pass.

[samnordmann]

I pushed another (hopefully) last commit to further cleanup stream_parallel_type.cpp

[samnordmann]

!test

[wujingyue]

Looking great!

FYI, I skipped stream_parallel_type.cpp because I may end up doing loop inlining differently and don't want to slow down your experiments.

[samnordmann]

!test

[samnordmann]

I may end up doing loop inlining differently

For my own learning, I'd be interested to hear your ideas