From b33b130916760fc86a027dd832cfd0b6ffe0dcdc Mon Sep 17 00:00:00 2001 From: Naoya Maruyama Date: Fri, 27 Feb 2026 09:54:40 -0800 Subject: [PATCH 1/7] Remove fallback --- csrc/index_compute.cpp | 189 +++++++++++++++-------------------------- csrc/index_compute.h | 1 + 2 files changed, 69 insertions(+), 121 deletions(-) diff --git a/csrc/index_compute.cpp b/csrc/index_compute.cpp index 1829afa8138..512104d777a 100644 --- a/csrc/index_compute.cpp +++ b/csrc/index_compute.cpp @@ -1413,52 +1413,42 @@ std::vector Index::getNonGlobalProducerStridedIndices( Val* Index::getLinearLogicalIndex( TensorView* consumer_tv, const std::vector& loops) { - if (!ir_utils::hasRootToLoopLinearTransformations(consumer_tv) || - ir_utils::isCpAsyncBulkLoad(consumer_tv->definition()) || - GpuLower::current()->idModelOptions().isTensorIndexerEnabled() || - GpuLower::current()->tmemInfo().hasTMemTensor()) { - const TensorIndexer& indexer = GpuLower::current()->tensorIndexer(); - auto per_dim_indices = indexer.getIndexFor( - consumer_tv->definition(), - /*as_consumer=*/true, - consumer_tv->getLogicalDomain(), - loops, - /*use_magic_zero=*/true); - Val* stride = consumer_tv->fusion()->oneVal(); - for (const auto [i, logical_id] : - enumerate(consumer_tv->getLogicalDomain()) | std::views::reverse) { - auto per_dim_index = per_dim_indices.at(i); - auto per_dim_strided_index = - SimplifyingIrBuilder::mulExpr(per_dim_index, stride); - per_dim_indices.at(i) = per_dim_strided_index; - stride = SimplifyingIrBuilder::mulExpr(stride, logical_id->extent()); - } - return sumVals(per_dim_indices); - } else { - auto guard = ir_utils::allocateToLogicalDomainGuard(consumer_tv, true); - return sumVals(getGlobalConsumerStridedIndices(consumer_tv, loops)); - } + NVF_ERROR( + GpuLower::current()->idModelOptions().isTensorIndexerEnabled(), + "Legacy indexer no longer available"); + + const TensorIndexer& indexer = GpuLower::current()->tensorIndexer(); + auto per_dim_indices = indexer.getIndexFor( + consumer_tv->definition(), + /*as_consumer=*/true, + consumer_tv->getLogicalDomain(), + loops, + /*use_magic_zero=*/true); + Val* stride = consumer_tv->fusion()->oneVal(); + for (const auto [i, logical_id] : + enumerate(consumer_tv->getLogicalDomain()) | std::views::reverse) { + auto per_dim_index = per_dim_indices.at(i); + auto per_dim_strided_index = + SimplifyingIrBuilder::mulExpr(per_dim_index, stride); + per_dim_indices.at(i) = per_dim_strided_index; + stride = SimplifyingIrBuilder::mulExpr(stride, logical_id->extent()); + } + return sumVals(per_dim_indices); } std::vector Index::getConsumerPerDimLogicalIndex( TensorView* consumer_tv, const std::vector& loops) { - if (!ir_utils::hasRootToLoopLinearTransformations(consumer_tv) || - GpuLower::current()->idModelOptions().isTensorIndexerEnabled() || - GpuLower::current()->tmemInfo().hasTMemTensor()) { - const TensorIndexer& indexer = GpuLower::current()->tensorIndexer(); - return indexer.getIndexFor( - consumer_tv->definition(), - /*as_consumer=*/true, - consumer_tv->getLogicalDomain(), - loops); - } else { - auto guard = ir_utils::allocateToLogicalDomainGuard(consumer_tv, false); - IndexFromIdGraph index_from_id_graph = - getTensorIndexFromIdGraph(loops, consumer_tv); - return getConsumerAllocationIndices( - consumer_tv, loops, index_from_id_graph); - } + NVF_ERROR( + GpuLower::current()->idModelOptions().isTensorIndexerEnabled(), + "Legacy indexer no longer available"); + + const TensorIndexer& indexer = GpuLower::current()->tensorIndexer(); + return indexer.getIndexFor( + consumer_tv->definition(), + /*as_consumer=*/true, + consumer_tv->getLogicalDomain(), + loops); } std::vector Index::getProducerPerDimLogicalIndex( @@ -1913,38 +1903,6 @@ Val* Index::getProducerStridedIndices( } } -namespace { - -bool shouldUseTensorIndexer( - const TensorView* producer, - const TensorView* consumer) { - // Check if TensorIndexer is definitely required - auto is_tensor_indexer_required = [&]() -> bool { - bool is_producer_tma_op = producer->definition() != nullptr && - producer->definition()->isA() && - ir_utils::isCpAsyncBulkLoad(producer->definition()); - bool is_consumer_tma_op = consumer->definition() != nullptr && - consumer->definition()->isA() && - ir_utils::isCpAsyncBulkLoad(consumer->definition()); - - return !ir_utils::hasRootToLoopLinearTransformations(producer) || - (consumer->definition()->isA() && - isHopper(consumer->definition()->as()->macro())) || - is_producer_tma_op || is_consumer_tma_op || - GpuLower::current()->tmemInfo().hasTMemTensor(); - }; - - // TensorIndexer is always used when required or if not disabled. - // Note: Previously, ldmatrix and stmatrix were first introduced - // with Ampere, their indexing were only implemented in the legacy - // indexer in a rather manual way. The current implementation uses - // the alternate loop domain to enable TensorIndexer-based indexing. - return is_tensor_indexer_required() || - GpuLower::current()->idModelOptions().isTensorIndexerEnabled(); -} - -} // namespace - // Producer is the inputs of an expression kir::TensorIndex* Index::getProducerIndex( TensorView* producer, @@ -1954,28 +1912,25 @@ kir::TensorIndex* Index::getProducerIndex( bool generate_pointer, DataType as_type, bool ld_st_matrix) { - Val* index = nullptr; - - if (shouldUseTensorIndexer(producer, consumer)) { - index = GpuLower::current()->tensorIndexer().getLinearIndex( - producer, consumer->definition(), loops, override_index, ld_st_matrix); - if (generate_pointer) { - auto address_offset = index; - if (producer->getMemoryType() == MemoryType::Shared) { - auto producer_dt = producer->getDataType(); - NVF_ERROR(producer_dt.has_value()); - auto index_dt = index->getDataType(); - NVF_ERROR(index_dt.has_value()); - address_offset = SimplifyingIrBuilder::mulExpr( - address_offset, - IrBuilder::create(dataTypeSizeByte(*producer_dt), *index_dt)); - } - index = SimplifyingIrBuilder::addExpr( - IrBuilder::baseAddressExpr(producer), address_offset); - } - } else { - index = getProducerStridedIndices( - producer, consumer, loops, override_index, generate_pointer); + NVF_ERROR( + GpuLower::current()->idModelOptions().isTensorIndexerEnabled(), + "Legacy indexer no longer available"); + + Val* index = GpuLower::current()->tensorIndexer().getLinearIndex( + producer, consumer->definition(), loops, override_index, ld_st_matrix); + if (generate_pointer) { + auto address_offset = index; + if (producer->getMemoryType() == MemoryType::Shared) { + auto producer_dt = producer->getDataType(); + NVF_ERROR(producer_dt.has_value()); + auto index_dt = index->getDataType(); + NVF_ERROR(index_dt.has_value()); + address_offset = SimplifyingIrBuilder::mulExpr( + address_offset, + IrBuilder::create(dataTypeSizeByte(*producer_dt), *index_dt)); + } + index = SimplifyingIrBuilder::addExpr( + IrBuilder::baseAddressExpr(producer), address_offset); } index = GpuLower::current()->commonScalarMap().hoistScalar(index, loops); @@ -2050,33 +2005,25 @@ kir::TensorIndex* Index::getConsumerIndex( bool generate_pointer, DataType as_type, bool ld_st_matrix) { - Val* index = nullptr; - if (!ir_utils::hasRootToLoopLinearTransformations(consumer) || - ir_utils::isCpAsyncBulkLoad(consumer->definition()) || - GpuLower::current()->idModelOptions().isTensorIndexerEnabled() || - GpuLower::current()->tmemInfo().hasTMemTensor()) { - index = GpuLower::current()->tensorIndexer().getLinearIndex( - consumer, consumer->definition(), loops, override_index, ld_st_matrix); - if (generate_pointer) { - auto address_offset = index; - if (consumer->getMemoryType() == MemoryType::Shared) { - auto consumer_dt = consumer->getDataType(); - NVF_ERROR(consumer_dt.has_value()); - auto index_dt = index->getDataType(); - NVF_ERROR(index_dt.has_value()); - address_offset = SimplifyingIrBuilder::mulExpr( - index, - IrBuilder::create(dataTypeSizeByte(*consumer_dt), *index_dt)); - } - index = SimplifyingIrBuilder::addExpr( - IrBuilder::baseAddressExpr(consumer), address_offset); + NVF_ERROR( + GpuLower::current()->idModelOptions().isTensorIndexerEnabled(), + "Legacy indexer no longer available"); + + Val* index = GpuLower::current()->tensorIndexer().getLinearIndex( + consumer, consumer->definition(), loops, override_index, ld_st_matrix); + if (generate_pointer) { + auto address_offset = index; + if (consumer->getMemoryType() == MemoryType::Shared) { + auto consumer_dt = consumer->getDataType(); + NVF_ERROR(consumer_dt.has_value()); + auto index_dt = index->getDataType(); + NVF_ERROR(index_dt.has_value()); + address_offset = SimplifyingIrBuilder::mulExpr( + index, + IrBuilder::create(dataTypeSizeByte(*consumer_dt), *index_dt)); } - } else { - NVF_ERROR( - override_index.empty(), - "Overriding of consumer indexing with the legacy indexer is not " - "supported"); - index = getConsumerStridedIndices(consumer, loops, generate_pointer); + index = SimplifyingIrBuilder::addExpr( + IrBuilder::baseAddressExpr(consumer), address_offset); } index = GpuLower::current()->commonScalarMap().hoistScalar(index, loops); diff --git a/csrc/index_compute.h b/csrc/index_compute.h index 544388557c0..8ee430bd373 100644 --- a/csrc/index_compute.h +++ b/csrc/index_compute.h @@ -480,6 +480,7 @@ class Index { const std::unordered_map& override_index = {}, bool generate_pointer = false); + // TODO: Remove //! Returns a vector of strided indices mapped onto the //! allocation domain of a consumer tensor. The size of the returned //! vector is guaranteed to be equal to the number of axes of the From c1dc46b11f731453ad35b5a5fcb2fa1121469ec8 Mon Sep 17 00:00:00 2001 From: Naoya Maruyama Date: Fri, 27 Feb 2026 11:21:22 -0800 Subject: [PATCH 2/7] cleanup --- csrc/index_compute.cpp | 704 +---------------------------------------- csrc/index_compute.h | 60 +--- 2 files changed, 11 insertions(+), 753 deletions(-) diff --git a/csrc/index_compute.cpp b/csrc/index_compute.cpp index 512104d777a..66454c7643e 100644 --- a/csrc/index_compute.cpp +++ b/csrc/index_compute.cpp @@ -1100,90 +1100,6 @@ std::unordered_map invertOneToOneMap( } // namespace -std::vector Index::getGlobalProducerStridedIndices( - TensorView* producer_tv, - const TensorView* consumer_tv, - const std::vector& loops, - const std::unordered_map& override_index) { - FUSER_PERF_SCOPE("GpuLower::Lower::getGlobalProducerIndex"); - - auto alloc_indices = getProducerAllocationIndices( - producer_tv, consumer_tv, loops, override_index); - - const auto& alloc_dom = producer_tv->getMaybeAllocationDomain(); - - // TODO: Abstract stride logic to reuse with consumer indexing - std::vector strides(alloc_dom.size(), nullptr); - { - int stride_i = 0; - for (const auto i : arange(alloc_dom.size())) { - if (alloc_dom[i]->isReduction()) { - strides[i] = GpuLower::current()->kernel()->oneVal(); - continue; - } - strides[i] = IrBuilder::getItemExpr( - IrBuilder::getAttrExpr( - IrBuilder::metadataExpr(producer_tv), "alloc_stride"), - (int64_t)stride_i++); - } - } - - NVF_ERROR(alloc_dom.size() == producer_tv->domain()->contiguity().size()); - Val* cur_contig_stride = GpuLower::current()->kernel()->oneVal(); - for (const auto i : arange(alloc_dom.size())) { - auto dim = alloc_dom.size() - i - 1; - if (alloc_dom[dim]->isReduction()) { - continue; - } - - auto producer_dim_contiguity = producer_tv->domain()->contiguity().at(dim); - if (alloc_dom[dim]->isBroadcast()) { - strides[dim] = cur_contig_stride->fusion()->zeroVal(); - NVF_ERROR(!producer_dim_contiguity.has_value()); - } else if (!producer_dim_contiguity.has_value()) { - NVF_THROW("Expected value for dimension contiguity"); - } else if (producer_dim_contiguity.value()) { - // If contig, used the stored stride which may be the previous - // dimensions stride * previous dimensions size - strides[dim] = cur_contig_stride; - // Prepare for the next dimension which may also be contiguous, multiply - // by extent of this dimension - auto alloc_dim_extent = getExtentOfRootAxis(alloc_dom[dim]); - cur_contig_stride = - SimplifyingIrBuilder::mulExpr(cur_contig_stride, alloc_dim_extent); - } else { - // If non contiguous dimension, keep local stride information, set cur - // stride to local stride * local raw extent - auto alloc_dim_extent = getExtentOfRootAxis(alloc_dom[dim]); - cur_contig_stride = - SimplifyingIrBuilder::mulExpr(strides[dim], alloc_dim_extent); - } - } - - auto vectorize_shift = - loops.empty() ? nullptr : loops.back()->vectorize_shift(); - - // Global striding - std::vector strided_inds( - alloc_dom.size(), GpuLower::current()->kernel()->zeroVal()); - for (const auto i : arange(alloc_dom.size())) { - Val* alloc_ind = alloc_indices.at(i); - if (alloc_ind->isZeroInt()) { - continue; - } - - auto strided_ind = SimplifyingIrBuilder::mulExpr(alloc_ind, strides[i]); - if (i == alloc_dom.size() - 1 && vectorize_shift != nullptr) { - strided_inds[i] = - SimplifyingIrBuilder::addExpr(strided_ind, vectorize_shift); - } else { - strided_inds[i] = strided_ind; - } - } - - return strided_inds; -} - namespace { // Maps all producer domains to consumer with broadcast @@ -1230,186 +1146,6 @@ Val* sumVals(std::vector vals) { } // namespace -// Producer index for either shared or local memory -std::vector Index::getNonGlobalProducerStridedIndices( - TensorView* producer_tv, - const TensorView* consumer_tv, - const std::vector& loops, - const std::unordered_map& override_index) { - bool is_mma_input = consumer_tv->definition()->isA(); - const auto gpu_lower = GpuLower::current(); - // Replay producer to look like consumer so we can index on producer since our - // loop nests look like consumer - auto pairwise_map = PairwiseLogicalDomainMap(producer_tv, consumer_tv); - // Resize ops can be and should be replayed. - auto producer_replayed_as_consumer = - TransformReplay::replayPasC( - producer_tv, - consumer_tv, - -1, - pairwise_map, - TransformReplayOptions().replayResize()) - .first; - - ir_utils::TVDomainGuard domain_guard( - producer_tv, producer_replayed_as_consumer); - const auto p2c_alloc_map = - mapAllProducerDomainsToConsumer(producer_tv, consumer_tv); - - // Map everything we can from reference to producer using compute at index - // map. All producer id's don't exist in the compute at map. The logical axes - // all may be, but since I haven't proven that to be the case, going to do a - // more conservative approach, which is to use the consumer as a proxy between - // producer to reference. - std::unordered_map index_map_ref_to_producer; - std::unordered_map c2p_index_map; - - // Map sent to best effort replay needs to match the exact incantation for - // compute_at_mode.cpp with MappingMode::Index - auto c2p_logical_map = PairwiseLogicalDomainMap(producer_tv, consumer_tv) - .mapBroadcast(false) - .mapConsumerToProducer(); - - // This replay has to be consistent with compute at index map. - BestEffortReplay replay_producer_as_consumer( - producer_tv->getLoopDomain(), - consumer_tv->getLoopDomain(), - c2p_logical_map); - - c2p_index_map = replay_producer_as_consumer.getReplay(); - - const auto& producer_indexing_from_idgraph = getTensorIndexFromIdGraph( - loops, consumer_tv, producer_tv, false, c2p_index_map); - - const auto& producer_indexing = producer_indexing_from_idgraph.index; - - // TODO: merge the two swizzle compute logic once the new one is ready. - // will need to replace cyclic shift swizzle with xor since swizzle2d - // doesn't have cyclic shift. - const auto& index_map = producer_indexing.indexMap(); - - const auto& extent_map = producer_indexing.extentMap(); - const auto& zero_domain_map = producer_indexing.zeroDomains(); - // Indices should now be mapped onto IterDomains in producer, so just grab - // and use them. - const auto& alloc_dom = producer_tv->getMaybeAllocationDomain(); - - // Figure out which alloc axes we don't need to index - std::unordered_set skip_indexing; - - for (auto alloc_id : alloc_dom) { - // Already taken care of because we can detect no indexing required - if (alloc_id->isBroadcast() || alloc_id->isReduction() || - alloc_id->isStride() || alloc_id->isDeviceDim() || - (alloc_id->isThread() && - producer_tv->getMemoryType() == MemoryType::Local)) { - skip_indexing.insert(alloc_id); - continue; - } - - // Already an entry for this allocation domain, continue - if (index_map.find(alloc_id) != index_map.end()) { - continue; - } - } - - std::vector strided_inds( - alloc_dom.size(), GpuLower::current()->kernel()->zeroVal()); - - // MMA operation op is a special operation that our automatic "zero domain" - // analysis of our current indexing approach does not work. So we need to - // manually specify which dimensions are used for MMA allocation. - std::function is_mma_allocation; - if (is_mma_input) { - int size = (int)alloc_dom.size(); - const IterDomain* allocation0 = alloc_dom.at(size - 3); - const IterDomain* allocation1 = alloc_dom.at(size - 2); - const IterDomain* allocation2 = alloc_dom.at(size - 1); - is_mma_allocation = [=](const IterDomain* id) { - return id == allocation0 || id == allocation1 || id == allocation2; - }; - } else { - is_mma_allocation = [](const IterDomain* id) { return false; }; - } - - for (const auto i : arange(alloc_dom.size())) { - if (skip_indexing.count(alloc_dom[i])) { - continue; - } - - auto override_it = override_index.find(alloc_dom[i]); - const bool is_overriden = override_it != override_index.end(); - - NVF_ERROR( - is_overriden || index_map.find(alloc_dom[i]) != index_map.end(), - "Couldn't find allocation mapping for ", - producer_tv->toString(), - " dim: ", - i, - " id: ", - alloc_dom[i]->toString()); - - auto alloc_ind_i = - is_overriden ? override_it->second : index_map.at(alloc_dom[i]); - - if (alloc_ind_i->isZeroInt()) { - continue; - } - - // Compute striding for this index. - Val* stride = nullptr; - for (const auto j : arange(i + 1, alloc_dom.size())) { - if (skip_indexing.count(alloc_dom[j])) { - continue; - } - - auto alloc_ext_j = (extent_map.find(alloc_dom[j]) == extent_map.end() || - is_mma_allocation(alloc_dom[j])) - ? alloc_dom[j]->extent() - : extent_map.at(alloc_dom[j]); - - alloc_ext_j = getExtentOfRootAxis(alloc_dom[j], alloc_ext_j); - - if (zero_domain_map.count(alloc_dom[j]) == 0 || - is_mma_allocation(alloc_dom[j])) { - if (stride == nullptr) { - stride = alloc_ext_j; - } else { - stride = SimplifyingIrBuilder::mulExpr(stride, alloc_ext_j); - } - } - } - - if (stride != nullptr) { - strided_inds[i] = SimplifyingIrBuilder::mulExpr(alloc_ind_i, stride); - } else { - strided_inds[i] = alloc_ind_i; - } - } - - if (producer_tv->isCircularBuffered()) { - auto db_loop = gpu_lower->circularBufferInfo().getCircularBufferLoop( - producer_tv, loops, true); - if (db_loop != nullptr) { - const int64_t stage_depth = - gpu_lower->circularBufferInfo() - .getCircularBufferOptionsFor(db_loop->iter_domain()) - .stage; - auto loop_index = db_loop->indexOrStartIfTrivial(); - auto db_switch_index = SimplifyingIrBuilder::modExpr( - loop_index, - SimplifyingIrBuilder::create(stage_depth, DataType::Index)); - auto original_alloc_size = - gpu_lower->circularBufferInfo().getOriginalAllocSize(producer_tv); - auto db_strided_index = - SimplifyingIrBuilder::mulExpr(db_switch_index, original_alloc_size); - strided_inds.push_back(db_strided_index); - } - } - - return strided_inds; -} - Val* Index::getLinearLogicalIndex( TensorView* consumer_tv, const std::vector& loops) { @@ -1456,20 +1192,16 @@ std::vector Index::getProducerPerDimLogicalIndex( const TensorView* consumer_tv, const std::vector& loops, const std::unordered_map& override_index) { - if (!ir_utils::hasRootToLoopLinearTransformations(producer_tv) || - GpuLower::current()->idModelOptions().isTensorIndexerEnabled() || - GpuLower::current()->tmemInfo().hasTMemTensor()) { - const TensorIndexer& indexer = GpuLower::current()->tensorIndexer(); - return indexer.getIndexFor( - consumer_tv->definition(), - /*as_consumer=*/false, - producer_tv->getLogicalDomain(), - loops); - } else { - auto guard = ir_utils::allocateToLogicalDomainGuard(producer_tv, false); - return getProducerAllocationIndices( - producer_tv, consumer_tv, loops, override_index); - } + NVF_ERROR( + GpuLower::current()->idModelOptions().isTensorIndexerEnabled(), + "Legacy indexer no longer available"); + + const TensorIndexer& indexer = GpuLower::current()->tensorIndexer(); + return indexer.getIndexFor( + consumer_tv->definition(), + /*as_consumer=*/false, + producer_tv->getLogicalDomain(), + loops); } std::vector Index::getStrides(TensorView* tv) { @@ -1525,384 +1257,6 @@ std::vector Index::getStrides(TensorView* tv) { return strides; } -std::vector Index::getConsumerAllocationIndices( - const TensorView* tv, - const std::vector& loops, - const IndexFromIdGraph& index_from_id_graph) { - const auto& alloc_dom = tv->getMaybeAllocationDomain(); - auto indexing = index_from_id_graph.index; - - std::vector alloc_inds( - alloc_dom.size(), GpuLower::current()->kernel()->zeroVal()); - for (const auto i : arange(alloc_dom.size())) { - // See a comment in indexing to allocation domains in - // getGlobalProducerIndex. - if (alloc_dom[i]->isReduction() || alloc_dom[i]->isBroadcast() || - alloc_dom[i]->isStride()) { - continue; - } - - NVF_ERROR( - indexing.indexMap().find(alloc_dom[i]) != indexing.indexMap().end(), - "Couldn't find allocation mapping for ", - tv->toString(), - " dim: ", - i, - " id: ", - alloc_dom[i]->toString()); - - auto alloc_ind = indexing.indexMap().at(alloc_dom[i]); - - alloc_inds[i] = alloc_ind; - } - return alloc_inds; -} - -std::vector Index::getProducerAllocationIndices( - TensorView* producer_tv, - const TensorView* consumer_tv, - const std::vector& loops, - const std::unordered_map& override_index) { - FUSER_PERF_SCOPE("GpuLower::Lower::getProducerAllocationIndices"); - // Replay producer to look like consumer so we can index on producer since - // our loop nests look like consumer - auto pairwise_map = - PairwiseLogicalDomainMap(producer_tv, consumer_tv).mapBroadcast(true); - - TensorDomain* producerAsC = TransformReplay::replayPasC( - producer_tv, - consumer_tv, - -1, - pairwise_map, - TransformReplayOptions().replayResize()) - .first; - - // Make the producer_tv look like consumer while performing indexing math - ir_utils::TVDomainGuard domain_guard(producer_tv, producerAsC); - - // Map sent to best effort replay needs to match the exact incantation for - // compute_at_mode.cpp with MappingMode::Index - auto c2p_logical_map = PairwiseLogicalDomainMap(producer_tv, consumer_tv) - .mapBroadcast(false) - .mapConsumerToProducer(); - - // This replay has to be consistent with compute at index map. - BestEffortReplay replay_producer_as_consumer( - producer_tv->getLoopDomain(), - consumer_tv->getLoopDomain(), - c2p_logical_map); - - auto c2p_map = replay_producer_as_consumer.getReplay(); - - // Make sure at least root domains are mapped even when extents may - // be different. This mapping is important for the indexing lookup - // tensors of PyTorch gather as a producer. The IDs of a lookup - // tensor may have larger extents than those of the corresponding - // output tensor, but the index expressions to those output IDs can - // still be used for the producer. Note that we always do not map - // the indirectly accessed ID and its corresponding output ID. The - // above relaxed mapping is only for the rest of the IDs. - // - // Note that when the consumer has swizzle, the swizzle are skipped. For - // example, if we have: - // consumer: - // root: I0, I1, I2 - // loop: I0, I3, I4 - // producer: - // root I5, I6, I7 - // where I3, I4 = swizzle(I1, I2) , then the c2p map will be I3->I6, I4->I7, - // I1 and I2 are not mapped. For this case, we should allow the root unmapped, - // If we add I1->I6 and I2->I7, the c2p map will no longer be injective, which - // is not what we want. - const auto p2c_map = invertOneToOneMap(c2p_map); - for (const auto& kv : PairwiseLogicalDomainMap(producer_tv, consumer_tv) - .mapBroadcast(false) - .mapDifferentExtents(true) - .mapConsumerToProducer()) { - auto consumer_root_id = kv.first; - auto producer_root_id = kv.second; - if (c2p_map.find(consumer_root_id) == c2p_map.end() && - p2c_map.find(producer_root_id) == p2c_map.end()) { - c2p_map.emplace(consumer_root_id, producer_root_id); - } - } - - const auto& producer_indexing_from_idgraph = - getTensorIndexFromIdGraph(loops, consumer_tv, producer_tv, true, c2p_map); - - auto producer_indexing = producer_indexing_from_idgraph.index; - - // Indices should now be mapped onto IterDomains in producer, so just grab - // and use them. - const auto& alloc_dom = producer_tv->getMaybeAllocationDomain(); - - std::vector alloc_inds( - alloc_dom.size(), GpuLower::current()->kernel()->zeroVal()); - - for (const auto i : arange(alloc_dom.size())) { - auto override_it = override_index.find(alloc_dom[i]); - const bool is_overriden = override_it != override_index.end(); - - if (alloc_dom[i]->isReduction() || - (alloc_dom[i]->isBroadcast() && !is_overriden)) { - continue; - } - - Val* alloc_ind = nullptr; - if (is_overriden) { - alloc_ind = override_it->second; - } else if ( - producer_indexing.indexMap().find(alloc_dom[i]) != - producer_indexing.indexMap().end()) { - alloc_ind = producer_indexing.indexMap().at(alloc_dom[i]); - } - - NVF_ERROR( - alloc_ind != nullptr, - "Couldn't find allocation mapping for ", - producer_tv->toString(), - " dim: ", - i, - " id: ", - alloc_dom[i]->toString()); - - alloc_inds.at(i) = alloc_ind; - } - - return alloc_inds; -} - -std::vector Index::getGlobalConsumerStridedIndices( - TensorView* consumer_tv, - const std::vector& loops) { - FUSER_PERF_SCOPE("GpuLower::Lower::getGlobalConsumerIndex"); - - auto index_from_id_graph = getTensorIndexFromIdGraph(loops, consumer_tv); - auto consumer_indexing = index_from_id_graph.index; - auto strides = getStrides(consumer_tv); - // if we need to override index, we need to generate the index from each - // allocation axis firstly. - auto alloc_inds = - getConsumerAllocationIndices(consumer_tv, loops, index_from_id_graph); - - // Global striding - auto vectorize_shift = - loops.empty() ? nullptr : loops.back()->vectorize_shift(); - std::vector strided_inds( - alloc_inds.size(), GpuLower::current()->kernel()->zeroVal()); - for (const auto i : arange(alloc_inds.size())) { - if (alloc_inds[i]->isZeroInt()) { - continue; - } else { - auto strided_ind = - SimplifyingIrBuilder::mulExpr(alloc_inds[i], strides[i]); - if (i == strides.size() - 1 && vectorize_shift != nullptr) { - strided_inds[i] = - SimplifyingIrBuilder::addExpr(strided_ind, vectorize_shift); - } else { - strided_inds[i] = strided_ind; - } - } - } - - NVF_ERROR( - strided_inds.size() == consumer_tv->getMaybeAllocationDomain().size()); - - return strided_inds; -} - -// Consumer index for either shared or local memory -std::vector Index::getNonGlobalConsumerStridedIndices( - const TensorView* consumer_tv, - const std::vector& loops, - const std::unordered_map& override_index) { - const auto gpu_lower = GpuLower::current(); - // At now, only ScatterOp set override_index, and the output of ScatterOp - // is on global memory, so in this method, the override_index must be empty. - NVF_ERROR(override_index.empty()); - auto consumer_indexing_from_idgraph = getTensorIndexFromIdGraph( - loops, - consumer_tv, - // Producer tv - nullptr, - // Index global - false); - - auto consumer_indexing = consumer_indexing_from_idgraph.index; - - const auto& index_map = consumer_indexing.indexMap(); - const auto& extent_map = consumer_indexing.extentMap(); - const auto& zero_domain_map = consumer_indexing.zeroDomains(); - - // Indices should now be mapped onto IterDomains in consumer, so just grab - // and use them. - const auto& alloc_dom = consumer_tv->getMaybeAllocationDomain(); - std::vector strided_inds( - alloc_dom.size(), GpuLower::current()->kernel()->zeroVal()); - for (const auto i : arange(alloc_dom.size())) { - if (alloc_dom[i]->isReduction() || alloc_dom[i]->isBroadcast() || - alloc_dom[i]->isStride() || alloc_dom[i]->isDeviceDim() || - (alloc_dom[i]->isThread() && - consumer_tv->getMemoryType() == MemoryType::Local)) { - continue; - } - - std::stringstream error_msg_loops; - if (index_map.find(alloc_dom[i]) == index_map.end()) { - for (auto loop : loops) { - error_msg_loops << " " << loop->iter_domain()->toString(); - } - } - - NVF_ERROR( - index_map.find(alloc_dom[i]) != index_map.end(), - "Couldn't find allocation mapping for ", - consumer_tv->toString(), - " dim: ", - i, - " id: ", - alloc_dom[i]->toString(), - ", loops: ", - error_msg_loops.str()); - - auto alloc_ind_i = index_map.at(alloc_dom[i]); - if (alloc_ind_i->isZeroInt()) { - continue; - } - - // Compute striding for this index. - Val* stride = nullptr; - for (const auto j : arange(i + 1, alloc_dom.size())) { - if (alloc_dom[j]->isBroadcast() || alloc_dom[j]->isReduction() || - alloc_dom[j]->isDeviceDim() || alloc_dom[j]->isStride()) { - continue; - } - - NVF_ERROR( - index_map.find(alloc_dom[j]) != index_map.end(), - "Couldn't find allocation mapping for ", - consumer_tv->toString(), - " dim: ", - j, - " id: ", - alloc_dom[j]->toString()); - - auto alloc_ext_j = extent_map.find(alloc_dom[j]) == extent_map.end() - ? alloc_dom[j]->extent() - : extent_map.at(alloc_dom[j]); - - alloc_ext_j = getExtentOfRootAxis(alloc_dom[j], alloc_ext_j); - - if (zero_domain_map.count(alloc_dom[j]) == 0) { - if (stride == nullptr) { - stride = alloc_ext_j; - } else { - stride = SimplifyingIrBuilder::mulExpr(stride, alloc_ext_j); - } - } - } - - if (stride != nullptr) { - strided_inds[i] = SimplifyingIrBuilder::mulExpr(alloc_ind_i, stride); - } else { - strided_inds[i] = alloc_ind_i; - } - } - - // This check was originally done in getConsumerStridedIndices, but - // the number of strided index values depends on the loop where the - // consumer tensor is located. If it's circular buffered and not in - // the prologue loop, strided_inds ends up having one more - // index, so it's just much simpler to check here before adding the - // additional index for circular buffering. - NVF_ERROR( - strided_inds.size() == consumer_tv->getMaybeAllocationDomain().size()); - - if (consumer_tv->isCircularBuffered()) { - auto db_loop = gpu_lower->circularBufferInfo().getCircularBufferLoop( - consumer_tv, loops); - const auto& opt = - gpu_lower->circularBufferInfo().getCircularBufferOptionsFor( - db_loop->iter_domain()); - bool is_circular_buffer_loop = opt.stage > 2; - bool is_prolog = - db_loop->circularBufferLoopStage() == CircularBufferLoopStage::Prolog; - - Val* db_switch_index = nullptr; - - // In circular buffered we don't materialize the prolog loop as there will - // be only one iteration. In circular buffer case we materialize the - // prolog loop as well covering the first N-1 iterations, N being the - // stage depth. - if (!is_prolog || is_circular_buffer_loop) { - if (is_prolog && is_circular_buffer_loop) { - // The buffer switching logic is the same as original index - // in the case of circular buffer prolog. - db_switch_index = db_loop->indexOrStartIfTrivial(); - } else { - auto loop_index = db_loop->indexOrStartIfTrivial(); - // Switching index generated for main loop or epilog component. - db_switch_index = SimplifyingIrBuilder::modExpr( - SimplifyingIrBuilder::addExpr( - loop_index, - SimplifyingIrBuilder::create( - opt.prefetch, DataType::Index)), - SimplifyingIrBuilder::create(opt.stage, DataType::Index)); - } - - // Use the generated switching buffer index to access the buffer space. - auto original_alloc_size = - gpu_lower->circularBufferInfo().getOriginalAllocSize(consumer_tv); - auto db_strided_index = - SimplifyingIrBuilder::mulExpr(db_switch_index, original_alloc_size); - strided_inds.push_back(db_strided_index); - } - } - return strided_inds; -} - -Val* Index::getProducerStridedIndices( - TensorView* producer, - const TensorView* consumer, - const std::vector& loops, - const std::unordered_map& override_index, - bool generate_pointer) { - FUSER_PERF_SCOPE("GpuLower::Lower::Index::getProducerStridedIndices"); - if (std::ranges::empty( - producer->getLoopDomain() | TensorDomain::kNoReductions)) { - if (generate_pointer) { - return IrBuilder::baseAddressExpr(producer); - } else { - return GpuLower::current()->kernel()->zeroVal(); - } - } - - if (producer->getMemoryType() == MemoryType::Global) { - auto index = sumVals(getGlobalProducerStridedIndices( - producer, consumer, loops, override_index)); - if (generate_pointer) { - return SimplifyingIrBuilder::addExpr( - IrBuilder::baseAddressExpr(producer), index); - } else { - return index; - } - } else { - auto index = sumVals(getNonGlobalProducerStridedIndices( - producer, consumer, loops, override_index)); - if (generate_pointer) { - auto index_bytes = IrBuilder::mulExpr( - index, - IrBuilder::create( - dataTypeSizeByte(*producer->getDataType()), - *index->getDataType())); - return IrBuilder::addExpr( - IrBuilder::baseAddressExpr(producer), index_bytes); - } else { - return index; - } - } -} - // Producer is the inputs of an expression kir::TensorIndex* Index::getProducerIndex( TensorView* producer, @@ -1959,44 +1313,6 @@ kir::TensorIndex* Index::getProducerIndex( return IrBuilder::create(producer, index, as_type); } -Val* Index::getConsumerStridedIndices( - TensorView* consumer, - const std::vector& loops, - bool generate_pointer) { - FUSER_PERF_SCOPE("GpuLower::Lower::Index::getConsumerStridedIndices"); - if (std::ranges::empty( - consumer->getLoopDomain() | TensorDomain::kNoReductions)) { - if (generate_pointer) { - return IrBuilder::baseAddressExpr(consumer); - } else { - return GpuLower::current()->kernel()->zeroVal(); - } - } - - if (consumer->getMemoryType() == MemoryType::Global) { - auto index = sumVals(getGlobalConsumerStridedIndices(consumer, loops)); - if (generate_pointer) { - return SimplifyingIrBuilder::addExpr( - IrBuilder::baseAddressExpr(consumer), index); - } else { - return index; - } - } else { - auto index = sumVals(getNonGlobalConsumerStridedIndices(consumer, loops)); - if (generate_pointer) { - auto index_bytes = IrBuilder::mulExpr( - index, - IrBuilder::create( - dataTypeSizeByte(*consumer->getDataType()), - *index->getDataType())); - return IrBuilder::addExpr( - IrBuilder::baseAddressExpr(consumer), index_bytes); - } else { - return index; - } - } -} - // Consumer is the output of an expression kir::TensorIndex* Index::getConsumerIndex( TensorView* consumer, diff --git a/csrc/index_compute.h b/csrc/index_compute.h index 8ee430bd373..3d37a7dc592 100644 --- a/csrc/index_compute.h +++ b/csrc/index_compute.h @@ -398,48 +398,11 @@ class PredicateInfo { // can make the below tensorviews const. class Index { private: - // Producer indexing if it's in shared or local memory - static std::vector getNonGlobalProducerStridedIndices( - TensorView* producer, - const TensorView* consumer, - const std::vector& loops, - const std::unordered_map& override_index = {}); - - // Consumer indexing if it's in shared or local memory - static std::vector getNonGlobalConsumerStridedIndices( - const TensorView* consumer, - const std::vector& loops, - const std::unordered_map& override_index = {}); - // get the strides of a tensor used for the index lowering + // Delete? static std::vector getStrides(TensorView* tv); - // get the allocation indices of a consumer tensor - static std::vector getConsumerAllocationIndices( - const TensorView* tv, - const std::vector& loops, - const IndexFromIdGraph& index_from_id_graph); - - // get the allocation indices of a producer tensor - static std::vector getProducerAllocationIndices( - TensorView* producer, - const TensorView* consumer, - const std::vector& loops, - const std::unordered_map& override_index = {}); - public: - // Producer if it's in global memory - static std::vector getGlobalProducerStridedIndices( - TensorView* producer, - const TensorView* consumer, - const std::vector& loops, - const std::unordered_map& override_index = {}); - - // Consumer indexing if it's in global memory - static std::vector getGlobalConsumerStridedIndices( - TensorView* consumer, - const std::vector& loops); - // Indexing functions // Consumer = Producer // i.e. T0 = T1... -> T0 is the consumer, T1 is the producer @@ -469,27 +432,6 @@ class Index { DataType as_type = DataType::Null, bool ld_st_matrix = false); - //! Returns a vector of strided indices mapped onto the - //! allocation domain of a producer tensor. The size of the returned - //! vector is guaranteed to be equal to the number of axes of the - //! indexing allocation domain. - static Val* getProducerStridedIndices( - TensorView* producer, - const TensorView* consumer, - const std::vector& loops, - const std::unordered_map& override_index = {}, - bool generate_pointer = false); - - // TODO: Remove - //! Returns a vector of strided indices mapped onto the - //! allocation domain of a consumer tensor. The size of the returned - //! vector is guaranteed to be equal to the number of axes of the - //! indexing allocation domain. - static Val* getConsumerStridedIndices( - TensorView* consumer, - const std::vector& loops, - bool generate_pointer = false); - //! Returns the logical index linearized from a multi-dimension address into a //! linear memory address a consumer tensor. The returned index is intended to //! be used for the computation of some tensor factories, such as: iota and From 7709bc2af258918f9cc69c15def9efece1221174 Mon Sep 17 00:00:00 2001 From: Naoya Maruyama Date: Mon, 6 Apr 2026 14:23:38 -0700 Subject: [PATCH 3/7] cleanup --- CMakeLists.txt | 1 - csrc/device_lower/analysis/index_compute.cpp | 1407 ---------------- csrc/device_lower/analysis/index_compute.h | 320 ---- .../analysis/sync_information.cpp | 1 - csrc/device_lower/pass/index.cpp | 1 - csrc/device_lower/pass/magic_zero.cpp | 124 -- csrc/device_lower/pass/magic_zero.h | 36 - csrc/device_lower/pass/rng.cpp | 1 - csrc/id_model/indexing.cpp | 2 - csrc/id_model/indexing_utils.h | 1 - csrc/index_compute.cpp | 1491 +---------------- csrc/index_compute.h | 380 ----- csrc/predicate_compute.cpp | 23 +- 13 files changed, 7 insertions(+), 3781 deletions(-) delete mode 100644 csrc/device_lower/analysis/index_compute.cpp delete mode 100644 csrc/device_lower/analysis/index_compute.h diff --git a/CMakeLists.txt b/CMakeLists.txt index 24f5bd2d59c..5b9241d253c 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -149,7 +149,6 @@ list(APPEND NVFUSER_SRCS ${NVFUSER_SRCS_DIR}/device_lower/analysis/divisible_split.cpp ${NVFUSER_SRCS_DIR}/device_lower/analysis/fused_reduction.cpp ${NVFUSER_SRCS_DIR}/device_lower/analysis/fusion_info.cpp - ${NVFUSER_SRCS_DIR}/device_lower/analysis/index_compute.cpp ${NVFUSER_SRCS_DIR}/device_lower/analysis/non_divisible_split.cpp ${NVFUSER_SRCS_DIR}/device_lower/analysis/padded_parallel_dimensions.cpp ${NVFUSER_SRCS_DIR}/device_lower/analysis/predicate_elimination.cpp diff --git a/csrc/device_lower/analysis/index_compute.cpp b/csrc/device_lower/analysis/index_compute.cpp deleted file mode 100644 index 859975030a4..00000000000 --- a/csrc/device_lower/analysis/index_compute.cpp +++ /dev/null @@ -1,1407 +0,0 @@ -// 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 -#include -#include -#include -#include -#include -#include -#include -#include - -namespace nvfuser { - -IndexFromIdGraph::IndexFromIdGraph( - IndexCompute index_, - IndexCompute concrete_index_, - std::unordered_map initial_concrete_index_map_, - std::vector loop_domains_) - : index(std::move(index_)), - concrete_index(std::move(concrete_index_)), - initial_concrete_index_map(std::move(initial_concrete_index_map_)), - resolved_loop_domains(std::move(loop_domains_)) {} - -namespace { - -// Maps all producer domains to consumer with broadcast -// forwarding. Used to find the allocation position. -// TODO: should this be an ir_util ? Didn't seem to be -// used too much though. -std::unordered_map mapAllProducerDomainsToConsumer( - const TensorView* producer_tv, - const TensorView* consumer_tv) { - // This map has forwarded broadcast axes, it should only be used to compute - // the allocation position of the producer - std::unordered_map p2c_alloc_map; - - // We want to replay producer as consumer instead of the other way around - // since consumer may have some broadcasted axes producer doesn't have - // merged into loops producer may use. If we did consumer as producer we - // wouldn't have this information in the mapping. - auto replay_PasC = BestEffortReplay::replayPasC( - producer_tv, - consumer_tv, - -1, - PairwiseLogicalDomainMap(producer_tv, consumer_tv)); - - // Grab consumer domain entries and reverse replay map. TODO: Maybe - // TransformReplay::replayPasC could return this map - for (auto id : consumer_tv->getLoopDomain()) { - const auto& c2p_map = replay_PasC.getReplay(); - auto c2p_it = c2p_map.find(id); - if (c2p_it != c2p_map.end()) { - auto c_id = c2p_it->first; - auto p_id = c2p_it->second; - p2c_alloc_map[p_id] = c_id; - } - } - - return p2c_alloc_map; -} - -std::unordered_map invertOneToOneMap( - const std::unordered_map& map) { - std::unordered_map inverted; - for (const auto& kv : map) { - bool inserted = inverted.emplace(kv.second, kv.first).second; - NVF_ERROR( - inserted, - "Multiple mappings to the same value detected: ", - kv.second->toString()); - } - return inverted; -} - -//! A struct to keep track of necessary parameters used in -//! configuring index compute pass. -//! These parameters are needed to propagate the indexing from the loop nodes of -//! the TVs and loop nests to the TVs logical domain during -//! index_compute.cpp::IndexCompute passes. -//! TODO: -//! Would expect this list to become shorter over time, -//! as more info can be determined holistically. -struct IndexingParameters { - //! Initial binding of index math to concrete iterdomain ids, - //! from the loop nest analysis. - std::unordered_map initial_concrete_id_index; - - //! (Used in non-global indexing) the concrete iterdomains that - //! we want to skip or merge into contiguous indexing paths. - std::unordered_set zero_domains; - - //! (Used in non-global indexing) the preferred path we would - //! be propagating contiguously merged indices backward. - std::unordered_set preferred_concrete_ids; - - //! Unswitched concrete domains. Back-traversing through the inner - //! domain of a merge may need to be replaced with the maximum of - //! the inner domain. - std::unordered_set unswitched_domains; -}; - -// Initial loop index map for global producer or consumer case. -IndexingParameters getLinearIndexParameters( - const LoopIndexing& loop_indexing, - bool index_producer = false) { - IndexingParameters index_parameters; - - auto& loops = loop_indexing.loops(); - auto& loop_domain = loop_indexing.loopDomains(); - auto& loop_index_map = index_parameters.initial_concrete_id_index; - - for (auto [loop_idx, loop] : enumerate(loops)) { - IterDomain* index_domain = - GpuLower::current()->info().caMap().getConcreteMappedID( - loop_domain[loop_idx], IdMappingMode::EXACT); - loop_index_map[index_domain] = loop->indexOrStartIfTrivial(); - } - - protectNonPredicateIndexWithMagicZero( - loops, - loop_indexing.loopDomains(), - index_parameters.initial_concrete_id_index); - - // Setup circular buffer increment for producer case: - // TODO: could unify these circular buffer index calculation - // in follow ups. - if (index_producer) { - auto circular_buffer_loop = - GpuLower::current()->circularBufferInfo().getCircularBufferLoop( - loop_indexing.consumerTv(), loops, true); - - for (auto loop_idx : arange(loops.size())) { - auto loop = loops[loop_idx]; - if (loop == circular_buffer_loop) { - auto loop_id = loop_indexing.loopDomains()[loop_idx]; - - auto concrete_loop_id = - GpuLower::current()->info().caMap().getConcreteMappedID( - loop_id, IdMappingMode::EXACT); - - auto prefetch_distance = - GpuLower::current() - ->circularBufferInfo() - .getCircularBufferOptionsFor(loop->iter_domain()) - .prefetch; - index_parameters.initial_concrete_id_index[concrete_loop_id] = - SimplifyingIrBuilder::addExpr( - index_parameters.initial_concrete_id_index[concrete_loop_id], - SimplifyingIrBuilder::create( - prefetch_distance, DataType::Index)); - } - } - } - - return index_parameters; -} - -// Initial index parameters for shared and local case -IndexingParameters getNonGlobalInitialIndexParameters( - const LoopIndexing& loop_indexing, - const TensorView* consumer_tv, - bool index_producer = false, - const TensorView* producer_tv = nullptr, - std::unordered_map p2c_map = {}) { - IndexingParameters index_parameters; - const auto& loops = loop_indexing.loops(); - const auto& loop_domains = loop_indexing.loopDomains(); - - // TODO: - // The non-global path should become shorter as we - // pull more info into id graph. - std::unordered_map alloc_id_map; - - if (index_producer) { - alloc_id_map = mapAllProducerDomainsToConsumer(producer_tv, consumer_tv); - } - - auto alloc_tv = index_producer ? producer_tv : consumer_tv; - auto alloc_info = lower_utils::getAllocPosInfo( - alloc_tv, loops, alloc_id_map, index_producer); - - std::unordered_map loop_to_ind_map; - std::unordered_set zero_loops; - - kir::ForLoop* circular_buffer_loop = nullptr; - - if (index_producer) { - circular_buffer_loop = - GpuLower::current()->circularBufferInfo().getCircularBufferLoop( - consumer_tv, loops, true); - } - - std::tie(loop_to_ind_map, zero_loops) = indexMapFromTV( - alloc_tv, - loops, - alloc_info.init_for_loop, - !index_producer, - circular_buffer_loop); - - ensureStaticIndexing(alloc_tv, alloc_info.init_for_loop, loops, alloc_id_map); - - NVF_ERROR( - loops.size() <= loop_domains.size(), - "Loop domain didn't replay all loops"); - - for (auto loop_idx : arange(loops.size())) { - auto loop = loops[loop_idx]; - auto loop_domain = loop_domains[loop_idx]; - - auto concrete_loop_domain = - GpuLower::current()->info().caMap().getConcreteMappedID( - loop_domain, IdMappingMode::EXACT); - - index_parameters.initial_concrete_id_index[concrete_loop_domain] = - loop_to_ind_map.at(loop); - - if (zero_loops.count(loop)) { - index_parameters.zero_domains.insert(concrete_loop_domain); - } - } - - // Derive preferred path from loop indexing result. - const TensorView* target_tv = index_producer ? producer_tv : consumer_tv; - index_parameters.preferred_concrete_ids = buildLoopIndexingPreferredPath( - target_tv, loop_indexing, index_producer, p2c_map); - - return index_parameters; -} - -// Check if this loop is actually unswitched, meaning an initial index -// of the maximum value from a non-size-one range is used. -bool trackUnswitchedDomain(kir::ForLoop* loop) { - // Loop index has only one valid value per thread, which means the - // loop is not actually unswitched - if (loop->isTrivial()) { - return false; - } - - // The same can be said as long as it's exactly mapped with a - // vectorized domain - const auto& id_exact_set = GpuLower::current() - ->info() - .caMap() - .getIdSets(IdMappingMode::EXACT) - .getDisjointSetOf(loop->iter_domain()); - - if (std::any_of(id_exact_set.begin(), id_exact_set.end(), [](auto id) { - return id->getParallelType() == ParallelType::Vectorize; - })) { - return false; - } - - return true; -} - -//! Initial index parameters for predicate, adjusts loop to indexing -//! may according to the information annotated on the loop nest. -//! -//! TODO: -//! This function is mostly copy pasted from previous implementation -//! at this step, further clean up is possible since: -//! 1. Much of the loop-to-ind adjustment will be issued from idgraph -//! 2. Much of the initial index logic could be shared across all -//! the 3 variants. -IndexingParameters getPredicateInitialIndexParameters( - const LoopIndexing& loop_indexing, - TensorView* consumer_tv, - kir::ForLoop* unswitch_or_vec_loop, - IterDomain* circular_buffer_axis, - bool is_start_predicate) { - IndexingParameters index_parameters; - const auto& loops = loop_indexing.loops(); - const auto& loop_domains = loop_indexing.loopDomains(); - - // This shouldn't be needed. - NVF_ERROR( - loops.size() <= loop_domains.size(), - "Loop domain didn't replay all loops"); - - std::unordered_map loop_to_ind_map; - - // Fill initial index with each forloop's index. - for (auto fl : loops) { - if (fl->isTrivial()) { - loop_to_ind_map[fl] = fl->start(); - } else { - loop_to_ind_map[fl] = fl->index(); - } - } - - bool unswitch_pred = unswitch_or_vec_loop != nullptr && - (unswitch_or_vec_loop->iter_domain()->getParallelType() == - ParallelType::Unswitch || - unswitch_or_vec_loop->iter_domain()->getParallelType() == - ParallelType::Unroll); - - // Vectorized predicates are different from unswitch. Unswitch predicates - // all loops within the unswitch (the outer most unswitch) are generated - // with loop->extent-1 as the index. With vectorized predicates, only the - // vectorized loop should be like this. - - bool within_unswitch = false; - - for (const auto loop_i : arange(loops.size())) { - auto loop = loops[loop_i]; - auto loop_id = loop->iter_domain(); - auto loop_pt = loop_id->getParallelType(); - auto ref_id = loop_domains.at(loop_i); - - if (!within_unswitch && unswitch_pred) { - within_unswitch = loop == unswitch_or_vec_loop; - } - - bool predicate_at_end = within_unswitch || loop == unswitch_or_vec_loop || - lower_utils::predicateAtEnd(loop); - - if (predicate_at_end) { - // Rely on the reference to check broadcasting. The for loop could be - // broadcasted on a constant value from an unroll split. Since reference - // may convert this to an iter domain, that for loop could be valid to - // generate predication from. - - // Note that loop->stop() is not used below. Instead, - // loop->iter_domain()->extent() is used, which is uniform - // across the mapped domains irrespective of halo. Predicates are - // compared with each to pick the most restrictive ones. The - // comparison is done by only using the offset, which is the - // term added to the index. So, the index term must be the - // same among all predicates, otherwise the comparison would - // be invalid. The effect by halo is added to the offset - // term. See getUnswitchStopOffset. - - if (ref_id->isBroadcast()) { - // Ignore indexing into broadcasted dimensions. - continue; - } else if (loop_id->isThread()) { - // When parallelized, if the loop stop is the same as the - // extent of the associated IterDomain, i.e., no extra - // iterations for halo, predicating with the threading index - // is sufficient for both the start and stop - // predicates. That isn't the case if the loop has halo, and - // in the case either the minimum and maximum values of the - // iteration domain needs to be used. - // - // Note: Better performance was obtained if using - // threadIdx in unswitch predicates was avoided. More - // specifically, in the Hdiff stencil example, instead of - // predicating with threadIdx.x for both the start and stop - // predicates, using zero and (blockDim.x - 1) for the start - // and stop predicates, respectively, resulted in less - // register pressure. The alternative codegen can be done by - // adding this to the first if condition: - // loop_id->isBlockDim(). This would not be a concern if the - // else part could be omitted, so canOmitElseClause should - // be used as well. - if (loop->stop() == loop_id->extent()) { - loop_to_ind_map[loop] = loop->start(); - } else if (is_start_predicate) { - loop_to_ind_map[loop] = GpuLower::current()->kernel()->zeroVal(); - } else { - // Note that the parallel dimension is used rather than - // loop-stop(). See the above comment. - loop_to_ind_map[loop] = - GpuLower::current()->info().parallelDimensionMap().get(loop_pt); - } - } else if (is_start_predicate) { - loop_to_ind_map[loop] = GpuLower::current()->kernel()->zeroVal(); - } else { - // Similar to the above, loop_id()->extent() is - // used here instead of loop->stop(). See the above comment. - loop_to_ind_map[loop] = SimplifyingIrBuilder::subExpr( - loop_id->extent(), GpuLower::current()->kernel()->oneVal()); - } - - // When predicating a loop at the maximum end, predicate - // expressions such as (extent-1) are used, which represent the - // maximum possible value of the loop range but are not - // guaranteed to result in the maximum index when traversing - // through merge inner domains as modulo is used. Keep track of - // those domains, which will be used by IndexCompute to make - // necessary adjustments. See also csrc/index_compute.h for more - // context. - if (!is_start_predicate && trackUnswitchedDomain(loop)) { - index_parameters.unswitched_domains.insert( - GpuLower::current()->info().caMap().getConcreteMappedID( - loop_id, IdMappingMode::EXACT)); - } - } - } - - // Increment circular buffer loop index - if (circular_buffer_axis != nullptr) { - auto db_loop = - GpuLower::current()->circularBufferInfo().getCircularBufferLoop( - circular_buffer_axis, loops, true); - if (db_loop != nullptr) { - auto loop_to_ind_map_it = loop_to_ind_map.find(db_loop); - NVF_ERROR(loop_to_ind_map_it != loop_to_ind_map.end()); - auto cur_index = loop_to_ind_map_it->second; - // if cur_index is not the same as the index of db_loop, it must - // be true that that index has been modified to support - // unswitch. In that case, it is not necessary to move ahead the - // index for circular buffering. - auto prefetch_distance = - (int64_t)GpuLower::current() - ->circularBufferInfo() - .getCircularBufferOptionsFor(db_loop->iter_domain()) - .prefetch; - bool is_same = cur_index == db_loop->indexOrStartIfTrivial(); - if (is_same) { - loop_to_ind_map[db_loop] = SimplifyingIrBuilder::addExpr( - cur_index, - SimplifyingIrBuilder::create( - prefetch_distance, DataType::Index)); - } - } - } - - // Convert loop-to-ind map to concrete-to-ind map - for (auto loop_idx : arange(loops.size())) { - auto loop = loops.at(loop_idx); - auto concrete_loop_domain = - GpuLower::current()->info().caMap().getConcreteMappedID( - loop_domains.at(loop_idx), IdMappingMode::EXACT); - index_parameters.initial_concrete_id_index[concrete_loop_domain] = - loop_to_ind_map.at(loop); - } - - // Note that, unlike non-predicate indexing, magic-zero insertion is - // not done at this point but is done individually for each indexed - // domain. See Index::getReferenceRootPredicates. - - return index_parameters; -} - -} // namespace - -LoopIndexing LoopIndexingAnalysis::fromLoopAndConsumer( - const std::vector& loops, - const TensorView* consumer_tv) { - LoopIndexingAnalysis analysis(loops, consumer_tv); - return analysis.getLoopIndexing(loops); -} - -VectorOfUniqueEntries LoopIndexingAnalysis:: - getReplayableConcreteIDs( - const std::vector& consumer_loop_ids, - const TensorView* consumer_tv) { - LoopIndexingAnalysis analysis(consumer_loop_ids, consumer_tv); - return analysis.replayed_concrete_ids_; -} - -LoopIndexingAnalysis::LoopIndexingAnalysis( - const std::vector& loops, - const TensorView* consumer_tv) - : consumer_tv_(consumer_tv) { - // Validate consistency in given loop nest - validateLoopStructure(loops); - - // Populate initial loop iter domains. - std::transform( - loops.begin(), - loops.end(), - std::back_inserter(initial_loop_domain_ids_), - [](kir::ForLoop* fl) { return fl->iter_domain(); }); - - run(); -} - -LoopIndexingAnalysis::LoopIndexingAnalysis( - const std::vector& consumer_loop_ids, - const TensorView* consumer_tv) - : consumer_tv_(consumer_tv) { - // Populate initial loop iter domains. - std::transform( - consumer_loop_ids.begin(), - consumer_loop_ids.end(), - std::back_inserter(initial_loop_domain_ids_), - [&](IterDomain* consumer_loop_id) { - // Make sure consumer_loop_id is indeed a consumer loop ID - NVF_ERROR( - std::find( - consumer_tv->getLoopDomain().begin(), - consumer_tv->getLoopDomain().end(), - consumer_loop_id) != consumer_tv->getLoopDomain().end(), - "Not a consumer loop ID: ", - consumer_loop_id->toString(), - ", consumer: ", - consumer_tv->toString()); - return GpuLower::current()->info().caMap().getConcreteMappedID( - consumer_loop_id, IdMappingMode::LOOP); - }); - - run(); -} - -void LoopIndexingAnalysis::run() { - // Collect consumer id's for view rfactor traversal. - all_consumer_id_vals_ = DependencyCheck::getAllValsBetween( - {consumer_tv_->getMaybeRootDomain().begin(), - consumer_tv_->getMaybeRootDomain().end()}, - {consumer_tv_->getLoopDomain().begin(), - consumer_tv_->getLoopDomain().end()}); - - // Resolve definition of each exact concrete id's involved in the whole loop - // nest transform history - traverseFromDomainVals(); - - // Construct concrete to consumer map. The replayed exprs are guaranteed to - // consume each concrete id once so this map is well defined. - for (auto expr : replayed_exprs_) { - for (auto input_id : ir_utils::filterByType(expr->inputs())) { - auto concrete_input_id = - GpuLower::current()->info().caMap().getConcreteMappedID( - input_id, IdMappingMode::EXACT); - concrete_id_to_consumer_[concrete_input_id] = expr; - } - } - - // Reconstruct the iterdomain view of the original loopnest after resolving - // the exact definition of each index. - constructLoopDomains(); - - //! Collect the set of indexing expressions that can be - //! resolved out of line. - collectOutOfLineExprs(); -} - -void LoopIndexingAnalysis::validateLoopStructure( - const std::vector& loops) { - // Throw an error when two loops are mapped with each other, which - // violates an assumption that unique mappings between concrete - // IterDomains and the IterDomains of the loop structure must be - // established. It should be a reasonable assumption, but fusions - // like below won't work: - // tv0 = [I0] - // tv1 = broadcast(tv0, {true, false}); - // tv2 = broadcast(tv0, {false, true}); - // tv3 = tv1 + tv2 - // Notice that the two axes of each of tv1, tv2 and tv3 are mapped - // with each other. We believe it is unlikely this limitation - // becomes a real concern in practice. - // Map concrete id to the original loop iter domain. - std::unordered_map concrete_to_loop; - for (auto for_loop : loops) { - // Largely duplicating original logic - auto loop_id = for_loop->iter_domain(); - auto concrete_loop_id = - GpuLower::current()->info().caMap().getConcreteMappedID( - loop_id, IdMappingMode::EXACT); - - NVF_ERROR( - !concrete_to_loop.count(concrete_loop_id), - "Unsupported loop structure. Two loops are mapped together.", - loop_id->toString(), - " and ", - concrete_to_loop.at(concrete_loop_id)->toString()); - - concrete_to_loop[concrete_loop_id] = loop_id; - } -} - -void LoopIndexingAnalysis::traverseFromDomainVals() { - // Order is really important here, start with outer most for loops in a - // depth first manner. The outer most loops are topologically closer to the - // outputs, so their broadcast dimensions are "more" resolved than those - // towards the inner most loops. - std::deque to_visit( - initial_loop_domain_ids_.begin(), initial_loop_domain_ids_.end()); - std::unordered_set visited_exprs; - std::unordered_set visited_ids; - - while (!to_visit.empty()) { - auto out_id = to_visit.front(); - to_visit.pop_front(); - - if (!visited_ids.emplace(out_id).second) { - continue; - } - auto expr = out_id->definition(); - - if (auto logical_id = - getLogicalIDToTraverse(out_id, all_consumer_id_vals_)) { - to_visit.emplace_front(logical_id); - } - - // ID's will be copied for the reference as we replay transformations. If - // there was no transformations on an iteration domain, a copy of the - // iteration domain for the reference is made here. - if (expr == nullptr) { - if (std::find( - initial_loop_domain_ids_.begin(), - initial_loop_domain_ids_.end(), - out_id) != initial_loop_domain_ids_.end()) { - concretizeAndVisitId(out_id); - } - continue; - } - - if (!visited_exprs.emplace(expr).second) { - continue; - } - - visitExpr(expr); - - auto inp_ids = ir_utils::filterByType(expr->inputs()); - // Make sure to put at the begining of the deque to maintain correct - // ordering. - to_visit.insert(to_visit.begin(), inp_ids.begin(), inp_ids.end()); - } -} - -IterDomain* LoopIndexingAnalysis::concretizeAndVisitId(IterDomain* id) { - auto concrete_id = GpuLower::current()->info().caMap().getConcreteMappedID( - id, IdMappingMode::EXACT); - if (replayed_concrete_ids_.pushBack(concrete_id)) { - concrete_to_original_id_[concrete_id] = id; - } - return concrete_id; -} - -namespace { -// Alias used for std::transform -IterDomain* exactConcreteId(IterDomain* id) { - return GpuLower::current()->info().caMap().getConcreteMappedID( - id, IdMappingMode::EXACT); -} -} // namespace - -void LoopIndexingAnalysis::visitExpr(Expr* expr) { - // Current implementation just tries to - // follow the exact behavior of reference replay - // except that no expr was actually "replayed". - - // Record all inputs, and stop if current expr - // duplicates id consumption or production. - if (visitIdsAndCheckDuplication(expr->inputs(), consumed_concrete_)) { - return; - } - if (visitIdsAndCheckDuplication(expr->outputs(), produced_concrete_)) { - return; - } - - // Record the expr if no duplication on input or output found - replayed_exprs_.push_back(expr); - - // Record the consumed and produced concrete ids by the newly - // recorded expression. - auto consumed_ids = ir_utils::filterByType(expr->inputs()); - std::transform( - consumed_ids.begin(), - consumed_ids.end(), - std::inserter(consumed_concrete_, consumed_concrete_.end()), - exactConcreteId); - - auto produced_ids = ir_utils::filterByType(expr->outputs()); - std::transform( - produced_ids.begin(), - produced_ids.end(), - std::inserter(produced_concrete_, produced_concrete_.end()), - exactConcreteId); -} - -bool LoopIndexingAnalysis::visitIdsAndCheckDuplication( - const std::vector& vals, - const std::unordered_set& existing_ids) { - bool duplication = false; - for (auto id : ir_utils::filterByType(vals)) { - duplication = duplication || existing_ids.count(concretizeAndVisitId(id)); - } - return duplication; -} - -void LoopIndexingAnalysis::constructLoopDomains() { - for (auto loop_id : initial_loop_domain_ids_) { - // Find the replayed_concrete_id mapping to the loop id. - auto ref_id_it = std::find_if( - replayed_concrete_ids_.vector().begin(), - replayed_concrete_ids_.vector().end(), - [&](IterDomain* concrete_id) { - return - // Make sure the replayed_concrete_id is a loop ID - !concrete_id_to_consumer_.count(concrete_id) && - // Use permissive map so the selected ID indeed represents the - // loop. - // This mapping look up is part of a staged indexing scheme. - // When we find a replayed exact id that exactly map to the loop - // id, this means that we can resolve indexing involved in this - // loop "locally", i.e. only with and with only the iterdomains - // on the given consumer tv. - // When we cannot find an exact mapping, the permissive mapping - // would help defering the indexing resolution for this loop nest - // level to other iterdomain expressions from tv's that are - // further concretized and usually they are further down the - // consumer chain of the given consumer tv. - GpuLower::current()->info().caMap().areMapped( - concrete_id, loop_id, IdMappingMode::PERMISSIVE); - }); - - NVF_ERROR( - ref_id_it != replayed_concrete_ids_.vector().end(), - "Could not find required iter domain in reference replay: ", - loop_id->toString()); - - auto ref_id = *ref_id_it; - loop_domains_.pushBack(concrete_to_original_id_.at(ref_id)); - } - - // Construct the root domain as the inputs of the replayed domain - auto loops_replayed_domain_vals = - ir_utils::filterByType(loop_domains_.vector()); - auto root_domain_vals = IterVisitor::getInputsTo( - {loops_replayed_domain_vals.begin(), loops_replayed_domain_vals.end()}); - - // Fill loop roots: - auto root_domain_ids = ir_utils::filterByType(root_domain_vals); - loop_root_domains_ = - std::vector(root_domain_ids.begin(), root_domain_ids.end()); - - // The domain may have dangling iteration domains, i.e. the inner output of - // a split but not the outer. Find which replayed vals are dependant on the - // root domains. - auto all_replayed_vals = - ir_utils::filterByType(replayed_concrete_ids_.vector()); - auto all_ids_from_root = DependencyCheck::getAllValsBetween( - {root_domain_vals.begin(), root_domain_vals.end()}, - {all_replayed_vals.begin(), all_replayed_vals.end()}); - - // Fill all dangling outputs as otherwise backwards visitor in index compute - // will complain for not having all outputs of the traversal. - for (auto id : ir_utils::filterByType(all_ids_from_root)) { - if (id->uses().empty()) { - loop_domains_.pushBack( - GpuLower::current()->info().caMap().getConcreteMappedID( - id, IdMappingMode::EXACT)); - } - } -} - -IndexFromIdGraph getTensorIndexFromIdGraph( - const std::vector& loops, - const TensorView* consumer_tv, - const TensorView* producer_tv, - bool is_global, - const std::unordered_map& c2p_map) { - bool index_producer = producer_tv != nullptr; - auto target_tv = index_producer ? producer_tv : consumer_tv; - - auto loop_indexing = - LoopIndexingAnalysis::fromLoopAndConsumer(loops, consumer_tv); - - IndexingParameters index_parameters; - - std::unordered_map p2c_map; - - // The p2c map is only needed when indexing producer - // as producer has replayed ids. - if (index_producer) { - p2c_map = invertOneToOneMap(c2p_map); - } - - if (is_global) { - index_parameters = getLinearIndexParameters(loop_indexing, index_producer); - } else { - index_parameters = getNonGlobalInitialIndexParameters( - loop_indexing, consumer_tv, index_producer, producer_tv, p2c_map); - } - - IndexCompute indexing( - index_parameters.initial_concrete_id_index, - index_parameters.zero_domains, - index_parameters.preferred_concrete_ids); - - // Run first backward traversal to generate - // loop nest based indexing math. - indexing.run(loop_indexing); - - // Populate indexing through exact map from initial indexing - auto consumer_root = index_producer ? consumer_tv->getMaybeRootDomain() - : consumer_tv->getMaybeAllocationDomain(); - - // First collect all iterdomains in consumer transform history. - auto all_consumer_vals = DependencyCheck::getAllValsBetween( - {consumer_root.begin(), consumer_root.end()}, - {consumer_tv->getLoopDomain().begin(), - consumer_tv->getLoopDomain().end()}); - - // Want update map to be based on almost exact, but indexing is on exact, make - // a map from one space to the other. - std::unordered_map> - almost_exact_2_target_ids; - - for (IterDomain* consumer_id : - ir_utils::filterByType(all_consumer_vals)) { - auto target_id = consumer_id; - - // use mapped producer id when indexing producer - if (index_producer) { - auto target_id_it = c2p_map.find(consumer_id); - if (target_id_it == c2p_map.end()) { - // consumer id not found in c2p map - // skip binding for this id. - continue; - } - target_id = target_id_it->second; - } - - auto almost_exact_concrete_id = - GpuLower::current()->info().caMap().getConcreteMappedID( - consumer_id, IdMappingMode::ALMOSTEXACT); - - auto almost_exact_2_target_ids_it = - almost_exact_2_target_ids.find(almost_exact_concrete_id); - if (almost_exact_2_target_ids_it == almost_exact_2_target_ids.end()) { - almost_exact_2_target_ids_it = - almost_exact_2_target_ids - .emplace( - almost_exact_concrete_id, - VectorOfUniqueEntries()) - .first; - } - auto& mapped_dims = almost_exact_2_target_ids_it->second; - mapped_dims.pushBack(target_id); - } - - // Map the concrete id indexing back to the producer or consumer tv - std::unordered_map> - index_update_map; - for (auto entry : indexing.indexMap()) { - auto ref_exact_id = entry.first; - auto almost_exact_concrete_id = - GpuLower::current()->info().caMap().getConcreteMappedID( - ref_exact_id, IdMappingMode::ALMOSTEXACT); - - if (almost_exact_2_target_ids.find(almost_exact_concrete_id) == - almost_exact_2_target_ids.end()) { - continue; - } - - auto consumer_ids = almost_exact_2_target_ids.at(almost_exact_concrete_id); - - for (auto consumer_id : consumer_ids) { - auto index_update_map_it = index_update_map.find(ref_exact_id); - if (index_update_map_it == index_update_map.end()) { - index_update_map_it = - index_update_map - .emplace(ref_exact_id, VectorOfUniqueEntries()) - .first; - } - auto& mapped_dims = index_update_map_it->second; - mapped_dims.pushBack(consumer_id); - } - } - - // No contig indexing was done in reference indexing - ContigIDs contig_finder( - target_tv->getLoopDomain(), - target_tv->getMaybeAllocationDomain(), - target_tv->domain()->contiguity(), - {}, - indexing.indexMap(), - GpuLower::current()->divisibleSplitSet(), - &GpuLower::current()->info().caMap(), - &GpuLower::current()->info().concretizedBroadcastDomains(), - p2c_map); - - auto target_indexing = indexing.updateIndexCompute( - target_tv->domain(), index_update_map, contig_finder); - - // Fill validation info. - // TODO: cleanup seems possible. - if (index_producer) { - fillProducerVectorizedContigAllocationDomains( - producer_tv, consumer_tv, contig_finder); - } else { - fillConsumerVectorizedContigAllocationDomains(consumer_tv, contig_finder); - } - - return IndexFromIdGraph( - target_indexing, - indexing, - index_parameters.initial_concrete_id_index, - loop_indexing.loopDomains()); -} - -IndexFromIdGraph getPredicateIndexingFromIdGraph( - const std::vector& loops, - TensorView* consumer_tv, - kir::ForLoop* unswitch_or_vec_loop, - IterDomain* circular_buffer_axis, - bool is_start_predicate) { - // Run replay pass on the loop nest to generate the deterministic - // traversal info from loop structure. - auto loop_indexing = - LoopIndexingAnalysis::fromLoopAndConsumer(loops, consumer_tv); - - // Bind initial index variables to the loop nodes and adjust - // according to loop and unswitch info. - auto index_parameters = getPredicateInitialIndexParameters( - loop_indexing, - consumer_tv, - unswitch_or_vec_loop, - circular_buffer_axis, - is_start_predicate); - - // Run first backward traversal to generate - // loop nest based indexing math. - IndexCompute indexing( - index_parameters.initial_concrete_id_index, - index_parameters.zero_domains, - index_parameters.preferred_concrete_ids, - index_parameters.unswitched_domains); - - indexing.run(loop_indexing); - - // First collect all iterdomains in consumer transform history. - auto all_consumer_vals = DependencyCheck::getAllValsBetween( - {consumer_tv->getMaybeAllocationDomain().begin(), - consumer_tv->getMaybeAllocationDomain().end()}, - {consumer_tv->getLoopDomain().begin(), - consumer_tv->getLoopDomain().end()}); - - // Want update map to be based on almost exact, but indexing is on exact, make - // a map from one space to the other. - std::unordered_map> - almost_exact_2_consumer_ids; - - for (IterDomain* consumer_id : - ir_utils::filterByType(all_consumer_vals)) { - auto almost_exact_concrete_id = - GpuLower::current()->info().caMap().getConcreteMappedID( - consumer_id, IdMappingMode::ALMOSTEXACT); - - auto almost_exact_2_consumer_ids_it = - almost_exact_2_consumer_ids.find(almost_exact_concrete_id); - if (almost_exact_2_consumer_ids_it == almost_exact_2_consumer_ids.end()) { - almost_exact_2_consumer_ids_it = - almost_exact_2_consumer_ids - .emplace( - almost_exact_concrete_id, - VectorOfUniqueEntries()) - .first; - } - auto& mapped_dims = almost_exact_2_consumer_ids_it->second; - mapped_dims.pushBack(consumer_id); - } - - // Map the concrete id indexing back to the consumer tv - std::unordered_map> - index_update_map; - for (auto entry : indexing.indexMap()) { - auto ref_exact_id = entry.first; - auto almost_exact_concrete_id = - GpuLower::current()->info().caMap().getConcreteMappedID( - ref_exact_id, IdMappingMode::ALMOSTEXACT); - - if (almost_exact_2_consumer_ids.find(almost_exact_concrete_id) == - almost_exact_2_consumer_ids.end()) { - continue; - } - auto consumer_ids = - almost_exact_2_consumer_ids.at(almost_exact_concrete_id); - - for (auto consumer_id : consumer_ids) { - auto index_update_map_it = index_update_map.find(ref_exact_id); - if (index_update_map_it == index_update_map.end()) { - index_update_map_it = - index_update_map - .emplace(ref_exact_id, VectorOfUniqueEntries()) - .first; - } - auto& mapped_dims = index_update_map_it->second; - mapped_dims.pushBack(consumer_id); - } - } - - // No contiguity info is used in the predicate indexing pass, the predicate - // generation logic that uses the index math generated here will take - // contiguity into account. Send an empty ContigID class so nothing is marked - // as contiguous. - auto contig_finder = ContigIDs::getNonContigIDs(); - - // Run second backward traversal to map back to the consumer_tv - auto target_indexing = indexing.updateIndexCompute( - consumer_tv->domain(), index_update_map, contig_finder); - - return IndexFromIdGraph( - target_indexing, - indexing, - index_parameters.initial_concrete_id_index, - loop_indexing.loopDomains()); -} - -namespace { - -class LoopIndexingTraversal { - enum class TraversalOrder { ForwardTopological, BackwardTopological }; - - public: - static std::vector forwardTopologicalOrder( - const std::vector& exprs) { - LoopIndexingTraversal traversal(exprs, TraversalOrder::ForwardTopological); - return traversal.getExprList(); - } - - static std::vector backwardTopologicalOrder( - const std::vector& exprs) { - LoopIndexingTraversal traversal(exprs, TraversalOrder::BackwardTopological); - return traversal.getExprList(); - } - - private: - explicit LoopIndexingTraversal( - const std::vector& exprs, - TraversalOrder traversal_order); - - // Returns the vals following the expression in either - // forward or backward order. - const std::vector& nextValsInTraversalOrder(Expr* expr); - - // Returns the vals that the expression follows in either - // forward or backward order. - const std::vector& prevValsInTraversalOrder(Expr* expr); - - // Returns the sorted list according to the given traversal order. - std::vector getExprList(); - - private: - // Reference to original un-sorted expression list. - const std::vector& exprs_; - - // The traversal order in this pass. - const TraversalOrder traversal_order_ = TraversalOrder::ForwardTopological; - - // Internal record of concrete id's and it's corresponding - // iterdomain expression that defines the exact index. - std::unordered_map concrete_id_to_dependency_; -}; - -LoopIndexingTraversal::LoopIndexingTraversal( - const std::vector& exprs, - TraversalOrder traversal_order) - : exprs_(exprs), traversal_order_(traversal_order) { - // Populate concrete id dependencies: - for (auto expr : exprs_) { - auto next_ids = - ir_utils::filterByType(nextValsInTraversalOrder(expr)); - for (auto id : next_ids) { - auto concrete_id = - GpuLower::current()->info().caMap().getConcreteMappedID( - id, IdMappingMode::EXACT); - NVF_ERROR( - concrete_id_to_dependency_.insert(std::make_pair(concrete_id, expr)) - .second, - "Repeated dependency, invalid iterdomain traversal."); - } - } -} - -const std::vector& LoopIndexingTraversal::nextValsInTraversalOrder( - Expr* expr) { - switch (traversal_order_) { - case TraversalOrder::ForwardTopological: - return expr->outputs(); - break; - case TraversalOrder::BackwardTopological: - return expr->inputs(); - break; - - default: - NVF_THROW("unimplemented traversal order"); - } - return expr->inputs(); -} - -const std::vector& LoopIndexingTraversal::prevValsInTraversalOrder( - Expr* expr) { - switch (traversal_order_) { - case TraversalOrder::ForwardTopological: - return expr->inputs(); - break; - case TraversalOrder::BackwardTopological: - return expr->outputs(); - break; - - default: - NVF_THROW("unimplemented traversal order"); - } - return expr->inputs(); -} - -std::vector LoopIndexingTraversal::getExprList() { - std::deque to_visit(exprs_.begin(), exprs_.end()); - - // pre-allocate result space. - std::vector result; - result.reserve(exprs_.size()); - - // Keeps track of visited and inserted expressions. - // An expr is visited if it has been placed in result list. - // An expr is inserted if the traversal has put the expr on - // the top of the stack once. Repeated insertion of the same - // expression would never be observed if the underlying - // dependency of the expressions is cycle free. - std::unordered_set visited, inserted; - - while (!to_visit.empty()) { - auto top = to_visit.front(); - if (visited.count(top)) { - to_visit.pop_front(); - continue; - } - - bool ready = true; - - for (auto prev_id : - ir_utils::filterByType(prevValsInTraversalOrder(top))) { - auto prev_expr_it = concrete_id_to_dependency_.find( - GpuLower::current()->info().caMap().getConcreteMappedID( - prev_id, IdMappingMode::EXACT)); - if (prev_expr_it != concrete_id_to_dependency_.end()) { - auto prev_expr = prev_expr_it->second; - if (!visited.count(prev_expr)) { - ready = false; - to_visit.push_front(prev_expr); - NVF_ERROR( - inserted.insert(prev_expr).second, - "Circular dependency in loop index expressions."); - break; - } - } - } - - if (ready) { - visited.insert(top); - result.emplace_back(top); - to_visit.pop_front(); - } - } - - return result; -} - -} // namespace - -void LoopIndexingAnalysis::collectOutOfLineExprs() { - // Keep track of all the id's that can be resolved without - // iterdomains on the left of ca axes. - std::unordered_set out_of_line_ids; - - // Start the set with all the loop ids. - std::transform( - consumer_tv_->getLoopDomain().begin() + - consumer_tv_->getComputeAtPosition(), - consumer_tv_->getLoopDomain().end(), - std::inserter(out_of_line_ids, out_of_line_ids.end()), - exactConcreteId); - - // Get the original selected list of index expressions - // in reverse topological order. - auto backward_expr_list = - LoopIndexingTraversal::backwardTopologicalOrder(replayed_exprs_); - - for (auto expr : backward_expr_list) { - auto id_outputs = ir_utils::filterByType(expr->outputs()); - if ( - // Check that all of the outputs are out of line - std::all_of( - id_outputs.begin(), - id_outputs.end(), - [&out_of_line_ids](IterDomain* id) { - return out_of_line_ids.count( - GpuLower::current()->info().caMap().getConcreteMappedID( - id, IdMappingMode::EXACT)); - })) { - // Record out of line expression - out_of_line_exprs_.push_back(expr); - - // Add all of the expression inputs as out of line id's. - auto id_inputs = ir_utils::filterByType(expr->inputs()); - std::transform( - id_inputs.begin(), - id_inputs.end(), - std::inserter(out_of_line_ids, out_of_line_ids.end()), - exactConcreteId); - } - } -} - -std::vector LoopIndexing::getForwardExprList() const { - return LoopIndexingTraversal::forwardTopologicalOrder(index_exprs_); -} - -std::vector LoopIndexing::getBackwardExprList() const { - return LoopIndexingTraversal::backwardTopologicalOrder(index_exprs_); -} - -std::unordered_set LoopIndexing::getAllExactConcreteIdSet() const { - std::unordered_set all_id_set; - for (auto expr : index_exprs_) { - auto out_ids = ir_utils::filterByType(expr->outputs()); - std::transform( - out_ids.begin(), - out_ids.end(), - std::inserter(all_id_set, all_id_set.end()), - exactConcreteId); - - auto in_ids = ir_utils::filterByType(expr->inputs()); - std::transform( - in_ids.begin(), - in_ids.end(), - std::inserter(all_id_set, all_id_set.end()), - exactConcreteId); - } - return all_id_set; -} - -namespace { - -//! Returns true if id is mapped together with any id in -//! the vector ids by permissive compute at map. -bool isPermissivelyMappedWithAny(IterDomain* id, const std::vector& ids) { - return std::any_of(ids.begin(), ids.end(), [&](Val* val) { - if (!(val->isA() && - GpuLower::current()->info().caMap().areMapped( - id, val->as(), IdMappingMode::PERMISSIVE))) { - return false; - } - // When id is an input to resize, make sure the resize argumens - // are compatible. This is important when, for example, a tensor - // is padded two times differently but to the same shape, and the - // pad outputs are exactly mapped. In such a case, there're two - // paths from the post logical ID to the original input ID, and - // the correct path depends on the path where this producer is - // used as a producer. See the FusionPad8 test for a concrete - // example. - if (auto id_resize = dynamic_cast(id->uses().at(0))) { - auto mapped_id_resize = - dynamic_cast(val->as()->uses().at(0)); - NVF_ERROR(mapped_id_resize != nullptr); - if (!(id_resize->leftExpand()->sameAs(mapped_id_resize->leftExpand()) && - id_resize->rightExpand()->sameAs( - mapped_id_resize->rightExpand()))) { - return false; - } - } - return true; - }); -} - -class LoopIndexingPreferredPathCompute : public IterVisitor { - public: - static std::unordered_set compute( - const TensorView* original_tv, - const LoopIndexing& loop_indexing, - bool use_replay_map, - const std::unordered_map& p2c_map) { - LoopIndexingPreferredPathCompute compute; - - auto all_concrete_ids = loop_indexing.getAllExactConcreteIdSet(); - - // Annotate all ids - auto all_original_ids = DependencyCheck::getAllValsBetween( - {original_tv->getMaybeAllocationDomain().begin(), - original_tv->getMaybeAllocationDomain().end()}, - {original_tv->getLoopDomain().begin(), - original_tv->getLoopDomain().end()}); - - for (auto original_id : - ir_utils::filterByType(all_original_ids)) { - auto mapped_id = original_id; - if (use_replay_map) { - auto c_id_it = p2c_map.find(original_id); - if (c_id_it == p2c_map.end()) { - continue; - } - mapped_id = c_id_it->second; - } - auto concrete_original_id = - GpuLower::current()->info().caMap().getConcreteMappedID( - mapped_id, IdMappingMode::EXACT); - if (all_concrete_ids.count(concrete_original_id)) { - if (original_id->isBroadcast() || original_id->isReduction() || - original_id->isStride()) { - continue; - } - compute.preferred_path_.insert(concrete_original_id); - } - } - - for (auto expr : loop_indexing.getForwardExprList()) { - compute.dispatch(expr); - } - - return compute.preferred_path_; - } - - private: - void dispatch(Expr* e) override { - // If an input ID is marked, propagate the marking to outputs of the - // expression - auto all_iter_inputs = ir_utils::filterByType(e->inputs()); - if (std::any_of( - all_iter_inputs.begin(), - all_iter_inputs.end(), - [&](IterDomain* inp_id) { - return this->preferred_path_.find( - GpuLower::current() - ->info() - .caMap() - .getConcreteMappedID( - inp_id, IdMappingMode::EXACT)) != - this->preferred_path_.end(); - })) { - auto all_iter_outputs = ir_utils::filterByType(e->outputs()); - - std::transform( - all_iter_outputs.begin(), - all_iter_outputs.end(), - std::inserter(preferred_path_, preferred_path_.end()), - exactConcreteId); - } - } - - std::unordered_set preferred_path_; -}; - -} // namespace - -// External interface for preferred path propagation. -std::unordered_set buildLoopIndexingPreferredPath( - const TensorView* original_tv, - const LoopIndexing& loop_indexing, - bool use_replay_map, - std::unordered_map p2c_map) { - return LoopIndexingPreferredPathCompute::compute( - original_tv, loop_indexing, use_replay_map, p2c_map); -} - -// Get an logical IterDomain that is mapped with an IterDomain. If -// multiple such IDs exist, select one whose input IDs are mapped with -// the consumer IDs. This is to ensure the path from the loop -// IterDomains to the root matches with the consumer tensor. -// Additionally, when none of the candidate iter domain has all of its -// inputs mapped with the consumer tensor, prefer one that has at -// least one mapped. This matters when the consumer tensor only has -// one of the merge inputs, for example. -IterDomain* getLogicalIDToTraverse( - IterDomain* id, - const std::vector& consumer_all_ids) { - const auto& logical_ids = - GpuLower::current()->info().caMap().getLogicalDomainsOfIdGroup( - id, IdMappingMode::PERMISSIVE); - if (logical_ids.empty()) { - return nullptr; - } - - // Keep track of an iter domain that has at least one input mapped. - IterDomain* fallback_candidate = nullptr; - - for (auto logical_id : logical_ids) { - auto def = logical_id->definition(); - if (def == nullptr) { - continue; - } - - auto logical_id_inputs = ir_utils::filterByType(def->inputs()); - if (std::all_of( - logical_id_inputs.begin(), - logical_id_inputs.end(), - [&](IterDomain* logical_id_input) { - return isPermissivelyMappedWithAny( - logical_id_input, consumer_all_ids); - })) { - return logical_id; - } - - if (std::any_of( - logical_id_inputs.begin(), - logical_id_inputs.end(), - [&](IterDomain* logical_id_input) { - return isPermissivelyMappedWithAny( - logical_id_input, consumer_all_ids); - })) { - if (fallback_candidate == nullptr) { - fallback_candidate = logical_id; - } - } - } - - if (fallback_candidate != nullptr) { - return fallback_candidate; - } - - // No mapped ID found, which means the consumer is a post-view - // tensor. In that case, it shouldn't matter which view path to - // traverse, so just return the first one. - return logical_ids.at(0); -} - -} // namespace nvfuser diff --git a/csrc/device_lower/analysis/index_compute.h b/csrc/device_lower/analysis/index_compute.h deleted file mode 100644 index 4f16dac4f71..00000000000 --- a/csrc/device_lower/analysis/index_compute.h +++ /dev/null @@ -1,320 +0,0 @@ -// clang-format off -/* - * SPDX-FileCopyrightText: Copyright (c) 2023-present NVIDIA CORPORATION & AFFILIATES. - * All rights reserved. - * SPDX-License-Identifier: BSD-3-Clause - */ -// clang-format on -#pragma once - -#include -#include -#include - -namespace nvfuser { - -// Struct to hold useful information from an index pass on iterdomain graph. -// Used to return the IndexCompute structure back to the indexing calls in -// index_compute.cpp. Other structurs are required to resolve the actual -// indexing math there. -struct IndexFromIdGraph { - IndexCompute index; - IndexCompute concrete_index; - std::unordered_map initial_concrete_index_map; - std::vector resolved_loop_domains; - - explicit IndexFromIdGraph( - IndexCompute index, - IndexCompute concrete_index, - std::unordered_map initial_concrete_index_map, - std::vector loop_domains); -}; - -//! Indexing interface, returns IndexFromIdGraph which the IndexCompute object -//! can be queried from directly for the produced indexing. If producer_tv != -//! nullptr producer will be indexed, if producer_tv == nullptr consumer will be -//! indexed. If is_global global indexing will be done, else shared memory or -//! local indexing will be performed. -IndexFromIdGraph getTensorIndexFromIdGraph( - const std::vector& loops, - const TensorView* consumer_tv, - const TensorView* producer_tv = nullptr, - bool is_global = true, - const std::unordered_map& c2p_map = {}); - -//! Indexing interface for calculating predicate index returns IndexFromIdGraph -//! which the IndexCompute object can be queried from directly for the produced -//! indexing If is_start_predicate, will produce indexing math for the start -//! predicates. -IndexFromIdGraph getPredicateIndexingFromIdGraph( - const std::vector& loops, - TensorView* consumer_tv, - kir::ForLoop* unswitch_or_vec_loop, - IterDomain* circular_buffer_axis, - bool is_start_predicate); - -//! getTensorIndexFromIdGraph is the function that index_compute will call very -//! straightforwardly. However, for implementing the new indexing logic that -//! starts to abstract some of the indexing away from index_compute we need to -//! move quite a bit of the intertwined indexing logic away from the -//! index_compute file and the index_reference_replay file. This is because we -//! want to separate out what has to be done on the fly, from what analysis we -//! can do early on with the iter domain graph and associated properties. -//! -//! getTensorIndexFromIdGraph places this analysis internally in -//! LoopIndexingAnalysis. LoopIndexingAnalysis though has to communicate to: -//! 1) index_compute.cpp::IndexCompute to tell IndexCompute which expressions -//! it needs to traverse to compute the indexing math. -//! -//! LoopIndexing is nothing but a mechanism for this communication. -//! -//! Holds information needed to produce indexing math. In the current version of -//! indexing pass, the iter domains combined with the loop nests are the source -//! of truth in terms of resolving the actual integer indexing math from the -//! sequence of iterdomain transforms. -//! -//! This information is crtiical in resolving indexing associated with complex -//! broadcast patterns. Check FusionComplexBCast* test cases as well as -//! Indexing* tests for examples where resolving indices from IterDomain -//! transformations can be challenging. -//! -//! The source of this challenge is due to inling patterns where the IterDomains -//! responsible for control flow are not local to a particular TensorView. -//! Broadcast, operations like view/reshape, and gather/shift can make indexing -//! local buffers complex because of the complex effects inlining into other -//! TensorViews produce. -//! -//! TODO: -//! The first iteration tries to match the semantics of reference -//! replay without any new logic. In a follow up iteration will -//! need to revisit a few further pathological patterns. -//! -//! Note: -//! The current implementation of loop indexing pass works on -//! equivalent classes defined by ComputeAt exact map. The -//! list of expressions stored in this class form a "reference", graph of -//! iterdomain expressions when all of their inputs and outputs are replaced -//! with their exact concrete mapped id's. -//! -//! Here an invariant in a graph of iterdomain expressions is that -//! each iterdomain is produced exactly once and is either a loop domain -//! or has been consumed exactly once by another expression. This makes sure -//! that a well defined indexing can be generated for each of the concrete ids -//! whenever we either forward or backward traverse the graph. -class LoopIndexing { - public: - //! Returns the original loop nest. - const auto& loops() const { - return loops_; - } - - //! Returns the vector of Iterdomains - //! that match the original loop pattern. - const auto& loopDomains() const { - return loop_domains_; - } - - const auto& loopRootDomains() const { - return loop_root_; - } - - //! Returns the consumer tv that the view info - //! was derived from. - auto consumerTv() const { - return consumer_tv_; - } - - //! Returns the set of Iterdomain transforms that - //! define the correct indexing path, in forward - //! topological order. - std::vector getForwardExprList() const; - - //! Returns the set of Iterdomain transforms that - //! define the correct indexing path, in backward - //! topological order. - std::vector getBackwardExprList() const; - - //! Returns the set of out of line expressions in - //! reverse topological order. - const std::vector& getBackwardOutOfLineExprList() const { - return out_of_line_exprs_; - } - - //! Returns all exact concrete id's that were produced - //! or consumed in the selected indexing expressions - std::unordered_set getAllExactConcreteIdSet() const; - - private: - friend class LoopIndexingAnalysis; - - //! The loop nest that this loop indexing is derived from. - std::vector loops_; - - //! Consumer tv, where the view related info was derived from. - const TensorView* consumer_tv_ = nullptr; - - //! The source iterdomains that all the Iterdomain transforms - //! in this loop nest originated from. - std::vector loop_root_; - - //! The loop iterdomains that the original loop nests correspond - //! to. May be longer than loops_ with the dangling iterdomains - //! appended towards the end. - std::vector loop_domains_; - - //! The selected sequence of expressions that should represent - //! the correct indexing math from the given loop nest. - std::vector index_exprs_; - - //! The subset of sequence of expressions that can be resolved - //! with only the iterdomains on the right of consumer tv's ca - //! axis. - //! Expressions are ordered in reverse topological order. - std::vector out_of_line_exprs_; -}; - -class LoopIndexingAnalysis { - public: - static LoopIndexing fromLoopAndConsumer( - const std::vector& loops, - const TensorView* consumer_tv); - - //! Return all concrete IDs that can be reachable from a given list - //! of consumer loop IDs. Reachability is defined as the existence - //! an indexing path from the the loop IDs - static VectorOfUniqueEntries getReplayableConcreteIDs( - const std::vector& consumer_loop_ids, - const TensorView* consumer_tv); - - private: - explicit LoopIndexingAnalysis( - const std::vector& loops, - const TensorView* consumer_tv); - - explicit LoopIndexingAnalysis( - const std::vector& consumer_loop_ids, - const TensorView* consumer_tv); - - void run(); - - //! Populate derived information into a LoopIndexing - //! data structure. - LoopIndexing getLoopIndexing(const std::vector& loops) { - LoopIndexing indexing; - indexing.loops_ = loops; - indexing.consumer_tv_ = consumer_tv_; - indexing.loop_root_ = loop_root_domains_; - indexing.loop_domains_ = loop_domains_.vector(); - indexing.index_exprs_ = replayed_exprs_; - indexing.out_of_line_exprs_ = out_of_line_exprs_; - return indexing; - } - - //! Validates that the current loop structure is well formed, in the sense - //! that ca_map would not map any two loops in the loop nest together. - void validateLoopStructure(const std::vector& loops); - - //! Start at the loop iter domains, and traverse back into history on the - //! concrete IDs in the exact map calling "visitExpr" expressions through the - //! history. - void traverseFromDomainVals(); - - //! Concretize the given iterdomain and record the visit (in deterministic - //! order) in terms of the exact mapped concrete id. Marks the mapping of the - //! id to the concrete id in "concrete_to_original_id_" and returns the - //! concrete id. - IterDomain* concretizeAndVisitId(IterDomain* id); - - //! If an equivalent expression has already been processed this function - //! simply returns. Otherwise puts the exact concrete IDs of inputs in - //! consumed_concrete_, and concrete IDs of outputs in produced_concrete_. - //! Then adds the expression to replayed_exprs_. - void visitExpr(Expr* expr); - - //! Iterates through provided vals, calls concretizeAndVisitId on them, and - //! returns if any of the returned vals are in existing_ids. This is used to - //! check if inputs or outputs of ID expressions have already been - //! produced/consumed in the traversal. Indexing only needs to consume/produce - //! one IterDomain per exact disjoint set. - bool visitIdsAndCheckDuplication( - const std::vector& vals, - const std::unordered_set& existing_ids); - - //! Fills loop_domains_ with the corresponding replayed_concrete_id mapping to - //! the provided loops. Must be done after the exact iterdomain "replay" - //! (traverseFromDomainVals). loop_domains_ are the original_id not the - //! concrete_id (translated with concrete_to_original_id). These iter domains - //! are used to grab the history that will be replayed in IndexCompute. We're - //! looking for "new" root domains and subsequent transformations, filling in - //! any missing "outputs" (or inputs for backward traversal). Then fills - //! loop_domains_ with all of these iter domains. - void constructLoopDomains(); - - //! Fills out_of_line_exprs_ by traversing the selected list of - //! expressions in reverse topological order and collect iterdomains - //! on the indexing paths that only involves loop id's on the right - //! of consumer's ca axis. - void collectOutOfLineExprs(); - - private: - //! Original consumer tv to derive view info from. - const TensorView* consumer_tv_ = nullptr; - - // Exact concrete domains that has been used - // in the traversal connection. - std::unordered_set produced_concrete_; - std::unordered_set consumed_concrete_; - - //! Iterdomains that the corresponding loops are generated from. - std::vector initial_loop_domain_ids_; - - //! All Id's in consumer's transform history - std::vector all_consumer_id_vals_; - - //! Concrete iterdomains visited in the domain traversal, - //! in the order they are visited in traverseFromDomainVals. - VectorOfUniqueEntries replayed_concrete_ids_; - - //! Keeping track of the original visited id's before they - //! were concretized. - std::unordered_map concrete_to_original_id_; - - //! Map from concrete id to its single consumer on the selected - //! iterdomain expression list. - std::unordered_map concrete_id_to_consumer_; - - //! Source domains that all the Iterdomain transforms - //! in the loop nest originated from. - std::vector loop_root_domains_; - - //! Leaf domains representing the original loop structure - VectorOfUniqueEntries loop_domains_; - - //! Selected list of exprs that will produce and consume each - //! of the exact concrete ids from the loop nest exactly once. - std::vector replayed_exprs_; - - //! Set of expressions from the selected list that can be - //! resolved from axes on the right of ca axes. - std::vector out_of_line_exprs_; -}; - -// When indexing there are sometimes an option to propagate an index down -// multiple paths. This will return the IterDomains in the history of the -// reference domain and mark which paths should be taken (if there's a -// preference) to reach the roots provided in preferred_roots. -std::unordered_set buildLoopIndexingPreferredPath( - const TensorView* original_tv, - const LoopIndexing& loop_indexing, - bool use_replay_map = false, - std::unordered_map p2c_map = {}); - -// Get an logical IterDomain that is mapped with an IterDomain. If -// multiple such IDs exist, select one whose input IDs are mapped with -// the consumer IDs. This is to ensure the path from the loop -// IterDomains to the root matches with the consumer tensor. -IterDomain* getLogicalIDToTraverse( - IterDomain* id, - const std::vector& consumer_all_ids); - -} // namespace nvfuser diff --git a/csrc/device_lower/analysis/sync_information.cpp b/csrc/device_lower/analysis/sync_information.cpp index da76301104c..f5bb93ce95b 100644 --- a/csrc/device_lower/analysis/sync_information.cpp +++ b/csrc/device_lower/analysis/sync_information.cpp @@ -6,7 +6,6 @@ */ // clang-format on #include -#include #include #include #include diff --git a/csrc/device_lower/pass/index.cpp b/csrc/device_lower/pass/index.cpp index d131ce73c11..b4326da1f87 100644 --- a/csrc/device_lower/pass/index.cpp +++ b/csrc/device_lower/pass/index.cpp @@ -5,7 +5,6 @@ * SPDX-License-Identifier: BSD-3-Clause */ // clang-format on -#include #include #include #include diff --git a/csrc/device_lower/pass/magic_zero.cpp b/csrc/device_lower/pass/magic_zero.cpp index 359c641d8a7..0e753195c9f 100644 --- a/csrc/device_lower/pass/magic_zero.cpp +++ b/csrc/device_lower/pass/magic_zero.cpp @@ -7,7 +7,6 @@ // clang-format on #include -#include #include #include #include @@ -126,127 +125,4 @@ bool needsMagicZero( return loop->isUnrolled() && (!ref_dom_simple || !ind_simple); } -void protectNonPredicateIndexWithMagicZero( - const std::vector& loops, - const std::vector& loop_domains, - std::unordered_map& concrete_loop_idx_map) { - if (!GpuLower::current()->isNvFuserZeroEnabled()) { - return; - } - - // Find magic zero insertion point - IterDomain* magic_zero_loop = nullptr; - - // Search for proper magic zero insertion point, - // prefer innermost. - for (auto idx : arange(loops.size())) { - auto loop = loops[idx]; - auto concrete_loop_id = - GpuLower::current()->info().caMap().getConcreteMappedID( - loop_domains[idx], IdMappingMode::EXACT); - auto loop_ind = concrete_loop_idx_map.at(concrete_loop_id); - - // Save the concrete id if this loop id is decided to - // be the insertion point by the magic zero util. - if (needsMagicZero(loop, concrete_loop_id, loop_ind)) { - magic_zero_loop = concrete_loop_id; - } - } - - // Insert magic zero if insertion point found - if (magic_zero_loop != nullptr && - concrete_loop_idx_map.count(magic_zero_loop)) { - auto& ind = concrete_loop_idx_map.at(magic_zero_loop); - ind = SimplifyingIrBuilder::addExpr( - ind, GpuLower::current()->kernel()->magicZeroVal()); - } -} - -namespace { - -//! Protect loop_index_to_protect appearing in overall_index_val -IndexMagicZeroInfo protectIndexByReplacingLoopIndex( - IterDomain* loop_id, - Val* overall_index_val, - Val* loop_index_to_protect) { - auto protected_loop_index = SimplifyingIrBuilder::addExpr( - loop_index_to_protect, GpuLower::current()->kernel()->magicZeroVal()); - - std::unordered_map replacement_map; - replacement_map[loop_index_to_protect] = protected_loop_index; - - auto protected_index = - ir_utils::replaceValRecursively(overall_index_val, replacement_map); - - IndexMagicZeroInfo info; - info.index = protected_index; - info.original_loop_index = loop_index_to_protect; - info.protected_loop_index = protected_loop_index; - info.loop_id = loop_id; - return info; -} - -} // namespace - -IndexMagicZeroInfo protectPredicateIndexWithMagicZero( - Val* index, - const IndexFromIdGraph& id_graph, - const std::vector& loops) { - if (!GpuLower::current()->isNvFuserZeroEnabled()) { - IndexMagicZeroInfo not_proteced; - not_proteced.index = index; - return not_proteced; - } - - // Gather the loop indices - std::unordered_set loop_indices; - for (auto loop_id : id_graph.resolved_loop_domains) { - auto concrete_loop_id = - GpuLower::current()->info().caMap().getConcreteMappedID( - loop_id, IdMappingMode::EXACT); - auto index_it = id_graph.initial_concrete_index_map.find(concrete_loop_id); - NVF_ERROR( - index_it != id_graph.initial_concrete_index_map.end(), - "Index not found for loop: ", - concrete_loop_id->toString()); - auto loop_index = index_it->second; - loop_indices.insert(loop_index); - } - - // Figure out which loop indices are used in index - const auto vals = DependencyCheck::getAllValsBetween(loop_indices, {index}); - - // Traverser from the inner-most loop and apply the magic-zero - // prorection if needed - for (int64_t i = static_cast(loops.size()) - 1; i >= 0; --i) { - auto loop = loops.at(i); - auto loop_id = id_graph.resolved_loop_domains.at(i); - NVF_ERROR(GpuLower::current()->info().caMap().areMapped( - loop_id, loop->iter_domain(), IdMappingMode::PERMISSIVE)); - IterDomain* concrete_loop_id = - GpuLower::current()->info().caMap().getConcreteMappedID( - loop_id, IdMappingMode::EXACT); - auto index_it = id_graph.initial_concrete_index_map.find(concrete_loop_id); - NVF_ERROR(index_it != id_graph.initial_concrete_index_map.end()); - auto loop_index = index_it->second; - - const auto is_loop_index_used = - std::find(vals.begin(), vals.end(), loop_index) != vals.end(); - - if (!is_loop_index_used) { - continue; - } - - if (needsMagicZero(loop, concrete_loop_id, loop_index)) { - return protectIndexByReplacingLoopIndex(loop_id, index, loop_index); - } - } - - // No loop is identified to require protection with magic zero. Just - // return the index argument as is - IndexMagicZeroInfo not_proteced; - not_proteced.index = index; - return not_proteced; -} - } // namespace nvfuser diff --git a/csrc/device_lower/pass/magic_zero.h b/csrc/device_lower/pass/magic_zero.h index b8ce6047d1e..c71921077b2 100644 --- a/csrc/device_lower/pass/magic_zero.h +++ b/csrc/device_lower/pass/magic_zero.h @@ -16,8 +16,6 @@ namespace nvfuser { -struct IndexFromIdGraph; - //! Insert magic zero definition at the begining of the kernel. Insert magic //! zero update after every (outer most) loop nest with a compile time extent. //! @@ -49,38 +47,4 @@ bool needsMagicZero( IterDomain* reference_domain = nullptr, Val* ind = nullptr); -struct IndexMagicZeroInfo { - //! Index that may be updated with magic zero - Val* index = nullptr; - //! Loop index that is protected by magic zero. nullptr if no loop - //! is protected - Val* original_loop_index = nullptr; - //! Protected loop index. nullptr if no loop is protected - Val* protected_loop_index = nullptr; - //! Protected loop. nullptr if no loop is protected - IterDomain* loop_id = nullptr; -}; - -//! Protect an index val of an IterDomain with magic zero -//! -//! This should be only used for predicate indexing. -//! -//! No protection is done if none of the loops is determined to require -//! protection by needsMagicZero. -IndexMagicZeroInfo protectPredicateIndexWithMagicZero( - Val* index, - const IndexFromIdGraph& id_graph, - const std::vector& loops); - -//! Protect an index val of a tensor with magic zero -//! -//! This should be only used for non-predicate indexing. -//! -//! No protection is done if none of the loops is determined to require -//! protection by needsMagicZero. -void protectNonPredicateIndexWithMagicZero( - const std::vector& loops, - const std::vector& loop_domains, - std::unordered_map& concrete_loop_idx_map); - } // namespace nvfuser diff --git a/csrc/device_lower/pass/rng.cpp b/csrc/device_lower/pass/rng.cpp index 3347dd5cc6a..7a83a14a9c6 100644 --- a/csrc/device_lower/pass/rng.cpp +++ b/csrc/device_lower/pass/rng.cpp @@ -7,7 +7,6 @@ // clang-format on #include -#include #include #include #include diff --git a/csrc/id_model/indexing.cpp b/csrc/id_model/indexing.cpp index 3fa1d6c4ab9..0d34db61eeb 100644 --- a/csrc/id_model/indexing.cpp +++ b/csrc/id_model/indexing.cpp @@ -10,7 +10,6 @@ #include #include "debug.h" -#include "device_lower/analysis/index_compute.h" #include "device_lower/analysis/non_divisible_split.h" #include "device_lower/lower2device.h" #include "device_lower/pass/magic_zero.h" @@ -705,7 +704,6 @@ std::vector TensorIndexer::getPredicates( std::unordered_set already_indexed_domains; - // Follow the same approach as Index::getReferenceRootPredicates. for (const auto& predicate_domain : predicate_domains) { const auto& predicate_domain_group = traversalGraph().toGroup(predicate_domain); diff --git a/csrc/id_model/indexing_utils.h b/csrc/id_model/indexing_utils.h index a3e737415a5..ed1ec32b0a8 100644 --- a/csrc/id_model/indexing_utils.h +++ b/csrc/id_model/indexing_utils.h @@ -7,7 +7,6 @@ // clang-format on #pragma once -#include "device_lower/analysis/index_compute.h" #include "device_lower/lower2device.h" #include "device_lower/utils.h" #include "id_model/id_model.h" diff --git a/csrc/index_compute.cpp b/csrc/index_compute.cpp index df382fd1247..84dd9c4da66 100644 --- a/csrc/index_compute.cpp +++ b/csrc/index_compute.cpp @@ -11,1128 +11,17 @@ #include -#include -#include #include -#include -#include -#include -#include -#include #include #include #include #include -#include #include -#include -#include -#include namespace nvfuser { -bool IndexCompute::hasUnswitchedDependentDomains(IterDomain* id) const { - auto concrete_id = maybeGetExactMapConcreteID(id); - auto it = unswitched_domain_map_.find(concrete_id); - return it != unswitched_domain_map_.end() && !it->second.empty(); -} - -void IndexCompute::initializeUnswitchDomainMap() { - NVF_ERROR(unswitched_domain_map_.empty()); - for (auto id : unswitched_loop_domains_) { - auto concrete_id = maybeGetExactMapConcreteID(id); - unswitched_domain_map_.emplace( - concrete_id, - std::vector>{ - std::deque{concrete_id}}); - } -} - -void IndexCompute::updateUnswitchedDomains(Expr* expr) { - if (auto split = dynamic_cast(expr)) { - auto split_in = maybeGetExactMapConcreteID(split->in()); - for (auto split_out : {split->inner(), split->outer()}) { - auto concrete_id = maybeGetExactMapConcreteID(split_out); - if (auto it = unswitched_domain_map_.find(concrete_id); - it != unswitched_domain_map_.end()) { - if (split_out == split->inner()) { - // In the case of upward traversal from the inner output, - // just copy the unswitched info - unswitched_domain_map_[split_in] = it->second; - } else { - // In the case of upward traversal from the outer output, - // prepend the inner domain to the lists - for (auto unswitched_dep_ids : it->second) { - unswitched_dep_ids.push_front( - maybeGetExactMapConcreteID(split->inner())); - unswitched_domain_map_[split_in].push_back(unswitched_dep_ids); - } - } - } - } - } else { - // Suppress a clang-tidy warning - NVF_ERROR(expr != nullptr); - // Propagate the unswitch info if any of outputs is - // unswitched. Unlike the split case, the propagated info - // is just reset as there's no obvious way to back-propagate the - // info through, e.g., merge - if (std::any_of( - expr->outputs().begin(), expr->outputs().end(), [this](Val* out) { - return out->isA() && - hasUnswitchedDependentDomains(out->as()); - })) { - for (auto inp : ir_utils::filterByType(expr->inputs())) { - auto inp_concrete = maybeGetExactMapConcreteID(inp); - unswitched_domain_map_.emplace( - inp_concrete, - std::vector>{ - std::deque{inp_concrete}}); - } - } - } -} - -void IndexCompute::handle(Split* split) { - auto in_id = maybeGetExactMapConcreteID(split->in()->as()); - auto outer_id = maybeGetExactMapConcreteID(split->outer()->as()); - auto inner_id = maybeGetExactMapConcreteID(split->inner()->as()); - - auto outer_it = index_map_.find(outer_id); - auto inner_it = index_map_.find(inner_id); - if (outer_it == index_map_.end() || inner_it == index_map_.end()) { - return; - } - - const auto outer_ind = outer_it->second; - const auto inner_ind = inner_it->second; - - const bool outer_zero = isZero(outer_id); - const bool inner_zero = isZero(inner_id); - - // We want to mark as zero merged in if we're working with shared or local - // memory, and the dimension we're working with is not part of the allocation, - // as we have special propagation rules for that scenario. - - // Maybe clear in_id as it could have been mapped over from another - // IndexCompute. Uncertain if this is needed but seems to be safe. - bool zero_merged_in = hasZeroMerged(in_id) || hasZeroMerged(inner_id) || - hasZeroMerged(outer_id); - - // If both are zero, the split input is also zero - if (inner_zero && outer_zero) { - zero_domains_.emplace(in_id); - } - - if (zero_merged_in) { - zero_merged_in_.emplace(in_id); - } - - if (isZero(in_id)) { - index_map_[in_id] = GpuLower::current()->kernel()->zeroVal(); - extent_map_[in_id] = GpuLower::current()->kernel()->zeroVal(); - } else if (zero_merged_in && outer_zero) { - index_map_[in_id] = inner_ind; - extent_map_[in_id] = getExtent(inner_id); - } else if (zero_merged_in && inner_zero) { - index_map_[in_id] = outer_ind; - extent_map_[in_id] = getExtent(outer_id); - } else { - index_map_[in_id] = SimplifyingIrBuilder::addExpr( - SimplifyingIrBuilder::mulExpr(outer_ind, getExtent(inner_id)), - inner_ind); - // The extent should be updated only when its allocation is - // partial, i.e., zero_merged_in is true. See PR #1270. - if (zero_merged_in) { - extent_map_[in_id] = SimplifyingIrBuilder::mulExpr( - getExtent(outer_id), getExtent(inner_id)); - } - } -} - -bool IndexCompute::isModuloInvalidUnswitchedIndex( - IterDomain* out_concrete_id, - Val* out_ind, - Val* inner_extent) const { - // If not in the unswitched domain map, this domain has no dependent - // unswitched domain - auto unswitched_domain_map_it = unswitched_domain_map_.find(out_concrete_id); - if (unswitched_domain_map_it == unswitched_domain_map_.end()) { - return false; - } - - for (const auto& unswitched_domain_list : unswitched_domain_map_it->second) { - NVF_ERROR(!unswitched_domain_list.empty()); - - // If the stride is a multiple of the inner extent, the loop - // unswitched index remains to be a valid maximum index as the - // module by the inner extent will be just zero. More - // specifically, the index for this unswitched domain would be (x - // - 1) * extent_of_inner_domain_0 * extent_of_inner_domain_1 - // ...., so if the stride component, i.e., the multiplication of all - // the inner extents is divisible by the merge inner extent, its - // contribution propagated to the inner path will be zero. This - // pattern is effectively the same as distributeDivisibleDivMod in - // the expr simplifier. - Val* stride = out_concrete_id->fusion()->oneVal(); - for (auto it = unswitched_domain_list.begin(); - it != unswitched_domain_list.end() - 1; - ++it) { - IterDomain* inner_id = *it; - stride = IrBuilder::mulExpr(stride, getExtent(inner_id)); - } - if (simplifyExpr(IrBuilder::modExpr(stride, inner_extent))->isZero()) { - continue; - } - - // Also, if the total extent including the inner domains is a - // divisible factor of the inner extent, the contribution by the - // unswitched domain is guaranteed to be still the maximum when - // propagated to the inner path. This pattern is effectively the - // same as distributeGcdRemainderDivMod in the expr simplifier. - Val* total_extent = - IrBuilder::mulExpr(stride, getExtent(unswitched_domain_list.back())); - if (simplifyExpr(IrBuilder::modExpr(inner_extent, total_extent)) - ->isZero()) { - continue; - } - - // Not proven to be safe. This does not mean it's proven to be - // invalid, but it's enough to make the generated code from the - // existing C++ tests and benchmarks remain unchanged - return true; - } - - return false; -} - -void IndexCompute::handle(Merge* merge) { - auto out_id = maybeGetExactMapConcreteID(merge->out()); - auto outer_id = maybeGetExactMapConcreteID(merge->outer()); - auto inner_id = maybeGetExactMapConcreteID(merge->inner()); - - auto out_it = index_map_.find(out_id); - if (out_it == index_map_.end()) { - return; - } - auto out_ind = out_it->second; - - auto zero = GpuLower::current()->kernel()->zeroVal(); - - if (isZero(out_id)) { - index_map_[outer_id] = zero; - index_map_[inner_id] = zero; - // TODO: Why do we set extent_map_ to zero? This has to be protected by zero - // merged in, but seems logical to me the extent would still be one. - extent_map_[outer_id] = zero; - extent_map_[inner_id] = zero; - zero_domains_.emplace(outer_id); - zero_domains_.emplace(inner_id); - return; - } - - if (!hasZeroMerged(out_id) && contig_ids_.find(out_id) != contig_ids_.end()) { - // Contiguous indexing path - auto input_ids = ir_utils::iterDomainInputsOfOrderedAs( - {merge->out()}, td_->maybeAllocation()); - - // Shouldn't hit this, but don't want to segfault if somehow we do. - NVF_ERROR(!input_ids.empty()); - - // Try to find the last non broadcast entry to put the index in if it's a - // contiguous merge. This isn't strictly necessary but there's implicit - // assumptions in the indexing logic that assume broadcasted allocation - // domains can be ignored. This logic is just to try and match that logic. - // Initialize everything to zero. - for (auto alloc_id : input_ids) { - index_map_[alloc_id] = zero; - } - - // If all are broadcast we can just send the index to the last entry. - if (std::ranges::all_of(input_ids, [](IterDomain* id) { - // I don't think reductions can be in here, but strictly matching the - // logic in the indexing functions like - // getNonGlobalConsumerStridedIndices - return id->isBroadcast() || id->isReduction() || id->isStride(); - })) { - index_map_[*(input_ids.end() - 1)] = out_ind; - } else { - for (auto id_it = input_ids.rbegin(); id_it != input_ids.rend(); - id_it++) { - auto id = *id_it; - if (id->isBroadcast() || id->isReduction() || id->isStride()) { - continue; - } else { - index_map_[id] = out_ind; - break; - } - } - } - - return; - } - - Val* inner_extent = getExtent(inner_id); - - const auto outer_extent = getExtent(outer_id); - - if (inner_id->isBroadcast() && inner_extent->isOneInt()) { - // Propagate away from broadcast dims - index_map_[outer_id] = out_ind; - index_map_[inner_id] = zero; - - extent_map_[outer_id] = getExtent(out_id); - if (hasZeroMerged(out_id)) { - zero_merged_in_.insert(outer_id); - } - } else if (outer_id->isBroadcast() && outer_extent->isOneInt()) { - // Propagate away from broadcast dims - index_map_[outer_id] = zero; - index_map_[inner_id] = out_ind; - - extent_map_[inner_id] = getExtent(out_id); - if (hasZeroMerged(out_id)) { - zero_merged_in_.insert(inner_id); - } - } else if (hasZeroMerged(out_id)) { - // Don't propagate to inner id if it's comprised of only broadcast - // allocation domains, unless outer is also all broadcast domains. Index - // shouldn't be anything but zero if both inner and outer are all broadcast - // domains, but didn't add a hard check for this. See Indexing5 test. - if (!inner_id->isBroadcast() && !outer_id->isBroadcast()) { - // If neither dimension is a broadcast (should be true for reference - // indexing) pick the preferred path or the inner path. - if (preferred_paths_.find(outer_id) != preferred_paths_.end() && - preferred_paths_.find(inner_id) == preferred_paths_.end()) { - // Marked that we should prop through outer, not inner. - index_map_[outer_id] = out_ind; - extent_map_[outer_id] = getExtent(out_id); - index_map_[inner_id] = zero; - extent_map_[inner_id] = zero; - zero_domains_.emplace(inner_id); - } else { - // Prop through inner - index_map_[inner_id] = out_ind; - extent_map_[inner_id] = getExtent(out_id); - index_map_[outer_id] = zero; - extent_map_[outer_id] = zero; - zero_domains_.emplace(outer_id); - } - } else if (inner_id->isBroadcast() && !outer_id->isBroadcast()) { - // Inner is broadcast and outer isn't, prop through outer - index_map_[outer_id] = out_ind; - extent_map_[outer_id] = getExtent(out_id); - index_map_[inner_id] = zero; - extent_map_[inner_id] = zero; - zero_domains_.emplace(inner_id); - } else { - // Default to propagating through inner - index_map_[inner_id] = out_ind; - extent_map_[inner_id] = getExtent(out_id); - index_map_[outer_id] = zero; - extent_map_[outer_id] = zero; - zero_domains_.emplace(outer_id); - } - zero_merged_in_.emplace(inner_id); - zero_merged_in_.emplace(outer_id); - } else { - index_map_[outer_id] = SimplifyingIrBuilder::divExpr(out_ind, inner_extent); - // Take the absolute maximum if module could result in an invalid - // index for an unswitched domain - index_map_[inner_id] = - isModuloInvalidUnswitchedIndex(out_id, out_ind, inner_extent) - ? SimplifyingIrBuilder::subExpr( - inner_extent, inner_extent->fusion()->oneVal()) - : SimplifyingIrBuilder::modExpr(out_ind, inner_extent); - } -} - -void IndexCompute::handle(Swizzle* swizzle) { - auto out_x_id = maybeGetExactMapConcreteID(swizzle->outX()); - auto out_y_id = maybeGetExactMapConcreteID(swizzle->outY()); - auto in_x_id = maybeGetExactMapConcreteID(swizzle->inX()); - auto in_y_id = maybeGetExactMapConcreteID(swizzle->inY()); - - auto out_x_it = index_map_.find(out_x_id); - auto out_y_it = index_map_.find(out_y_id); - - if (out_x_it == index_map_.end() || out_y_it == index_map_.end()) { - return; - } - - const auto out_x_ind = out_x_it->second; - const auto out_y_ind = out_y_it->second; - - std::pair swizzled_index = dispatchSwizzle( - swizzle->swizzleType(), - out_x_ind, - out_y_ind, - getExtent(out_x_id), - getExtent(out_y_id)); - index_map_[in_x_id] = swizzled_index.first; - index_map_[in_y_id] = swizzled_index.second; -} - -void IndexCompute::handle(Resize* resize) { - auto out_id = maybeGetExactMapConcreteID(resize->out()); - auto in_id = maybeGetExactMapConcreteID(resize->in()); - - auto out_it = index_map_.find(out_id); - - if (out_it == index_map_.end()) { - return; - } - - const auto out_ind = out_it->second; - - if (isZero(out_id) || hasZeroMerged(out_id)) { - // When the out ID is (partially) zero, the in ID is not indexable. Don't - // add any new mapping to the index and extent maps. This is fine since when - // a resize shows up as part of root to logical transformations, the input - // to the resize is not indexed as the indexing is done using the logical - // domain. This could be an issue when a resize is shows up outside of - // rfactor transfomations, but currently that only can happen when a - // producer tensor is transformed to look like a consumer. Since inlining is - // not allowed with resize, the out ID should never be a zero domain in that - // case. - return; - } else { - index_map_[in_id] = sub(out_ind, resize->leftExpand()); - extent_map_[in_id] = sub( - sub(getExtent(out_id), resize->leftExpand()), resize->rightExpand()); - } -} - -void IndexCompute::dispatch(Expr* e) { - NVF_ERROR( - (e->isOneOf()), - "Invalid expr type found in transform traversal."); - updateUnswitchedDomains(e); - BackwardVisitor::dispatch(e); -} - -IndexCompute::IndexCompute( - const TensorDomain* _td, - std::unordered_map initial_index_map, - std::unordered_map extent_map, - std::unordered_set zero_domains, - std::unordered_set zero_merged_in, - std::unordered_set preferred_paths) - : IndexCompute( - _td, - std::move(initial_index_map), - std::move(extent_map), - std::move(zero_domains), - std::move(zero_merged_in), - ContigIDs::getNonContigIDs(), - std::move(preferred_paths)) {} - -IndexCompute::IndexCompute( - const TensorDomain* _td, - std::unordered_map initial_index_map, - std::unordered_map extent_map, - std::unordered_set zero_domains, - std::unordered_set zero_merged_in, - const ContigIDs& contig_finder, - std::unordered_set preferred_paths, - std::unordered_set unswitched_loop_domains) - : td_(_td), - index_map_(std::move(initial_index_map)), - extent_map_(std::move(extent_map)), - zero_domains_(std::move(zero_domains)), - zero_merged_in_(std::move(zero_merged_in)), - contig_ids_{contig_finder.contigIDs()}, - preferred_paths_(std::move(preferred_paths)), - unswitched_loop_domains_(std::move(unswitched_loop_domains)) { - FUSER_PERF_SCOPE("GpuLower::Lower::IndexCompute::IndexCompute"); - - if (isOptionDisabled(DisableOption::ContigIndexing)) { - contig_ids_.clear(); - } - - // Make sure we recompute any indices we can that map to a contiguous access - // in physical memory. - const auto& within_contig = contig_finder.withinContigIDs(); - for (auto contig_id : contig_ids_) { - if (index_map_.find(contig_id) != index_map_.end()) { - NVF_ERROR(within_contig.find(contig_id) != within_contig.end()); - for (auto id : within_contig.at(contig_id)) { - index_map_.erase(id); - } - } - } - - initializeUnswitchDomainMap(); -} - -IndexCompute::IndexCompute( - std::unordered_map initial_index_map, - std::unordered_set zero_domains, - std::unordered_set preferred_paths, - std::unordered_set unswitched_loop_domains) - : td_{nullptr}, - index_map_(std::move(initial_index_map)), - zero_domains_(std::move(zero_domains)), - preferred_paths_(std::move(preferred_paths)), - concrete_id_pass_{true}, - swizzle_mode_{SwizzleMode::Loop}, - unswitched_loop_domains_(std::move(unswitched_loop_domains)) { - FUSER_PERF_SCOPE("GpuLower::Lower::IndexCompute::IndexCompute"); - initializeUnswitchDomainMap(); -} - -void IndexCompute::run(const LoopIndexing& loop_indexing) { - NVF_ERROR(concrete_id_pass_, "concrete pass only for this option"); - // Apply loop swizzles if there are any that outputs to - // the loop domains. - // Currently only support loop swizzles that directly output - // to concrete loop domains and these are validated in - // validate swizzle pass. - // TODO: - // will gradually enable replaying and mapping of loop - // swizzles in the IR infrastructure and once that's piped - // through this part of logic will be removed. - - // Resolve the index vals that could be resolved with only - // the loops that consumer_tv doesn't share with any of its - // consumers, i.e. the not-inlined loops that define consumer_tv - // values. - collectIndexIntoPermissiveMap(loop_indexing); - - // Run through the loop indexing expressions and generate - // the indexing integer math for the concrete ids. - for (auto expr : loop_indexing.getBackwardExprList()) { - // Resolve missing values from permissive map. - updateIndexMapFromPermissiveMap(expr); - - dispatch(expr); - } -} - -void IndexCompute::collectIndexIntoPermissiveMap( - const LoopIndexing& loop_indexing) { - // Visit the expressions that only produces un-inlined iterdomains, - // in reverse topological order. - for (auto expr : loop_indexing.getBackwardOutOfLineExprList()) { - // Compute indexing vals for the expression inputs. - // - // This stage should run before any indexing computation so it could be - // made sure that all index values computed at this stage are - // the ones that can be resolved only with the not-inlined - // iterdomains. - // - auto id_outputs = ir_utils::filterByType(expr->outputs()); - if (std::ranges::all_of(id_outputs, [this](IterDomain* id) { - return index_map_.count( - GpuLower::current()->info().caMap().getConcreteMappedID( - id, IdMappingMode::EXACT)); - })) { - // Visit this expression: - // LoopIndexingAnalysis::traverseFromDomainVals made sure that each - // concrete index is bound exactly once so computing these expressions - // early should still be consistent. - dispatch(expr); - - auto id_inputs = ir_utils::filterByType(expr->inputs()); - for (auto id : id_inputs) { - // Collect backward pass results from this expression if they are - // made available in by this expression. - auto idx_it = index_map_.find( - GpuLower::current()->info().caMap().getConcreteMappedID( - id, IdMappingMode::EXACT)); - - if (idx_it != index_map_.end()) { - permissive_index_map_ - [GpuLower::current()->info().caMap().getConcreteMappedID( - id, IdMappingMode::PERMISSIVE)] = idx_it->second; - } - } - } - } -} - -void IndexCompute::updateIndexMapFromPermissiveMap(const Expr* id_expr) { - auto id_outputs = ir_utils::filterByType(id_expr->outputs()); - for (auto id : id_outputs) { - auto concrete_id = GpuLower::current()->info().caMap().getConcreteMappedID( - id, IdMappingMode::EXACT); - // Only try to copy index val from permissive map when - // the index is missing. - if (!index_map_.count(concrete_id)) { - auto permissive_id = - GpuLower::current()->info().caMap().getConcreteMappedID( - id, IdMappingMode::PERMISSIVE); - // Write the permissive index val into index_map_ if the - // missing value is found here. - auto permissive_it = permissive_index_map_.find(permissive_id); - if (permissive_it != permissive_index_map_.end()) { - index_map_[concrete_id] = permissive_it->second; - } - } - } -} - -void IndexCompute::run() { - const std::vector domain_vals(td_->loop().begin(), td_->loop().end()); - traverseTo(domain_vals, false); -} - -IterDomain* IndexCompute::maybeGetExactMapConcreteID(IterDomain* id) const { - if (concrete_id_pass_) { - return GpuLower::current()->info().caMap().getConcreteMappedID( - id, IdMappingMode::EXACT); - } - return id; -} - -Val* IndexCompute::getExtent(IterDomain* id) const { - // Pick from extent_map_ if available. Previously parallel - // dimensions were ued (e.g., blockDim.x), however, it would result - // in out-of-bounds errors when the extent of IterDomain is smaller - // than the threading dimension. - if (extent_map_.find(id) != extent_map_.end()) { - return extent_map_.at(id); - } else { - return id->extent(); - } -} - -bool IndexCompute::hasZeroMerged(IterDomain* id) const { - return zero_merged_in_.find(id) != zero_merged_in_.end() || isZero(id); -} - -bool IndexCompute::isZero(IterDomain* id) const { - return zero_domains_.find(id) != zero_domains_.end(); -} - -IndexCompute IndexCompute::updateIndexCompute( - const TensorDomain* new_td, - const std::unordered_map>& - id_map, - const ContigIDs& contig_finder) const { - FUSER_PERF_SCOPE("GpuLower::Lower::updateIndexCompute"); - - std::unordered_map updated_index_map; - std::unordered_map updated_extent_map; - std::unordered_set updated_zero_domains; - std::unordered_set updated_zero_merged_in; - std::unordered_set updated_unswitched_domains; - - // Multile IDs can map to the same ID, so loop over the mappings in - // a deterministic order to have deterministic indexing results - for (auto prev_id : getSortedKeys(id_map, Statement::lessThan)) { - const auto& new_ids = id_map.at(prev_id); - for (auto new_id : new_ids.vector()) { - if (index_map_.find(prev_id) != index_map_.end()) { - updated_index_map[new_id] = index_map_.at(prev_id); - } - - if (extent_map_.find(prev_id) != extent_map_.end()) { - updated_extent_map[new_id] = getExtent(prev_id); - } - - if (zero_domains_.find(prev_id) != zero_domains_.end()) { - updated_zero_domains.emplace(new_id); - } - - if (zero_merged_in_.find(prev_id) != zero_merged_in_.end()) { - updated_zero_merged_in.emplace(new_id); - } - - if (auto it = unswitched_loop_domains_.find(prev_id); - it != unswitched_loop_domains_.end()) { - updated_unswitched_domains.emplace(new_id); - } - } - } - - IndexCompute updated_index_compute( - new_td, - updated_index_map, - updated_extent_map, - updated_zero_domains, - updated_zero_merged_in, - contig_finder, - {}, - updated_unswitched_domains); - - updated_index_compute.run(); - - return updated_index_compute; -} - -namespace { -// Map indices down to the loop domains for applying swizzle -class UpdateLeafIndices : public IterVisitor { - public: - UpdateLeafIndices( - const TensorDomain* td, - std::unordered_map initial_index_map, - std::unordered_map extent_map) - : td_(td), - index_map_(std::move(initial_index_map)), - extent_map_(std::move(extent_map)) { - const std::vector domain_vals(td_->loop().begin(), td_->loop().end()); - - traverseTo(domain_vals, false); - } - - const std::unordered_map& indexMap() const { - return index_map_; - } - - const std::unordered_map& extentMap() const { - return extent_map_; - } - - private: - using IterVisitor::handle; - - void handle(Split* split) override { - auto in_id = split->in(); - auto outer_id = split->outer(); - auto inner_id = split->inner(); - - // Nothing need to be done when mappings for the output axes - // already exist. - if (index_map_.find(outer_id) != index_map_.end()) { - return; - } - - if (!index_map_.count(in_id)) { - // Reduction axes on producer side could be visited on forward - // propagation pass and current implementation does not yet - // support reduction on swizzled iterdomains, so un-indexed - // reduction iterdomains are just ignored for now. It is the same - // for broadcast iterdomains. - NVF_ERROR( - in_id->isReduction() || in_id->isBroadcast(), - "Undefined index for ", - in_id->toString()); - return; - } - - auto factor = split->factor(); - index_map_[inner_id] = - SimplifyingIrBuilder::modExpr(index_map_[in_id], factor); - extent_map_[inner_id] = factor; - index_map_[outer_id] = - SimplifyingIrBuilder::divExpr(index_map_[in_id], factor); - extent_map_[outer_id] = - SimplifyingIrBuilder::ceilDivExpr(getExtent(in_id), factor); - } - - void handle(Merge* merge) override { - auto out_id = merge->out(); - auto outer_id = merge->outer(); - auto inner_id = merge->inner(); - - // Nothing need to be done when mappings for the output axes - // already exist. - if (index_map_.find(out_id) != index_map_.end()) { - return; - } - - if (outer_id->isBroadcast()) { - if (!index_map_.count(inner_id)) { - // Reduction axes on producer side could be visited on forward - // propagation pass and current implementation does not yet - // support reduciton on swizzled iterdomains, so un-indexed - // reduction iterdomains are just ignored for now. The same applies to - // BroadcastOp. - NVF_ERROR( - inner_id->isReduction() || inner_id->isBroadcast(), - "Undefined index for ", - inner_id->toString()); - return; - } - - NVF_ERROR( - index_map_.find(inner_id) != index_map_.end(), "Inner ID not found"); - - index_map_[out_id] = index_map_[inner_id]; - extent_map_[out_id] = getExtent(inner_id); - return; - } else if (inner_id->isBroadcast()) { - if (!index_map_.count(outer_id)) { - // Reduction axes on producer side could be visited on forward - // propagation pass and current implementation does not yet - // support reduciton on swizzled iterdomains, so un-indexed - // reduction iterdomains are just ignored for now. - NVF_ERROR( - outer_id->isReduction() || outer_id->isBroadcast(), - "Undefined index for ", - outer_id->toString()); - return; - } - - NVF_ERROR( - index_map_.find(outer_id) != index_map_.end(), "Outer ID not found"); - - index_map_[out_id] = index_map_[outer_id]; - extent_map_[out_id] = getExtent(outer_id); - return; - } - - if (!index_map_.count(outer_id) || !index_map_.count(inner_id)) { - // Reduction axes on producer side could be visited on forward - // propagation pass and current implementation does not yet - // support reduciton on swizzled iterdomains, so un-indexed - // reduction iterdomains are just ignored for now. - NVF_ERROR( - (outer_id->isReduction() || outer_id->isBroadcast()) && - (inner_id->isReduction() || inner_id->isBroadcast()), - "Undefined index for ", - outer_id->toString(), - " and ", - inner_id->toString()); - return; - } - - NVF_ERROR( - index_map_.find(outer_id) != index_map_.end(), "Outer ID not found"); - NVF_ERROR( - index_map_.find(inner_id) != index_map_.end(), "Inner ID not found"); - - index_map_[out_id] = SimplifyingIrBuilder::addExpr( - index_map_[inner_id], - SimplifyingIrBuilder::mulExpr( - index_map_[outer_id], getExtent(inner_id))); - - extent_map_[out_id] = - SimplifyingIrBuilder::mulExpr(getExtent(outer_id), getExtent(inner_id)); - } - - // return extent_map_[id] if exists, else return id->extent() - Val* getExtent(IterDomain* id) { - if (extent_map_.find(id) != extent_map_.end()) { - return extent_map_.at(id); - } else { - return id->extent(); - } - } - - private: - const TensorDomain* td_; - std::unordered_map index_map_; - std::unordered_map extent_map_; -}; - -Val* getExtentOfRootAxis(IterDomain* id, Val* normal_extent = nullptr) { - // If id is device dim, ignore the extent which holds the unsharded extent. - if (id->isDeviceDim()) { - normal_extent = GpuLower::current()->kernel()->oneVal(); - } else if (normal_extent == nullptr) { - normal_extent = id->extent(); - } - - return normal_extent; -} - -} // namespace - namespace { -//! Check if the index of a parallel loop should be substituted with -//! zero. -//! -//! Zero substitution only happens with the BID parallel types with -//! Local Or Shared tensors or the TID parallel types with Local -//! tensors. -//! -//! This check is straightforward for consumers, but for producers -//! the substitution is only done when the producer uses the same -//! parallel type as the loop parallel type. -//! -//! If there's a mapped producer IterDoamin and that ID is -//! parallelized, there are a couple of cases depending on the -//! parallel type and the producer memory type: -//! -//! Loop PT, producer PT, producer mem -> index -//! - BID, TID/Serial, Shared / Local -> use BID -//! - BID, BID, Shared / Local -> use zero when loop PT == producer PT -//! - BID, BID, Shared / Local -> invalid when loop PT != producer PT -//! - TID, Serial, Local -> use TID -//! - TID, TID, Local -> use zero when loop PT == producer PT -//! - TID, TID, Local -> invalid when loop PT != producer PT -//! -//! The invalid cases should not happen here as they should be already -//! detected as invalid parallelization. Thus, we just need to find if -//! the producer has a mapped IterDomain that has the same parallel -//! type as the loop IterDomain. -bool isParallelLoopIndexSubstitutedAsZero( - const TensorView* tv, - IterDomain* loop_id, - bool as_consumer, - bool within_mma_loops) { - const auto& ca_map = GpuLower::current()->info().caMap(); - - // MMA operands are currently indexed in units of "fragments", - // so each mma tensor domain would be zero-ed and the tensor index - // calculated here would be the fragment index. - // TODO: This is a quick WAR to enable iterating over a register array - // of MMA fragments, so we could generate unrolled mma loops. - // Eventually we still want IdGraph to be able to analyze the - // in-register layout of mma fragments for more unified indexing math - // as well as more flexibility in swizzling loops. - if (loop_id->isMma() && !as_consumer) { - return true; - } - - const bool is_shared = tv->getMemoryType() == MemoryType::Shared; - const bool is_local = tv->getMemoryType() == MemoryType::Local; - - if (!((loop_id->isBlockDim() && (is_shared || is_local)) || - (loop_id->isThread() && is_local))) { - return false; - } - - // If this is for a consumer, the above check is sufficient - if (as_consumer) { - return true; - } - - // Note && TODO: - // mma swizzled lane_id does not map naturally from producer - // to consumer but they should still be detected as same - // parallel type. In a follow up may want to extend - // find_matching_parallel_domain to cover this case. - if ((within_mma_loops || ir_utils::isLdMatrixOp(tv->definition())) && - loop_id->getParallelType() == ParallelType::TIDx) { - return true; - } - - // When indexing a producer, additional checks are required as - // mentioned above - auto producer_tv = tv; - auto it = std::find_if( - tv->getLoopDomain().begin(), - tv->getLoopDomain().end(), - [&](IterDomain* tv_id) { - // Matching is done using the index and loop maps. See - // validateParallelize as well. - return ca_map.areMapped(loop_id, tv_id, IdMappingMode::EXACT) || - ca_map.areMapped(loop_id, tv_id, IdMappingMode::PERMISSIVE); - }); - - // There's no mapped producer ID. Zero substitution shouldn't be - // done. - if (it == tv->getLoopDomain().end()) { - return false; - } - - // Producer ID that corresponds to the loop ID - IterDomain* producer_id = *it; - - // If the loop PT and producer PT are the same, the producer ID can - // be indexed as just zero. Otherwise, it must use the loop parallel - // type as its index. - - // Sanity check when not substituted, i.e., when the producer ID - // uses a different as the loop PT. Not necessary as these - // conditions are already validated, but just double checking. - - if (loop_id->getParallelType() != producer_id->getParallelType()) { - NVF_ERROR( - (loop_id->isBlockDim() && !producer_id->isBlockDim()) || - (loop_id->isThreadDim() && !producer_id->isThread()), - "Found invalid parallelization that should have been detected by the " - "parallel validation: loop ID: ", - loop_id->toString(), - ", producer: ", - producer_tv->toString()); - } - - return producer_id->getParallelType() == loop_id->getParallelType(); -} - -} // namespace - -// Used for local and shared index mapping. Returns a map from loops -// to loop indices as well as a set of loops that do not contribute to -// indexing. -std::pair< - std::unordered_map, - std::unordered_set> -indexMapFromTV( - const TensorView* tv, - const std::vector& loops, - kir::ForLoop* alloc_loop, - bool as_consumer, - kir::ForLoop* circular_buffer_loop) { - bool within_alloc = false; - if (alloc_loop == nullptr) { - within_alloc = true; - } - - const bool is_global = tv->getMemoryType() == MemoryType::Global; - const bool is_shared = tv->getMemoryType() == MemoryType::Shared; - - std::unordered_map loop_to_ind_map; - - // Check if the current op has an implicit loop implemented - // within an mma instruction. - bool within_mma_loops = std::ranges::any_of( - loops, [](kir::ForLoop* fl) { return fl->iter_domain()->isMma(); }); - - // Track domains that do not contibute to the resulting - // index. Previously, index->isZeroInt() was used to detect such - // domains, but that's not a reliable method as we may set an - // initial index to zero for unswitch. - std::unordered_set zero_loops; - - for (auto loop : loops) { - Val* idx = nullptr; - // See also LoopNestGenerator::pushAlloc. - // NOLINTNEXTLINE(bugprone-branch-clone) - if (!within_alloc) { - if ((loop->iter_domain()->isThreadDim() && is_shared) || - (loop->iter_domain()->isThread() && is_global)) { - idx = loop->indexOrStartIfTrivial(); - } else { - idx = GpuLower::current()->kernel()->zeroVal(); - zero_loops.insert(loop); - } - } else if (isParallelLoopIndexSubstitutedAsZero( - tv, loop->iter_domain(), as_consumer, within_mma_loops)) { - idx = GpuLower::current()->kernel()->zeroVal(); - zero_loops.insert(loop); - } else { - idx = loop->indexOrStartIfTrivial(); - } - - if (loop == circular_buffer_loop) { - const int64_t prefetch_distance = - GpuLower::current() - ->circularBufferInfo() - .getCircularBufferOptionsFor(loop->iter_domain()) - .prefetch; - idx = SimplifyingIrBuilder::addExpr( - idx, - SimplifyingIrBuilder::create( - prefetch_distance, DataType::Index)); - } - - loop_to_ind_map[loop] = idx; - - if (!within_alloc && loop == alloc_loop) { - within_alloc = true; - } - } - // NOLINTNEXTLINE(clang-analyzer-cplusplus.NewDeleteLeaks) - return {loop_to_ind_map, zero_loops}; -} - -//! Set "pragma unroll" required for loops that indexing of Local -//! tensors depends on. -//! -//! \param tv Indexed tensor -//! \param alloc_loop Allocation loop of tv -//! \param loops The current loop structure -//! \param id_map Producer-to-consumer map in case of indexing as producer -void ensureStaticIndexing( - const TensorView* tv, - kir::ForLoop* alloc_loop, - const std::vector& loops, - const std::unordered_map& id_map) { - if (tv->getMemoryType() != MemoryType::Local) { - return; - } - - bool within_alloc = false; - if (alloc_loop == nullptr) { - within_alloc = true; - } - - for (auto loop : loops) { - if (!within_alloc) { - if (loop == alloc_loop) { - within_alloc = true; - } - continue; - } - IterDomain* loop_id = loop->iter_domain(); - if (loop->vectorize() || - nvfuser::ir_utils::isMemoryPartitionedAcross( - tv->getMemoryType(), loop_id->getParallelType())) { - continue; - } - // Look for a domain that is mapped with the loop. If mapped in - // the loop map, the loop index should be used for indexing of the - // tensor, except for broadcast and reduction domains. - auto it = std::find_if( - tv->getLoopDomain().begin(), - tv->getLoopDomain().end(), - [loop_id, &id_map](IterDomain* id) { - if (id->isBroadcast() || id->isReduction() || id->isStride()) { - return false; - } - auto id_replacement = id_map.find(id); - if (id_replacement != id_map.end()) { - id = id_replacement->second; - } - return GpuLower::current()->info().caMap().areMapped( - loop_id, id, IdMappingMode::PERMISSIVE); - }); - if (it != tv->getLoopDomain().end()) { - loop->requireUnroll(); - } - } -} - -namespace { - -std::unordered_map invertOneToOneMap( - const std::unordered_map& map) { - std::unordered_map inverted; - for (const auto& kv : map) { - bool inserted = inverted.emplace(kv.second, kv.first).second; - NVF_ERROR( - inserted, - "Multiple mappings to the same value detected: ", - kv.second->toString()); - } - return inverted; -} - -} // namespace - -namespace { - -// Maps all producer domains to consumer with broadcast -// forwarding. Used to find the allocation position. -std::unordered_map mapAllProducerDomainsToConsumer( - TensorView* producer_tv, - const TensorView* consumer_tv) { - // This map has forwarded broadcast axes, it should only be used to compute - // the allocation position of the producer - std::unordered_map p2c_alloc_map; - - // We want to replay producer as consumer instead of the other way around - // since consumer may have some broadcasted axes producer doesn't have - // merged into loops producer may use. If we did consumer as producer we - // wouldn't have this information in the mapping. - auto replay_PasC = BestEffortReplay::replayPasC( - producer_tv, - consumer_tv, - -1, - PairwiseLogicalDomainMap(producer_tv, consumer_tv)); - - // Grab consumer domain entries and reverse replay map. TODO: Maybe - // TransformReplay::replayPasC could return this map - for (auto id : consumer_tv->getLoopDomain()) { - const auto& c2p_map = replay_PasC.getReplay(); - auto c2p_it = c2p_map.find(id); - if (c2p_it != c2p_map.end()) { - auto c_id = c2p_it->first; - auto p_id = c2p_it->second; - p2c_alloc_map[p_id] = c_id; - } - } - - return p2c_alloc_map; -} - Val* sumVals(std::vector vals) { Val* result_index = GpuLower::current()->kernel()->zeroVal(); for (auto v : vals) { @@ -1201,59 +90,6 @@ std::vector Index::getProducerPerDimLogicalIndex( loops); } -std::vector Index::getStrides(TensorView* tv) { - // Indices should now be mapped onto IterDomains in consumer, so just grab - // and use them. - const auto& alloc_dom = tv->getMaybeAllocationDomain(); - - std::vector strides( - alloc_dom.size(), GpuLower::current()->kernel()->oneVal()); - { - int stride_i = 0; - for (const auto i : arange(alloc_dom.size())) { - if (alloc_dom[i]->isReduction() || alloc_dom[i]->isStride()) { - strides[i] = GpuLower::current()->kernel()->oneVal(); - continue; - } - strides[i] = IrBuilder::getItemExpr( - IrBuilder::getAttrExpr(IrBuilder::metadataExpr(tv), "alloc_stride"), - (int64_t)stride_i++); - } - } - - NVF_ERROR(alloc_dom.size() == tv->domain()->contiguity().size()); - Val* cur_contig_stride = GpuLower::current()->kernel()->oneVal(); - for (const auto i : arange(alloc_dom.size())) { - auto dim = alloc_dom.size() - i - 1; - if (alloc_dom[dim]->isReduction() || alloc_dom[dim]->isStride()) { - continue; - } - - auto dim_contiguity = tv->domain()->contiguity().at(dim); - if (alloc_dom[dim]->isBroadcast()) { - strides[dim] = cur_contig_stride->fusion()->zeroVal(); - NVF_ERROR(!dim_contiguity.has_value()); - } else if (!dim_contiguity.has_value()) { - NVF_THROW("Expected value for dimension contiguity"); - } else if (dim_contiguity.value()) { - // If contig, used the stored stride which may be the previous - // dimensions stride * previous dimensions size - strides[dim] = cur_contig_stride; - // Prepare for the next dimension which may also be contiguous, multiply - // by extent of this dimension - auto alloc_dim_extent = getExtentOfRootAxis(alloc_dom[dim]); - cur_contig_stride = - SimplifyingIrBuilder::mulExpr(cur_contig_stride, alloc_dim_extent); - } else { - // If non contiguous dimension, keep local stride information, set cur - // stride to local stride * local raw extent - cur_contig_stride = SimplifyingIrBuilder::mulExpr( - strides[dim], getExtentOfRootAxis(alloc_dom[dim])); - } - } - return strides; -} - // Producer is the inputs of an expression kir::TensorIndex* Index::getProducerIndex( TensorView* producer, @@ -1340,330 +176,6 @@ kir::TensorIndex* Index::getConsumerIndex( consumer, index, as_type); } -namespace { - -// Find iteration domains in the history of a consumer to predicate comprised -// only of merge operations. Only return iteration domains that are subsequently -// fed into a split, or are in the provided domain. In other words, we don't -// want to return every IterDomain that's contiguous, just the one closest to -// the loop domain. Predicates are not associated with physical memory so we can -// treat all of them as contiguous merges. -// -// TODO: This seems to have a large overlap with ContigIDs. Consider -// refactoring. -std::vector getPredicateContigIds( - TensorView* consumer_tv, - const std::unordered_map& consumer_index_map) { - const auto gpu_lower = GpuLower::current(); - - // When there's a resize expr between the root and the logical - // domains, predicate the logical domain. Otherwise, predicate the - // root domain. The actual size of an IterDomain after resize - // changes, and the output IterDomain needs to be used to generate - // its predicate. - const auto& consumer_root_domain = ir_utils::hasResizedRfactor(consumer_tv) - ? consumer_tv->getLogicalDomain() - : consumer_tv->getMaybeRootDomain(); - - if (consumer_root_domain.empty()) { - return std::vector(); - } - - std::unordered_map concrete_index_map; - for (auto entry : consumer_index_map) { - auto c_id = gpu_lower->info().caMap().getConcreteMappedID( - entry.first, IdMappingMode::EXACT); - concrete_index_map[c_id] = entry.second; - } - - std::unordered_set final_ids; - for (auto root_i : arange(consumer_root_domain.size())) { - auto root_id = consumer_root_domain[root_i]; - if (root_id->maybePartial()) { - final_ids.insert(root_id); - continue; - } - } - - ContigIDs contig_finder( - consumer_tv->getLoopDomain(), - consumer_root_domain, - TensorDomain::getContiguityFilledWith(consumer_root_domain, true), - final_ids, - concrete_index_map, - GpuLower::current()->divisibleSplitSet(), - &GpuLower::current()->info().caMap(), - &GpuLower::current()->info().concretizedBroadcastDomains(), - {}, - false, - true); - - std::vector contig_id_infos; - std::unordered_set covered_roots; - - // Create entries and return them - for (auto root_id : consumer_root_domain) { - if (covered_roots.count(root_id) > 0) { - continue; - } - - if (root_id->isBroadcast()) { - continue; - } - - auto contig_id_it = contig_finder.allocToIndexedID().find(root_id); - - NVF_ERROR( - contig_id_it != contig_finder.allocToIndexedID().end(), - "Error in predicate contiguity analysis, missing index for root ", - root_id->toString()); - - auto contig_id = contig_id_it->second; - - // Pick inputs from the starting domains, i.e., - // reference_predicated_root_domain. - auto contig_alloc_ids = contig_finder.indexedAllocIDs(contig_id); - covered_roots.insert(contig_alloc_ids.begin(), contig_alloc_ids.end()); - PredicateDomainInfo contig_id_info; - contig_id_info.id = contig_id; - contig_id_info.covered_ids = std::unordered_set( - contig_alloc_ids.begin(), contig_alloc_ids.end()); - contig_id_infos.push_back(contig_id_info); - } - return contig_id_infos; -} - -// Get the start and stop limit offsets that define the valid range to -// compute. In the simplest case, they are just 0 and -// IterDomain::extent. However, IterDomain may have non-zero start and -// stop that's different from extent. -std::pair getStartAndStopLimitOffsets(IterDomain* consumer_id) { - NVF_ERROR(consumer_id != nullptr); - - Val* start_limit = consumer_id->start(); - Val* stop_limit = SimplifyingIrBuilder::negExpr(consumer_id->stopOffset()); - - return {start_limit, stop_limit}; -} - -// Get the offsets for the start and stop predicates. The offsets -// are to be added to the index. -std::pair getStartAndStopOffsets( - IterDomain* consumer_id, - TensorView* consumer_tv, - const std::unordered_map& consumer_start_index_map, - const std::unordered_map& consumer_stop_index_map, - bool unswitch, - bool intermediate_domain_pred) { - // By default, the offsets for the start and stop predicates are - // just zero. All halo-related adjustments are done at root domains, - // so consumer_id is not a root domain, no adjustment is required. - if (consumer_id->definition() != nullptr && !intermediate_domain_pred) { - return { - GpuLower::current()->kernel()->zeroVal(), - GpuLower::current()->kernel()->zeroVal()}; - } - - // Get the boundaries of two ends - auto limits = getStartAndStopLimitOffsets(consumer_id); - - // At this point, we have everything to create both start and stop - // predicates as: - // - // index + start_offset >= start_limit - // index + stop_offset < extent + stop_limit - // - // start_offset and stop_limit are both zero (was not the case with shift) - // - // In order to enable consolidating unswitch predicates, organize - // the predicates as: - // - // index + (start_offset - start_limit) >= 0 - // index + (stop_offset - stop_limit) < extent - - auto start_offset = SimplifyingIrBuilder::negExpr(limits.first); - auto stop_offset = SimplifyingIrBuilder::negExpr(limits.second); - - return {start_offset, stop_offset}; -} - -// Updates a loop index map with a loop index protected by magic zero -std::unordered_map updateInitialLoopIndexMap( - const std::unordered_map& initial_loop_index_map, - const IndexMagicZeroInfo& magic_zero_info) { - if (magic_zero_info.original_loop_index != nullptr) { - NVF_ERROR(magic_zero_info.protected_loop_index != nullptr); - auto concrete_loop_id = - GpuLower::current()->info().caMap().getConcreteMappedID( - magic_zero_info.loop_id, IdMappingMode::EXACT); - auto updated_map = initial_loop_index_map; - updated_map[concrete_loop_id] = magic_zero_info.protected_loop_index; - return updated_map; - } else { - return initial_loop_index_map; - } -} - -} // namespace - -std::vector getNonDivisibleConsumerDomainsToPredicate( - TensorView* consumer_tv) { - const auto& non_divisible_split_info = - GpuLower::current()->nonDivisibleSplitInfo(); - - std::vector pred_info_vec; - - auto it = non_divisible_split_info.splitsToPredicate().find(consumer_tv); - if (it == non_divisible_split_info.splitsToPredicate().end()) { - return {}; - } - - const auto& splits_to_predicate = it->second; - - for (auto split : splits_to_predicate) { - PredicateDomainInfo info{ - .id = split->in(), - .covered_ids = {split->in()}, - .is_intermediate_domain = true}; - pred_info_vec.emplace_back(info); - } - - return pred_info_vec; -} - -// Returns predicates and the concrete (by loop map) root domains they cover -std::vector Index::getReferenceRootPredicates( - TensorView* consumer_tv, - const std::vector& loops, - kir::ForLoop* unswitch_or_vec_loop) { - FUSER_PERF_SCOPE("GpuLower::Lower::Index::getReferenceRootPredicates"); - - const auto gpu_lower = GpuLower::current(); - - const bool is_unswitch = unswitch_or_vec_loop != nullptr; - - auto db_axis = - gpu_lower->circularBufferInfo().getCircularBufferAxis(consumer_tv); - - // Generate start and stop indexing from idgraph. - // - // Both start and stop positions may need to be predicated. Indexing - // differs when generating predicates for unswitch. - // NOTE: If we could find-and-replace KIR nodes, we could just - // generate one index map, clone it and replace the loop-to-index - // mappings of unswitched loops for the start predicate. - - auto stop_indexing_from_idgraph = getPredicateIndexingFromIdGraph( - loops, consumer_tv, unswitch_or_vec_loop, db_axis, false); - const auto consumer_stop_indexing = stop_indexing_from_idgraph.index; - const auto& consumer_stop_index_map = consumer_stop_indexing.indexMap(); - - // If not unswitch, share the same indexing map as the stop index - // map - const auto start_indexing_from_idgraph = is_unswitch - ? getPredicateIndexingFromIdGraph( - loops, consumer_tv, unswitch_or_vec_loop, db_axis, true) - : stop_indexing_from_idgraph; - const auto consumer_start_indexing = start_indexing_from_idgraph.index; - const auto& consumer_start_index_map = consumer_start_indexing.indexMap(); - - // Get the contiguous ids we need to generate predicates for - auto contig_id_infos = - getPredicateContigIds(consumer_tv, consumer_stop_index_map); - - auto non_divisible_splits = - getNonDivisibleConsumerDomainsToPredicate(consumer_tv); - contig_id_infos.insert( - contig_id_infos.end(), - non_divisible_splits.begin(), - non_divisible_splits.end()); - - std::vector pred_info_vec; - - for (const auto& contig_id_entry : contig_id_infos) { - auto contig_id = contig_id_entry.id; - // No predicates needed for braodcasted indices. - if (contig_id->isBroadcast()) { - continue; - } - - auto root_ids = contig_id_entry.covered_ids; - - const auto consumer_stop_indexing_it = - consumer_stop_index_map.find(contig_id); - - // First condition below happens with Misaligned predicates, where - // inner-most vectorized loops are not included in the loops - // parameter. Predicates involving vectorized loops are separately - // generated in lower_misaligned_vectorization. - // - // Can not omit stop index even if it is zero. This is important for empty - // tensor support, because in empty tensor the extent of an ID can be zero - if (consumer_stop_indexing_it == consumer_stop_index_map.end()) { - continue; - } - - PredicateInfo info; - - // The final predicates will look like: - // (index + start_offset) >= 0 && (index + stop_offset) < extent. - - std::tie(info.start_offset_, info.stop_offset_) = getStartAndStopOffsets( - contig_id, - consumer_tv, - consumer_start_index_map, - consumer_stop_index_map, - unswitch_or_vec_loop != nullptr, - contig_id_entry.is_intermediate_domain); - - auto stop_index = consumer_stop_indexing_it->second; - auto start_index = consumer_start_index_map.at(contig_id); - - IndexMagicZeroInfo start_magic_zero_info; - IndexMagicZeroInfo stop_magic_zero_info; - - // When the start and stop indices are not the same, apply the - // magic-zero protection separately for both of them. - if (stop_index != start_index) { - start_magic_zero_info = protectPredicateIndexWithMagicZero( - start_index, start_indexing_from_idgraph, loops); - stop_magic_zero_info = protectPredicateIndexWithMagicZero( - stop_index, stop_indexing_from_idgraph, loops); - } else { - stop_magic_zero_info = protectPredicateIndexWithMagicZero( - stop_index, stop_indexing_from_idgraph, loops); - start_magic_zero_info = stop_magic_zero_info; - } - - start_index = start_magic_zero_info.index; - stop_index = stop_magic_zero_info.index; - - // Build predicates for start positions as: - // start_index + start_offset >= 0 - auto offsetted_start_index = - SimplifyingIrBuilder::addExpr(start_index, info.start_offset_); - auto start_pred = SimplifyingIrBuilder::geExpr( - offsetted_start_index, GpuLower::current()->kernel()->zeroVal()); - info.start_predicate_ = start_pred; - - // Build predicates for stop positions as: - // stop_index + stop_offset < IterDomain::extent - auto stop_offset = info.stop_offset_; - auto offsetted_stop_index = - SimplifyingIrBuilder::addExpr(stop_index, stop_offset); - auto stop_pred = - SimplifyingIrBuilder::ltExpr(offsetted_stop_index, contig_id->extent()); - info.stop_predicate_ = stop_pred; - - for (auto consumer_id : contig_id_entry.covered_ids) { - info.predicated_domains_.insert(consumer_id); - } - pred_info_vec.emplace_back(info); - } - - return pred_info_vec; -} - PredicateInfo PredicateInfo::getFalseInfo() { PredicateInfo info; info.start_predicate_ = GpuLower::current()->kernel()->falseVal(); @@ -1743,7 +255,8 @@ std::pair Index::getCpAsyncBulkGmemIndex( } else { std::stringstream ss; ss << "Hopper::CpAsyncBulkS2GIndex"; - auto gmem_address = getConsumerIndex(consumer_tv, loops, {}, true); + auto gmem_address = + getConsumerIndex(consumer_tv, loops, {}, true); index = IrBuilder::structExpr( {{"raw_gmem_addr", gmem_address}, {"bytes", expected_bytes}}, ss.str()); diff --git a/csrc/index_compute.h b/csrc/index_compute.h index 3d37a7dc592..59ae778561c 100644 --- a/csrc/index_compute.h +++ b/csrc/index_compute.h @@ -8,326 +8,18 @@ #pragma once #include -#include -#include #include #include #include -/* - * Index compute takes in a list of indices typically generated from the - * surrounding for loop nest. The number of indicies are intended to match the - * number of dimensions of the incomming TensorView which may have less or more - * dimensions than its allocation domain due to split/merge operations. - * Split/merge operations are then replayed backwards produce resulting - * indices (based on input indices) that match the allocation dimension. - * - * For example with GLOBAL tensor: - * TV[I, K] - * TV[Io, Ii{4}, K] = TV.split(I, factor=4) - * ALLOC: NONE - * INDEX: indexCompute {i, j, k} -> {i * 4 + j, k} - * FLATTENED_INDEX: {i * 4 + j, k} -> {(i * 4 + j) * K + k} - * PREDICATE: {i * 4 + j, k} -> i * 4 + j < I - * - * - * For example with SHARED tensor: - * - * global_TV[I, K] - * global_TV[Io, Ii{4}, K] = global_TV.split(I, factor=4) - * smem_TV.compute_at(global_TV, 1) - * global_TV.parallelize(1, threadIDx.x) - * - * ALLOC: alloc(smem_TV, 4 x K) - * INDEX: indexCompute(smem_TV, {threadIdx.x, k}) -> {threadIdx.x, k} - * FLATTENED_INDEX: {threadIdx.x * 4 + j, k} -> {(threadIdx.x * 4 + j) * K + k} - * PREDICATE: {threadIdx.x * 4 + j, k} -> threadIdx.x * 4 + j < I // Same as if - * global - * - * - * For example with LOCAL tensor: - * global_TV[I, K, L] - * global_TV[Io, Ii{4}, K, L] = global_TV.split(I, factor=4) - * reg_TV.compute_at(global_TV, 2) - * global_TV.parallelize(1, threadIDx.x) - * global_TV{i, j, k, l} -> { i * 4 + j, k, l } - * global_TV{ i * 4 + j, k, l } -> { (i * 4 + j) * K * L + k * L + l} - * - * ALLOC: alloc(reg_TV, K x L) - * INDEX: {k, l} -> {k, l} - * FLATTENED_INDEX: {k, l} -> {k * L + l} - * PREDICATE: i * 4 + j < I && k < K && l < L -> // Same as if global - * - * These indices can then be flattened later based on strides. - */ - namespace nvfuser { -class ContigIDs; -class LoopIndexing; -struct IndexFromIdGraph; class TensorIndexer; namespace kir { class ForLoop; } -class IndexCompute : public BackwardVisitor { - protected: - using BackwardVisitor::handle; - - void dispatch(Expr*) override; - - void handle(Split*) override; - void handle(Merge*) override; - void handle(Swizzle*) override; - void handle(Resize*) override; - - // return extent_map_[id] if exists, else return id->extent() - Val* getExtent(IterDomain* id) const; - - //! True if a domain is not used to index - bool isZero(IterDomain* id) const; - //! True if any dependent of a domain is not used to index - bool hasZeroMerged(IterDomain* id) const; - - //! Returns the concrete ID from the compute at EXACT mode map if - //! concrete_id_pass == true, otherwise returns id passed in. - //! Helps unify the expr handling logic in reference domain and concrete id - //! based traversal. - IterDomain* maybeGetExactMapConcreteID(IterDomain* id) const; - - //! (Concrete indexing pass only) - //! Collect permissive index binding from the given expression. - //! See also permissive_map_ and LoopIndexing::getBackwardOutOfLineExprList. - void collectIndexIntoPermissiveMap(const LoopIndexing& loop_indexing); - - //! (Concrete indexing pass only) - //! Iterate through id_expr's input and pull index vals from permissive - //! map, when both of the following are true: - //! 1. the output id is missing in index_map_. - //! 2. the output id is found in permissive map. - void updateIndexMapFromPermissiveMap(const Expr* id_expr); - - //! Initialize unswitched_domain_map_ from the loop unswitched - //! domains - void initializeUnswitchDomainMap(); - - //! Propagate unswitched map info from expr outputs to inputs - void updateUnswitchedDomains(Expr* expr); - - //! Query if an IterDomain has a dependent unswitched domain - bool hasUnswitchedDependentDomains(IterDomain* id) const; - - //! Query if the usual modulo propagation may be invalid for a merge - //! inner path - bool isModuloInvalidUnswitchedIndex( - IterDomain* out_concrete_id, - Val* out_ind, - Val* inner_extent) const; - - // Tensor domain we're mapping back to allocation - const TensorDomain* td_; // NOLINT - - // Map we update as we propagate backward, containing all IDs in the - // propagation. Initial indices are mapped with this map at tv->domain() - // and are back propagated to tv->getMaybeAllocationDomain(). This index_map_ - // keeps the indices at intermediate IterDomain's in that back propagation. - std::unordered_map index_map_; // NOLINT - - // Map from IterDomain to their broadcasted extent. If a TV has I0*I1 but its - // producer has B0*I1 this map will contain a mapping from the ID{B0*I1} to - // the extent I0*I1. Also contains updated extents if we merge in a 0 index. - // See zero_merged_in_. - std::unordered_map extent_map_; // NOLINT - - // Keeps track of domains that do not contribute to indexing - std::unordered_set zero_domains_; // NOLINT - - // This set keeps track of IterDomain's that have had a zero index merged into - // them. This happens if we do something like tv->axis(0)->split(4) then - // tv->computeAt(1, ...) if this tensor is in smem or lmem the backward - // indexing would be (0, i) then when we do the backward computation that zero - // and i would attempt to be merged together. We handle indices like these - // specially. - std::unordered_set zero_merged_in_; - - // IDs that are a result of contiguous merges - std::unordered_set contig_ids_; - - // Mentions if we should propagate an index down a particular IterDomain path - // if there's an option - std::unordered_set preferred_paths_; - - // Temporary flag which tells IndexCompute to use concrete id's from the exact - // map rather than the actual IDs used in the ID expressions. - bool concrete_id_pass_ = false; - - // Mode of swizzle that are activated in this index compute - // instance. Will treat swizzles of different mode as no-op. - // Currently data mode swizzles are handled same as before in IndexSwizzle - // pass, while loop mode swizzles are handled early on in concrete indexing - // pass. See also [Note on swizzle mode] - SwizzleMode swizzle_mode_ = SwizzleMode::NoSwizzle; - - // (Concrete id pass only) - // Contains the indexing math that could be resolved with only the - // iterdomains on the right of the consumer_tv's ca axis, i.e. the - // ones that corresponding to the loops that consumer_tv would not - // share with any of its consumers. - // These indexing vals should be kept separate from index_map_ and - // should only be used when the indexing traversal follows the - // order defined in LoopIndexingAnalysis::traverseFromDomainVals. - std::unordered_map permissive_index_map_; - - //! Leaf domains that have maximum index values for unswitch - //! predicates. These domains need extra adjustments when going - //! through module operations for merge inner domains as module does - //! not always guarantee to preserve the maximum-ness property - std::unordered_set unswitched_loop_domains_; - - //! Mapppings from unswitched IterDomains to their unswitched - //! domains and their inner domains. Used to figure out if a module - //! could invalidate the maximum-ness property of an unswitched index. - //! - //! Mappings are created in a bottom-up fashion from loop to root - //! such that fine-grained domain mappings are kept as much as - //! possible for making the modulo analysis most precise. - //! - //! Specifically, for the loop domains, this just maps unswitched - //! domains, i.e., those included in unswitched_loop_domains_, to - //! themselves. There'll be no mapping for those loop domains that - //! are not included in unswitched_loop_domains_. The mappings of - //! all other domains are defined based on their consumer - //! domains. By default, they are also just mapped - //! to themselves if any of the consumers are also mapped. However, - //! when a domain is the input to a split, the mappings of the split output - //! domains are tracked separately and the split input will be - //! mapped to two sets of unswitched domains, one from the inner - //! output and another from the outer output. The mapping info from - //! the inner output is propagated as is, whereas the mapping info - //! from the outer output is prepended with the inner output - //! domain so that the unswitched domain list includes its inner - //! domain. Note that the semantic of inner domains is defined based - //! on split operations since they define propagated index math. - //! - //! The reason of tracking the information from split outer domains - //! separately is to avoid adjusting the unswitched predicate index - //! as much as possible. For example, here's a common transpose - //! scheduling pattern: - //! - //! // Initial 2D tensor - //! [i0, i1] - //! // Create a square tile of 32x32 - //! -> [i0 / 32, 32, i1 / 32, 32] - //! -> [i0 / 32 * i1 / 32, 32 * 32] - //! // Factor out a small domain (commonly vectorized) - //! -> [i0 / 32 * i1 / 32, 32 * 32 / 4, 4] - //! // Factor out another domain (commonly parallelized by TIDx) - //! -> [i0 / 32 * i1 / 32, 32 * 32 / 4 / 128, 128, 4] - //! - //! Notice that the merge of "32 * 32" is not contiguous, so we need - //! to predicate its input domains by propagating index exprs - //! through the merge inner path with "% 32". If any of the final - //! loop domains are unswitched, we need to make sure the index expr - //! sent through "% 32" is the maximum for the domain of extent - //! "32". Conservatively, this can just be 31, however, that isn't - //! always strictly required. For example, suppose the innermost - //! domain of extent 4 is unswitched. Its initial index is - //! 3. Propagating it through the merge inner path as usual is - //! guaranteed to be correct. More generally, it's always the case - //! when the inner extent of a merge is divisible by the extent of - //! an unswitched output and its domains. Suppose also the third - //! innermost domain is also unswitched, its initial index is 1. Its - //! contribution through the merge inner path is zero as the initial - //! index is multiplied by the extents of its inner domains, i.e., - //! 128 and 4, and they are divisible by the extent of the merge - //! inner domain. Again, more generally, if the stride of an - //! unswitched domain is a multiple of the inner extent of the merge - //! operation producing the unswitched domain, there's no - //! contribution from the unswitched domain, so it doesn't matter if - //! it's maximum or not. - //! - //! In the above pattern, the second innermost domain is commonly - //! parallelized with TIDx. Suppose it's also unswitched. Notice - //! that there's no concern for that domain of invalding the - //! maximum-ness property as threadIdx.x is the only valid initial - //! index value for each thread. However, this is the reason we keep track - //! of the split output contributions separately. More specifically, - //! the intermediate domain of (32 * 32 / 4) will have an index of - //! (1 * 128 + threadIdx.x), and the domain of (32 * 32) will have - //! (1 * 128 * 4 + threadIdx.x * 4 + 3). As discussed above, we can - //! reason about that the first and third components of this - //! unswitched expression is safe with respect to the propagation - //! with modulo by 32. The second component is also safe as that's - //! the only valid index for the domain. If not separately tracked, - //! all we could know would be that the extent of (32 * 32) is - //! 1024. Since part of the dependent domains are parallelized the - //! propagated index is not guaranteed to be 1023, so we would need - //! to make a conservative decision to send 1023 to the merge inner - //! path. - std::unordered_map>> - unswitched_domain_map_; - - public: - const std::unordered_map& indexMap() const { - return index_map_; - } - - const std::unordered_map& extentMap() const { - return extent_map_; - } - - const std::unordered_set& zeroDomains() const { - return zero_domains_; - } - - const std::unordered_set& zeroMergedIn() const { - return zero_merged_in_; - } - - // Propagate back from _td using initial_index_map - IndexCompute( - const TensorDomain* _td, - std::unordered_map initial_index_map, - std::unordered_map _extent_map, - std::unordered_set zero_domains, - std::unordered_set _zero_merged_in, - std::unordered_set preferred_paths = {}); - - IndexCompute( - const TensorDomain* _td, - std::unordered_map initial_index_map, - std::unordered_map _extent_map, - std::unordered_set zero_domains, - std::unordered_set _zero_merged_in, - const ContigIDs& contig_finder, - std::unordered_set preferred_paths = {}, - std::unordered_set unswitched_domains = {}); - - // Entry point used for using concrete id based traversal. This traversal is - // assumed to start at loop IDs provided by initial_index_map. - IndexCompute( - std::unordered_map initial_index_map, - std::unordered_set zero_domains, - std::unordered_set preferred_paths, - std::unordered_set unswitched_domains = {}); - - // Updates index_map, extent_map, and zero_merged_in based on id_map and - // returns a new IndexCompute ready to be used. - IndexCompute updateIndexCompute( - const TensorDomain* new_td, - const std::unordered_map>& - id_map, - const ContigIDs& contig_finder) const; - - // Interface to run index traversal through loop indexing analysis result to - // be used with the entry point for concrete id based traversal. - void run(const LoopIndexing& loop_indexing); - - virtual void run(); -}; - //! Information about a predicate. By default, it corresponds to a //! single logical domain but may cover multiple logial domains due to //! contigous indexing. @@ -394,26 +86,12 @@ class PredicateInfo { }; // Simple interface for IndexCompute -// If getComputeAtAxis and more generally TensorView const model is fixed, we -// can make the below tensorviews const. class Index { - private: - // get the strides of a tensor used for the index lowering - // Delete? - static std::vector getStrides(TensorView* tv); - public: // Indexing functions // Consumer = Producer // i.e. T0 = T1... -> T0 is the consumer, T1 is the producer // Producer indexing dispatch - // The argument `generate_pointer` specifies whether to generate pointer for - // the tensor. If global tensor, then generate T1.data. If shared memory - // tensor, then use `cvta` ptx to convert shared memory address to unsigned - // int for indexing. Search `toSmem` in the codebase for additional - // information. This argument is effective only if the indexed tensor is a - // shared memory or global tensor. On other memory type, this argument will - // cause an error. static kir::TensorIndex* getProducerIndex( TensorView* producer, const TensorView* consumer, @@ -456,33 +134,6 @@ class Index { const std::vector& loops, const std::unordered_map& override_index = {}); - //! Take a consumer tensorview and loop nest and generates predicates - //! associated with the concrete roots of the loop nest. Returns a list of - //! predicates, and a list of concrete roots they're associated with. It - //! is assumed that no predicate is required if index[i] is an index - //! directly from a for loop. This will not catch all cases if we actually - //! have static size information for example: - //! - //! TV[I].split(4) - //! would produce the code: - //! for(i : I/4) - //! for(j : 4) - //! if( i * 4 + j < TV.size(0)) - //! TV[i * 4 + j]... - //! - //! However if we had TV.size[0] = 16 at "compile time" then we wouldn't - //! need the predicate. This will be caught by canOmitPredicate in the - //! predicate lowering - //! - //! unswitch_or_vec_loop is the for loop to start the unswitch like - //! predicate, this is not a bool value as if we have an unswitch loop - //! with a vectorized loop inside, we only want to base the "unswitch" - //! like predicate on the vectorized loop. - static std::vector getReferenceRootPredicates( - TensorView* consumer_tv, - const std::vector& loops, - kir::ForLoop* unswitch_or_vec_loop); - //! Compute the result for iota static Val* iota( TensorView* consumer_tv, @@ -505,33 +156,6 @@ class Index { const std::vector& loops); }; -// Used for local and shared index mapping. Returns a map from loops -// to loop indices as well as a set of loops that do not contribute to -// indexing. -// TODO: could be cleaned up further. -std::pair< - std::unordered_map, - std::unordered_set> -indexMapFromTV( - const TensorView* tv, - const std::vector& loops, - kir::ForLoop* alloc_loop, - bool as_consumer, - kir::ForLoop* circular_buffer_loop = nullptr); - -//! Set "pragma unroll" required for loops that indexing of Local -//! tensors depends on. -//! -//! \param tv Indexed tensor -//! \param alloc_loop Allocation loop of tv -//! \param loops The current loop structure -//! \param id_map Producer-to-consumer map in case of indexing as producer -void ensureStaticIndexing( - const TensorView* tv, - kir::ForLoop* alloc_loop, - const std::vector& loops, - const std::unordered_map& id_map = {}); - struct PredicateDomainInfo { public: // Iteration domain to predicate @@ -546,8 +170,4 @@ struct PredicateDomainInfo { bool is_intermediate_domain = false; }; -// Get all domains that need to be predicated due to non-divisible splits -std::vector getNonDivisibleConsumerDomainsToPredicate( - TensorView* consumer_tv); - } // namespace nvfuser diff --git a/csrc/predicate_compute.cpp b/csrc/predicate_compute.cpp index 532f2ff0d6e..0ee9fcc4280 100644 --- a/csrc/predicate_compute.cpp +++ b/csrc/predicate_compute.cpp @@ -876,14 +876,8 @@ Val* PredicateCompute::getInlinePredicate( RECORD_AND_RETURN(parallel_dom_pred); } - std::vector pred_info_vec; - if (!ir_utils::hasRootToLoopLinearTransformations(out_tv) || - GpuLower::current()->idModelOptions().isTensorIndexerEnabled()) { - pred_info_vec = - gpu_lower->tensorIndexer().getPredicates(out_tv, expr, loops); - } else { - pred_info_vec = Index::getReferenceRootPredicates(out_tv, loops, nullptr); - } + std::vector pred_info_vec = + gpu_lower->tensorIndexer().getPredicates(out_tv, expr, loops); std::vector preds; @@ -978,16 +972,9 @@ void UnswitchPredicate::predicateOn(Expr* tv_expr) { auto out_tv = ir_utils::getTvOutput(tv_expr); NVF_ERROR(out_tv != nullptr, "Missing TensorView output"); - std::vector ref_pred_info; - - if (!ir_utils::hasRootToLoopLinearTransformations(out_tv) || - GpuLower::current()->idModelOptions().isTensorIndexerEnabled()) { - ref_pred_info = gpu_lower->tensorIndexer().getPredicates( - out_tv, tv_expr, for_loops_, unrolled_loop_); - } else { - ref_pred_info = - Index::getReferenceRootPredicates(out_tv, for_loops_, unrolled_loop_); - } + std::vector ref_pred_info = + gpu_lower->tensorIndexer().getPredicates( + out_tv, tv_expr, for_loops_, unrolled_loop_); // If RootPredicateInfo has a static predicate that is more // restrictive than the current one, replace the current with the From d29bae0b56a8d660d99438fe47a9313eae4870fa Mon Sep 17 00:00:00 2001 From: Naoya Maruyama Date: Mon, 6 Apr 2026 16:27:59 -0700 Subject: [PATCH 4/7] format --- csrc/index_compute.cpp | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/csrc/index_compute.cpp b/csrc/index_compute.cpp index 84dd9c4da66..6e979cc3e35 100644 --- a/csrc/index_compute.cpp +++ b/csrc/index_compute.cpp @@ -255,8 +255,7 @@ std::pair Index::getCpAsyncBulkGmemIndex( } else { std::stringstream ss; ss << "Hopper::CpAsyncBulkS2GIndex"; - auto gmem_address = - getConsumerIndex(consumer_tv, loops, {}, true); + auto gmem_address = getConsumerIndex(consumer_tv, loops, {}, true); index = IrBuilder::structExpr( {{"raw_gmem_addr", gmem_address}, {"bytes", expected_bytes}}, ss.str()); From fc1212df4cd476c4ee1b4961ce654cddd2e5b3e7 Mon Sep 17 00:00:00 2001 From: Naoya Maruyama Date: Mon, 6 Apr 2026 16:38:46 -0700 Subject: [PATCH 5/7] missing header --- csrc/index_compute.h | 2 ++ 1 file changed, 2 insertions(+) diff --git a/csrc/index_compute.h b/csrc/index_compute.h index 59ae778561c..44c75a07919 100644 --- a/csrc/index_compute.h +++ b/csrc/index_compute.h @@ -13,6 +13,8 @@ #include #include +#include "type.h" + namespace nvfuser { class TensorIndexer; From 7d40b2fae25990c82093a29e63a92b9fc858c310 Mon Sep 17 00:00:00 2001 From: Naoya Maruyama Date: Mon, 6 Apr 2026 17:01:32 -0700 Subject: [PATCH 6/7] build fix --- csrc/index_compute.h | 1 + 1 file changed, 1 insertion(+) diff --git a/csrc/index_compute.h b/csrc/index_compute.h index 44c75a07919..c4c6c739976 100644 --- a/csrc/index_compute.h +++ b/csrc/index_compute.h @@ -13,6 +13,7 @@ #include #include +#include "ir/all_nodes.h" #include "type.h" namespace nvfuser { From 94cf3e1ff94a71e21c4551c7c9f7634fb59dd1c5 Mon Sep 17 00:00:00 2001 From: Naoya Maruyama Date: Mon, 6 Apr 2026 21:29:55 -0700 Subject: [PATCH 7/7] clang-tidy --- .../analysis/sync_information.cpp | 16 ++-- csrc/device_lower/pass/rng.cpp | 7 +- csrc/id_model/indexing_utils.h | 4 +- csrc/predicate_compute.cpp | 78 +++++++------------ 4 files changed, 40 insertions(+), 65 deletions(-) diff --git a/csrc/device_lower/analysis/sync_information.cpp b/csrc/device_lower/analysis/sync_information.cpp index f5bb93ce95b..d720b3c0310 100644 --- a/csrc/device_lower/analysis/sync_information.cpp +++ b/csrc/device_lower/analysis/sync_information.cpp @@ -553,17 +553,13 @@ std::string SyncMap::toString() const { std::stringstream ss; ss << "SyncMap:"; std::vector sorted_tvs; - std::transform( - needs_raw_sync_.begin(), - needs_raw_sync_.end(), - std::back_inserter(sorted_tvs), - [](auto kv) { return kv.first; }); - std::sort( - sorted_tvs.begin(), - sorted_tvs.end(), - [](TensorView* tv1, TensorView* tv2) { - return tv1->name() < tv2->name(); + std::ranges::transform( + needs_raw_sync_, std::back_inserter(sorted_tvs), [](auto kv) { + return kv.first; }); + std::ranges::sort(sorted_tvs, [](TensorView* tv1, TensorView* tv2) { + return tv1->name() < tv2->name(); + }); bool is_first = true; for (auto tv : sorted_tvs) { if (!is_first) { diff --git a/csrc/device_lower/pass/rng.cpp b/csrc/device_lower/pass/rng.cpp index 7a83a14a9c6..c757f9ec7b7 100644 --- a/csrc/device_lower/pass/rng.cpp +++ b/csrc/device_lower/pass/rng.cpp @@ -38,6 +38,7 @@ class RNGInserter : public kir::ExprMutator { Val* rng_subseq_ = nullptr; Val* rng_offset_ = nullptr; TensorView* rng_result_ = nullptr; + // NOLINTNEXTLINE(cppcoreguidelines-avoid-const-or-ref-data-members) const std::vector& exprs; struct InsertionInfo { @@ -192,10 +193,8 @@ std::vector addRNG(const std::vector& exprs) { // Check if magic zero was even used, if not we don't have to define it or // update it. auto kernel_exprs = GpuLower::current()->kernel()->exprs(); - const bool has_rng = - std::any_of(kernel_exprs.begin(), kernel_exprs.end(), [](Expr* expr) { - return expr->isA(); - }); + const bool has_rng = std::ranges::any_of( + kernel_exprs, [](Expr* expr) { return expr->isA(); }); if (!has_rng) { return exprs; diff --git a/csrc/id_model/indexing_utils.h b/csrc/id_model/indexing_utils.h index ed1ec32b0a8..6d77eb0db66 100644 --- a/csrc/id_model/indexing_utils.h +++ b/csrc/id_model/indexing_utils.h @@ -22,8 +22,8 @@ inline kir::ForLoop* getForLoop( IterDomain* loop_id, const std::vector& for_loops, const ValGraph& loop_graph) { - auto it = std::find_if( - for_loops.begin(), for_loops.end(), [&](kir::ForLoop* for_loop) -> bool { + auto it = + std::ranges::find_if(for_loops, [&](kir::ForLoop* for_loop) -> bool { IterDomain* for_loop_id = for_loop->iter_domain(); return loop_graph.disjointValSets().strictAreMapped( loop_id, for_loop_id); diff --git a/csrc/predicate_compute.cpp b/csrc/predicate_compute.cpp index 0ee9fcc4280..2683c2d6ad6 100644 --- a/csrc/predicate_compute.cpp +++ b/csrc/predicate_compute.cpp @@ -41,7 +41,7 @@ bool isOutputLocal(const Expr* expr) { bool ParallelizedDomainPredicate::PredicateInfo::addDomain(IterDomain* id) { auto concrete_id = lower_utils::getConcreteMappedId(id); - if (std::find(ids_.begin(), ids_.end(), concrete_id) == ids_.end()) { + if (std::ranges::find(ids_, concrete_id) == ids_.end()) { ids_.push_back(concrete_id); return true; } else { @@ -91,11 +91,8 @@ std::vector getUnswitchProtectedParallelLoopIds( std::vector loop_ids; loop_ids.reserve(loops.size()); - std::transform( - loops.begin(), - loops.end(), - std::back_inserter(loop_ids), - [&](kir::ForLoop* loop) { + std::ranges::transform( + loops, std::back_inserter(loop_ids), [&](kir::ForLoop* loop) { return getLoopPromotion(loop->iter_domain(), id_model); }); @@ -128,7 +125,7 @@ std::vector getUnswitchProtectedParallelLoopIds( for (const auto& [expr_g, dir] : predicate_path) { const auto inputs = getInputsOfExprGroup(indexing_graph, expr_g, dir); const auto outputs = getOutputsOfExprGroup(indexing_graph, expr_g, dir); - if (std::any_of(inputs.begin(), inputs.end(), [&](const ValGroup& input) { + if (std::ranges::any_of(inputs, [&](const ValGroup& input) { return non_unswitch_dep_ids.has(input); })) { // Depends on non-unswitched ids @@ -176,10 +173,9 @@ std::vector getUnswitchProtectedParallelLoopIds( // If none of the inputs depends on unswitched_loop_id and its // dependents, this expr should not matter. - if (std::none_of( - inputs.begin(), inputs.end(), [&](const ValGroup& input) { - return unswitch_dep_ids.has(input); - })) { + if (std::ranges::none_of(inputs, [&](const ValGroup& input) { + return unswitch_dep_ids.has(input); + })) { continue; } @@ -189,7 +185,7 @@ std::vector getUnswitchProtectedParallelLoopIds( // unswitched_loop_id itself. Use of unswitched_loop_id and its // dependents should not make unswitched_loop_id not fully // unswitched. - if (std::any_of(inputs.begin(), inputs.end(), [&](const ValGroup& input) { + if (std::ranges::any_of(inputs, [&](const ValGroup& input) { return non_unswitch_dep_ids.has(input) && !unswitch_dep_ids.has(input); })) { @@ -272,10 +268,8 @@ ParallelizedDomainPredicate::getPredicateMap( // the other output is assigned with the maximum index, this // predicate is sufficient even when blockDim.x > K. if (within_unswitch && - std::find( - unswitch_protected_loop_ids.begin(), - unswitch_protected_loop_ids.end(), - loop_id) != unswitch_protected_loop_ids.end()) { + std::ranges::find(unswitch_protected_loop_ids, loop_id) != + unswitch_protected_loop_ids.end()) { continue; } @@ -349,8 +343,7 @@ Val* ParallelizedDomainPredicate::getPredicate( RECORD_AND_RETURN(pred); } -UnswitchPredicateKey::UnswitchPredicateKey() - : predicated_concrete_id_(nullptr) { +UnswitchPredicateKey::UnswitchPredicateKey() { for (auto pt : kParallelTypeThreads) { parallel_concrete_ids_.insert({pt, nullptr}); } @@ -434,10 +427,8 @@ UnswitchPredicateKey::UnswitchPredicateKey( consumer_tv->getLoopDomain().end(), std::back_inserter(parallelized_consumer_loop_ids), [&](IterDomain* x) { - return std::find( - all_parallelized_consumer_ids.begin(), - all_parallelized_consumer_ids.end(), - x) != all_parallelized_consumer_ids.end(); + return std::ranges::find(all_parallelized_consumer_ids, x) != + all_parallelized_consumer_ids.end(); }); if (parallelized_consumer_loop_ids.empty()) { @@ -607,11 +598,9 @@ Val* createSingleExpressionElectSync( auto pred_map = ParallelizedDomainPredicate::getPredicateMap(pred->expr(), loops); - bool is_async_warp = - std::any_of(loops.begin(), loops.end(), [](kir::ForLoop* fl) { - return fl->circularBufferLoopStage() == - CircularBufferLoopStage::AsyncWarp; - }); + bool is_async_warp = std::ranges::any_of(loops, [](kir::ForLoop* fl) { + return fl->circularBufferLoopStage() == CircularBufferLoopStage::AsyncWarp; + }); Val* parallel_dom_pred = GpuLower::current()->kernel()->trueVal(); for (auto pt : {ParallelType::TIDx, ParallelType::TIDy, ParallelType::TIDz}) { @@ -677,11 +666,9 @@ Val* createMultipleExpressionElectSync( // Determine if warp specialized tma load expression. ParallelType async_warp_on = ParallelType::Serial; - auto async_warp_loop_it = - std::find_if(loops.begin(), loops.end(), [](kir::ForLoop* fl) { - return fl->circularBufferLoopStage() == - CircularBufferLoopStage::AsyncWarp; - }); + auto async_warp_loop_it = std::ranges::find_if(loops, [](kir::ForLoop* fl) { + return fl->circularBufferLoopStage() == CircularBufferLoopStage::AsyncWarp; + }); if (async_warp_loop_it != loops.end()) { auto circular_buffer_type = std::get( GpuLower::current() @@ -735,11 +722,9 @@ OneDimTmaPredicateInfo PredicateCompute::OneDimTmaLoadExpectArrive( // zero. std::unordered_map replace_map; const auto& loops = pred->tma1dLoadLoops(); - auto circular_loop_iter = - std::find_if(loops.begin(), loops.end(), [](kir::ForLoop* fl) { - return fl->circularBufferLoopStage() == - CircularBufferLoopStage::AsyncWarp; - }); + auto circular_loop_iter = std::ranges::find_if(loops, [](kir::ForLoop* fl) { + return fl->circularBufferLoopStage() == CircularBufferLoopStage::AsyncWarp; + }); for (auto it = circular_loop_iter; it != loops.end(); it++) { auto fl = *it; // save circular buffer loop index, will be replaced when generating @@ -747,12 +732,9 @@ OneDimTmaPredicateInfo PredicateCompute::OneDimTmaLoadExpectArrive( // tma1dLoadLoops() returns all the loops above the actual tma load expr. // skip the loops that are already in the current loop nest since their // indices are accessible. - if (std::any_of( - current_loops.begin(), - current_loops.end(), - [&](kir::ForLoop* loop) { - return loop->iter_domain() == fl->iter_domain(); - })) { + if (std::ranges::any_of(current_loops, [&](kir::ForLoop* loop) { + return loop->iter_domain() == fl->iter_domain(); + })) { one_dim_tma_pred_info.loop_indices_circular_to_predicate.push_back( fl->index()); continue; @@ -795,8 +777,8 @@ Val* PredicateCompute::OneDimTmaWaitParity( // predicate OneDimTmaLoadExpectArrive . NVF_ERROR(expr->isA()) auto inline_pred_1d_tma = one_dim_tma_pred_info.inline_pred_val; - auto circular_loop_iter = std::find_if( - current_loops.begin(), current_loops.end(), [](kir::ForLoop* fl) { + auto circular_loop_iter = + std::ranges::find_if(current_loops, [](kir::ForLoop* fl) { return fl->circularBufferLoopStage() == CircularBufferLoopStage::ComputeWarp; }); @@ -1054,10 +1036,8 @@ void UnswitchPredicate::predicateOn(Expr* tv_expr) { pending_predicates_.begin() + (int64_t)pending_predicates_.size() - 1; } else if (root_ids.size() == 1) { // If not new, try to find a corresponding MergedPredicates. - merged_pred_it = std::find_if( - pending_predicates_.begin(), - pending_predicates_.end(), - [&first_key](const auto& merged_predicates) { + merged_pred_it = std::ranges::find_if( + pending_predicates_, [&first_key](const auto& merged_predicates) { return merged_predicates.predicate_key == first_key; }); // Note: It is possible that no matching merged predicate info