From c5aab207b90f1e30ccdedd62d07e9caba7ec79e9 Mon Sep 17 00:00:00 2001 From: shmsong Date: Thu, 23 Jun 2022 14:48:06 -0700 Subject: [PATCH 01/14] add peeling attribute --- build_variables.bzl | 1 + torch/csrc/jit/codegen/cuda/kernel_ir.h | 8 ++++ .../codegen/cuda/lower_predicate_peeling.cpp | 44 +++++++++++++++++++ .../codegen/cuda/lower_predicate_peeling.h | 40 +++++++++++++++++ torch/csrc/jit/codegen/cuda/type.cpp | 16 +++++++ torch/csrc/jit/codegen/cuda/type.h | 6 +++ 6 files changed, 115 insertions(+) create mode 100644 torch/csrc/jit/codegen/cuda/lower_predicate_peeling.cpp create mode 100644 torch/csrc/jit/codegen/cuda/lower_predicate_peeling.h diff --git a/build_variables.bzl b/build_variables.bzl index c8da690886ad5..fddc7559b2057 100644 --- a/build_variables.bzl +++ b/build_variables.bzl @@ -683,6 +683,7 @@ libtorch_cuda_core_sources = [ "torch/csrc/jit/codegen/cuda/lower_fusion_simplifier.cpp", "torch/csrc/jit/codegen/cuda/lower_index.cpp", "torch/csrc/jit/codegen/cuda/lower_index_hoist.cpp", + "torch/csrc/jit/codegen/cuda/lower_predicate_peeling.cpp", "torch/csrc/jit/codegen/cuda/lower_insert_syncs.cpp", "torch/csrc/jit/codegen/cuda/lower_instrument.cpp", "torch/csrc/jit/codegen/cuda/lower_loops.cpp", diff --git a/torch/csrc/jit/codegen/cuda/kernel_ir.h b/torch/csrc/jit/codegen/cuda/kernel_ir.h index 042d4856bcec6..48f7def310ec5 100644 --- a/torch/csrc/jit/codegen/cuda/kernel_ir.h +++ b/torch/csrc/jit/codegen/cuda/kernel_ir.h @@ -429,6 +429,8 @@ struct LoopTransformInfo { DoubleBufferLoopStage double_buffer_loop_stage = DoubleBufferLoopStage::NotApplicable; + PredicatePeelStage predicate_peel_stage = PredicatePeelStage::NoApplicable; + //! Tracks if this for loop is for base index calculation for //! lifted memory address. bool is_base_index_loop = false; @@ -444,6 +446,12 @@ struct LoopTransformInfo { is_base_index_loop = true; return *this; } + + //! Setter API + LoopTransformInfo& predicatePeelStage(PredicatePeelStage stage) { + predicate_peel_stage = stage; + return *this; + } }; //! ForLoop provides scoping around an int iterator from 0 to range. Exprs diff --git a/torch/csrc/jit/codegen/cuda/lower_predicate_peeling.cpp b/torch/csrc/jit/codegen/cuda/lower_predicate_peeling.cpp new file mode 100644 index 0000000000000..fa756e8b6aa66 --- /dev/null +++ b/torch/csrc/jit/codegen/cuda/lower_predicate_peeling.cpp @@ -0,0 +1,44 @@ +#include +#include +#include +#include +#include +#include +#include + +#include + +namespace torch { +namespace jit { +namespace fuser { +namespace cuda { + +bool PredicatePeeling::supportedPeelingLoop(IterDomain* id) { + if (id->isParallelized()) { + return false; + } + + auto id_def = id->definition(); + + if (id_def == nullptr) { + // This case is not profitable so skip peeling support. + return false; + } else if (auto split = dynamic_cast(id_def)) { + auto split_in = split->in(); + + // This is typical case that we want to peel. + // where we advance a serial iteration through + // constant sized tiles. + return split_in->definition() == nullptr && id == split->outer() && + split->factor()->isConstInt(); + } + + // TODO: + // should extend to use separability defined in anohter PR. + return false; +} + +} // namespace cuda +} // namespace fuser +} // namespace jit +} // namespace torch diff --git a/torch/csrc/jit/codegen/cuda/lower_predicate_peeling.h b/torch/csrc/jit/codegen/cuda/lower_predicate_peeling.h new file mode 100644 index 0000000000000..2904627b34f15 --- /dev/null +++ b/torch/csrc/jit/codegen/cuda/lower_predicate_peeling.h @@ -0,0 +1,40 @@ +#pragma once + +#include + +#include +#include +#include + +#include + +namespace torch { +namespace jit { +namespace fuser { +namespace cuda { + +class PredicatePeelingInfo { + public: + bool shouldPeelLoop(kir::ForLoop* forloop); + + void markLoopPeeled(kir::ForLoop* forloop); + + private: + private: + std::unordered_set concrete_loops_to_peel_; +}; + +namespace PredicatePeeling { + +// User space check that makes sure the loop can +// actually be peeled to remove predicates. +bool supportedPeelingLoop(IterDomain* id); + +std::vector peelPredicatedLoop(const std::vector exprs); + +} // namespace PredicatePeeling + +} // namespace cuda +} // namespace fuser +} // namespace jit +} // namespace torch diff --git a/torch/csrc/jit/codegen/cuda/type.cpp b/torch/csrc/jit/codegen/cuda/type.cpp index b9b1231d20142..b14cdfeb7df3c 100644 --- a/torch/csrc/jit/codegen/cuda/type.cpp +++ b/torch/csrc/jit/codegen/cuda/type.cpp @@ -1199,6 +1199,22 @@ TORCH_CUDA_CU_API std::ostream& operator<<( default: TORCH_INTERNAL_ASSERT(false, "unknown double buffer stage"); } +} + +std::ostream& operator<<(std::ostream& os, PredicatePeelStage peel_stage) { + switch (peel_stage) { + case PredicatePeelStage::NoApplicable: + break; + case PredicatePeelStage::Prolog: + os << "{PeeledProlog}"; + break; + case PredicatePeelStage::Main: + os << "{PeeledMain}"; + break; + default: + TORCH_INTERNAL_ASSERT(false, "unsupported loop attribute"); + break; + } return os; } diff --git a/torch/csrc/jit/codegen/cuda/type.h b/torch/csrc/jit/codegen/cuda/type.h index 7250888ebb0e0..1fa04c519f461 100644 --- a/torch/csrc/jit/codegen/cuda/type.h +++ b/torch/csrc/jit/codegen/cuda/type.h @@ -334,6 +334,12 @@ enum class DoubleBufferLoopStage { LowerProlog }; +enum class PredicatePeelStage { + NoApplicable, + Prolog, + Main +} predicate_peel_stage; + //! Supported swizzle types, //! corresponds to swizzles functions on the runtime cuda //! naming it swizzle_2d to reserve the options to have a swizzle_1d. From ffbdb439846c463934b942f445955ebfb3395b39 Mon Sep 17 00:00:00 2001 From: shmsong Date: Sun, 26 Jun 2022 18:06:25 -0700 Subject: [PATCH 02/14] Add predicate peeling pass --- torch/csrc/jit/codegen/cuda/codegen.cpp | 9 +- torch/csrc/jit/codegen/cuda/index_compute.cpp | 58 ++++- .../jit/codegen/cuda/ir_interface_nodes.h | 7 + torch/csrc/jit/codegen/cuda/ir_iostream.cpp | 3 + torch/csrc/jit/codegen/cuda/kernel_ir.cpp | 13 + torch/csrc/jit/codegen/cuda/kernel_ir.h | 2 + torch/csrc/jit/codegen/cuda/lower2device.cpp | 7 +- torch/csrc/jit/codegen/cuda/lower2device.h | 10 + torch/csrc/jit/codegen/cuda/lower_index.cpp | 2 +- .../codegen/cuda/lower_predicate_peeling.cpp | 230 ++++++++++++++++++ .../codegen/cuda/lower_predicate_peeling.h | 22 +- .../codegen/cuda/runtime/random_numbers.cu | 8 +- torch/csrc/jit/codegen/cuda/tensor_view.cpp | 11 + torch/csrc/jit/codegen/cuda/test/test_gpu.cpp | 39 +++ torch/csrc/jit/codegen/cuda/type.cpp | 5 +- torch/csrc/jit/codegen/cuda/type.h | 9 +- 16 files changed, 418 insertions(+), 17 deletions(-) diff --git a/torch/csrc/jit/codegen/cuda/codegen.cpp b/torch/csrc/jit/codegen/cuda/codegen.cpp index 3bfbce473d5b6..498f8d407bd38 100644 --- a/torch/csrc/jit/codegen/cuda/codegen.cpp +++ b/torch/csrc/jit/codegen/cuda/codegen.cpp @@ -2289,7 +2289,14 @@ class CudaKernelGenerator : private OptOutConstDispatch { } indent() << "for(nvfuser_index_t " << gen_index; - if (loop->iter_domain()->isParallelized()) { + + // TODO: need to revisit this one, + // a predicate peeled loop would be guaranteed not to be + // a degenerate loop. So the comments on the else block + // should not apply here. + if (loop->iter_domain()->isParallelized() || + loop->loopTransformInfo().predicate_peel_stage == + PredicatePeelStage::Main) { code_ << " = " << gen_start << "; "; } else { // Do not start at the start of the ID when not parallelized. Instead, diff --git a/torch/csrc/jit/codegen/cuda/index_compute.cpp b/torch/csrc/jit/codegen/cuda/index_compute.cpp index 16345258b049c..9d2102d56ca16 100644 --- a/torch/csrc/jit/codegen/cuda/index_compute.cpp +++ b/torch/csrc/jit/codegen/cuda/index_compute.cpp @@ -1602,6 +1602,16 @@ std::vector Index::getGlobalProducerStridedIndices( if (root_ind->isZeroInt()) { continue; } else { + if (auto tile_entry = + gpu_lower->predicatePeelingInfo().getMaybePeeledTileEntry( + loops, root_dom[i])) { + if (tile_entry.value().peel_stage != PredicatePeelStage::Prolog) { + root_ind = SimplifyingIrBuilder::subExpr( + root_ind, + PredicatePeeling::getSplitTileMainOffset( + root_dom[i], tile_entry.value().inner_factor)); + } + } auto strided_ind = SimplifyingIrBuilder::mulExpr(root_ind, strides[i]); if (i == root_dom.size() - 1 && vectorize_shift != nullptr) { strided_inds[i] = @@ -2090,6 +2100,16 @@ std::vector Index::getGlobalConsumerStridedIndices( if (root_ind->isZeroInt()) { continue; } else { + if (auto tile_entry = + gpu_lower->predicatePeelingInfo().getMaybePeeledTileEntry( + loops, root_dom[i])) { + if (tile_entry.value().peel_stage != PredicatePeelStage::Prolog) { + root_ind = SimplifyingIrBuilder::subExpr( + root_ind, + PredicatePeeling::getSplitTileMainOffset( + root_dom[i], tile_entry.value().inner_factor)); + } + } auto strided_ind = SimplifyingIrBuilder::mulExpr(root_ind, strides[i]); if (i == root_dom.size() - 1 && vectorize_shift != nullptr) { strided_inds[i] = @@ -3167,10 +3187,18 @@ std::vector Index::getReferenceRootPredicates( info.start_predicate_ = start_pred; } + auto maybe_tiled_entry = + gpu_lower->predicatePeelingInfo().getMaybePeeledTileEntry( + loops, contig_id); + + bool has_peeled_prolog = maybe_tiled_entry.has_value() && + maybe_tiled_entry.value().peel_stage == PredicatePeelStage::Prolog; + // Build predicates for stop positions as: // stop_index + stop_offset < IterDomain::extent auto stop_offset = info.stop_offset_; - if (canOmitStopPredicate(stop_index, stop_offset, contig_id)) { + if (canOmitStopPredicate(stop_index, stop_offset, contig_id) && + !has_peeled_prolog) { info.stop_predicate_ = GpuLower::current()->kernel()->trueVal(); } else { auto offsetted_stop_index = @@ -3178,6 +3206,34 @@ std::vector Index::getReferenceRootPredicates( auto stop_pred = SimplifyingIrBuilder::ltExpr( offsetted_stop_index, contig_id->extent()) ->as(); + auto maybe_tiled_entry = + gpu_lower->predicatePeelingInfo().getMaybePeeledTileEntry( + loops, contig_id); + + if (maybe_tiled_entry.has_value()) { + auto tile_entry = maybe_tiled_entry.value(); + if (tile_entry.peel_stage == PredicatePeelStage::Prolog) { + stop_pred = SimplifyingIrBuilder::ltExpr( + offsetted_stop_index, + PredicatePeeling::getSplitTileOffset( + contig_id, tile_entry.inner_factor)) + ->as(); + } else if ( + tile_entry.for_loop->doubleBufferLoopStage() == + DoubleBufferLoopStage::Main) { + auto db_index = SimplifyingIrBuilder::addExpr( + tile_entry.for_loop->index(), + IrBuilder::create( + gpu_lower->doubleBufferInfo().getStageDepthFor( + tile_entry.for_loop->iter_domain()) - + 1)); + stop_pred = SimplifyingIrBuilder::ltExpr( + db_index, tile_entry.for_loop->stop()) + ->as(); + } else { + stop_pred = gpu_lower->kernel()->trueVal(); + } + } info.stop_predicate_ = stop_pred; } diff --git a/torch/csrc/jit/codegen/cuda/ir_interface_nodes.h b/torch/csrc/jit/codegen/cuda/ir_interface_nodes.h index a60555498e30a..e246d8274e0c5 100644 --- a/torch/csrc/jit/codegen/cuda/ir_interface_nodes.h +++ b/torch/csrc/jit/codegen/cuda/ir_interface_nodes.h @@ -532,6 +532,12 @@ class TORCH_CUDA_CU_API TensorView : public Val { friend TORCH_CUDA_CU_API void groupReductions( const std::vector&); + void peelPredicatedLoop(int axis_id); + + const auto& peeledSerialId() const { + return peeled_serial_id_; + } + protected: void setDomain(TensorDomain* td) { domain_ = td; @@ -570,6 +576,7 @@ class TORCH_CUDA_CU_API TensorView : public Val { //! Indicates the circular buffering stage depth if applicable. unsigned int circular_buffer_stage_ = 0; + std::vector peeled_serial_id_; // special handling for CPU based zero-dim tensors (i.e. CPU Tensors that only // have one value). This is only used if on an input value, otherwise ignored. // This is important as special handling because these "scalars" should be diff --git a/torch/csrc/jit/codegen/cuda/ir_iostream.cpp b/torch/csrc/jit/codegen/cuda/ir_iostream.cpp index cb6373428f0fa..6c6b8e02b2f29 100644 --- a/torch/csrc/jit/codegen/cuda/ir_iostream.cpp +++ b/torch/csrc/jit/codegen/cuda/ir_iostream.cpp @@ -627,6 +627,9 @@ void IrPrinter::handle(const kir::ForLoop* node) { handle(node->index()); os_ << " in "; handle(node->iter_domain()); + os_ << ", start = " << node->start(); + os_ << ", stop = " << node->stop(); + os_ << ", step = " << node->step(); os_ << ":\n"; handleScope(node->body()); } diff --git a/torch/csrc/jit/codegen/cuda/kernel_ir.cpp b/torch/csrc/jit/codegen/cuda/kernel_ir.cpp index c3fd4edca1b3d..d020144cc7c06 100644 --- a/torch/csrc/jit/codegen/cuda/kernel_ir.cpp +++ b/torch/csrc/jit/codegen/cuda/kernel_ir.cpp @@ -50,6 +50,19 @@ Predicate::Predicate(IrBuilderPasskey passkey, Bool* value) TORCH_INTERNAL_ASSERT(value != nullptr); } +Predicate::Predicate(IrBuilderPasskey passkey, const Predicate* other) + : Val(passkey, ValType::Predicate, DataType::Bool), + ptype_(other->ptype_), + expr_(other->expr_), + thread_pred_(other->thread_pred_), + unrolled_loop_(other->unrolled_loop_) { + TORCH_INTERNAL_ASSERT( + passkey.ir_container_->isA(), + "IR type only valid for Kernel container."); + TORCH_INTERNAL_ASSERT( + other->value_ == nullptr, "No support yet for predicate deep copy"); +} + TensorIndex::TensorIndex( IrBuilderPasskey passkey, const TensorView* view, diff --git a/torch/csrc/jit/codegen/cuda/kernel_ir.h b/torch/csrc/jit/codegen/cuda/kernel_ir.h index 48f7def310ec5..c1db32e026193 100644 --- a/torch/csrc/jit/codegen/cuda/kernel_ir.h +++ b/torch/csrc/jit/codegen/cuda/kernel_ir.h @@ -82,6 +82,8 @@ class TORCH_CUDA_CU_API Predicate final : public Val { explicit Predicate(IrBuilderPasskey passkey, Bool* value); + explicit Predicate(IrBuilderPasskey passkey, const Predicate* other); + PredicateType predicate_type() const { return ptype_; } diff --git a/torch/csrc/jit/codegen/cuda/lower2device.cpp b/torch/csrc/jit/codegen/cuda/lower2device.cpp index 5b85cc27b197d..dfeda60aeb20c 100644 --- a/torch/csrc/jit/codegen/cuda/lower2device.cpp +++ b/torch/csrc/jit/codegen/cuda/lower2device.cpp @@ -18,6 +18,7 @@ #include #include #include +#include #include #include #include @@ -316,6 +317,7 @@ void GpuLower::lower(Fusion* fusion, DataType index_type) { compute_at_map_->allocateIndexVariables(); addressComputeInfo().build(fusion_); + predicatePeelingInfo().build(fusion_); // Run our passes keeping the lowered expressions and forwarding // them @@ -350,12 +352,15 @@ void GpuLower::lower(Fusion* fusion, DataType index_type) { const auto exprs_double_buffered = DoubleBufferPass::run(exprs_with_precompute_address); + const auto predicate_peeled = + PredicatePeeling::peelPredicatedLoop(exprs_double_buffered); + // This pass inserts predicates as well as branches in the code. Up until now // the code is explicitly single shot for loop based. Need to be careful in // later passes when doing any kind of insertions in loop nest structure as // insertions could be on if then or else instead of directly on a for loop. const auto exprs_unrolled_loops = - UnrollPass::runPass(fusion_, exprs_double_buffered); + UnrollPass::runPass(fusion_, predicate_peeled); const auto exprs_unrolled_mv_loops = processMisalignedVectorization(exprs_unrolled_loops); diff --git a/torch/csrc/jit/codegen/cuda/lower2device.h b/torch/csrc/jit/codegen/cuda/lower2device.h index 9b96997630c6b..4f9f7ac471d97 100644 --- a/torch/csrc/jit/codegen/cuda/lower2device.h +++ b/torch/csrc/jit/codegen/cuda/lower2device.h @@ -13,6 +13,7 @@ #include #include #include +#include #include #include #include @@ -177,6 +178,14 @@ class TORCH_CUDA_CU_API GpuLower : public NonCopyable { return profile_; } + auto& predicatePeelingInfo() { + return predicate_peeling_info_; + } + + const auto& predicatePeelingInfo() const { + return predicate_peeling_info_; + } + // This is an interface to propagate information after expression // replacement on the kernel IR. E.g.: // for ... @@ -223,6 +232,7 @@ class TORCH_CUDA_CU_API GpuLower : public NonCopyable { FusedReductionInfo fused_reduction_info_; SyncMap sync_map_; AddressComputeInfo address_compute_info_; + PredicatePeelingInfo predicate_peeling_info_; kir::KernelPerformanceProfile profile_; // Track which tensor views are inputs or outputs of a vectorized operation diff --git a/torch/csrc/jit/codegen/cuda/lower_index.cpp b/torch/csrc/jit/codegen/cuda/lower_index.cpp index 10ed09f0761cf..210d57d41a04c 100644 --- a/torch/csrc/jit/codegen/cuda/lower_index.cpp +++ b/torch/csrc/jit/codegen/cuda/lower_index.cpp @@ -778,7 +778,7 @@ void IndexLowering::handle(const LoadStoreOp* ldst) { pushBack(IrBuilder::create(ldst->opType(), out, in)); GpuLower::current()->propagateExprInfo(ldst, back()); if (ldst->predicate()) { - back()->setPredicate(ldst->predicate()); + back()->setPredicate(IrBuilder::create(ldst->predicate())); } } diff --git a/torch/csrc/jit/codegen/cuda/lower_predicate_peeling.cpp b/torch/csrc/jit/codegen/cuda/lower_predicate_peeling.cpp index fa756e8b6aa66..8bdadb08c8efb 100644 --- a/torch/csrc/jit/codegen/cuda/lower_predicate_peeling.cpp +++ b/torch/csrc/jit/codegen/cuda/lower_predicate_peeling.cpp @@ -14,6 +14,7 @@ namespace fuser { namespace cuda { bool PredicatePeeling::supportedPeelingLoop(IterDomain* id) { + // Not meaningful to peel a parallel loop if (id->isParallelized()) { return false; } @@ -38,6 +39,235 @@ bool PredicatePeeling::supportedPeelingLoop(IterDomain* id) { return false; } +void PredicatePeelingInfo::build(Fusion* fusion) { + auto used_vals = fusion->usedMathVals(); + for (auto tv : ir_utils::filterByType(used_vals)) { + // Only visit tensorviews with peeling serial id info. + if (!tv->peeledSerialId().empty()) { + // Create a set of leaf ids for validation. + std::unordered_set leaf_id_set{ + tv->domain()->domain().begin(), tv->domain()->domain().end()}; + + for (auto peeled_id : tv->peeledSerialId()) { + // Quick check that the peeled id at schedule + // time is still leaf. + TORCH_INTERNAL_ASSERT( + leaf_id_set.count(peeled_id), + "only exisiting leaf domain supported for peeling\n", + tv->toString(), + " does not have\n", + peeled_id->toString()); + + // Insert the peeled concrete id to the recorded map. + concrete_id_of_peeled_loops_.insert( + GpuLower::current()->caMap()->getConcreteMappedID( + peeled_id, IdMappingMode::LOOP)); + } + } + } +} + +c10::optional PredicatePeelingInfo::getMaybePeeledTileEntry( + const std::vector& loops, + IterDomain* root_id) { + auto gpu_lower = GpuLower::current(); + + std::unordered_map concrete_id_to_loop_map; + for (auto fl : loops) { + auto concrete_loop_id = gpu_lower->caMap()->getConcreteMappedID( + fl->iter_domain(), IdMappingMode::LOOP); + concrete_id_to_loop_map[concrete_loop_id] = fl; + } + + for (auto peeled_id : concrete_id_of_peeled_loops_) { + // Need to see the peeled loop to validate this check + auto matching_loop_it = concrete_id_to_loop_map.find(peeled_id); + if (matching_loop_it != concrete_id_to_loop_map.end()) { + // This is the only supported case at the initial stage. + auto split = peeled_id->definition()->as(); + if (gpu_lower->caMap()->areMapped( + split->in(), root_id, IdMappingMode::EXACT)) { + // This means the given id has been peeled. + PeeledTileEntry entry; + entry.peel_stage = + matching_loop_it->second->loopTransformInfo().predicate_peel_stage; + entry.inner_factor = split->factor(); + entry.for_loop = matching_loop_it->second; + return entry; + } + } + } + return c10::nullopt; +} + +bool PredicatePeelingInfo::shouldPeelLoop(kir::ForLoop* forloop) const { + auto loop_concrete_id = GpuLower::current()->caMap()->getConcreteMappedID( + forloop->iter_domain(), IdMappingMode::LOOP); + + return concrete_id_of_peeled_loops_.count(loop_concrete_id); +} + +namespace { +class LoopNestDeepCloner { + public: + static void clone(kir::ForLoop* original_fl, kir::ForLoop* new_fl) { + LoopNestDeepCloner cloner; + cloner.cloned_scopes_.push_back(&new_fl->body()); + for (auto expr : original_fl->body().exprs()) { + cloner.handle(expr); + } + } + + private: + void handle(Expr* expr) { + if (auto fl = dynamic_cast(expr)) { + handle(fl); + } else { + cloned_scopes_.back()->push_back(expr); + } + } + + void handle(kir::ForLoop* fl) { + auto new_fl = IrBuilder::create(fl); + cloned_scopes_.push_back(&new_fl->body()); + for (auto expr : fl->body().exprs()) { + handle(expr); + } + cloned_scopes_.pop_back(); + + cloned_scopes_.back()->push_back(new_fl); + } + + private: + std::vector cloned_scopes_; +}; + +class PredicatePeeledLoops : kir::ExprMutator { + public: + static std::vector run(const std::vector exprs) { + PredicatePeeledLoops peeled_loops; + peeled_loops.traverseAndInsert(exprs); + return peeled_loops.exprs_; + } + + private: + using kir::ExprMutator::handle; + + kir::ForLoop* createPeeledLoop(kir::ForLoop* fl) { + // Make clone of the outermost loop, but + // only the first iteration. + auto peeled_loop = IrBuilder::create( + fl->iter_domain(), + fl->index(), + fl->kernel()->zeroVal(), + fl->kernel()->oneVal(), + fl->kernel()->oneVal(), + false, + nullptr, + fl->isUnrollRequired(), + fl->loopTransformInfo().predicatePeelStage(PredicatePeelStage::Prolog)); + + LoopNestDeepCloner::clone(fl, peeled_loop); + return peeled_loop; + } + + kir::ForLoop* createMainLoop(kir::ForLoop* fl) { + auto start = + SimplifyingIrBuilder::addExpr(fl->start(), fl->kernel()->oneVal()); + + auto main_loop = IrBuilder::create( + fl->iter_domain(), + fl->index(), + start, + fl->stop(), + fl->step(), + fl->vectorize(), + fl->vectorize_shift(), + fl->isUnrollRequired(), + fl->loopTransformInfo().predicatePeelStage(PredicatePeelStage::Main)); + + LoopNestDeepCloner::clone(fl, main_loop); + return main_loop; + } + + void handle(kir::ForLoop* fl) final { + kir::ExprMutator::handle(fl); + if (GpuLower::current()->predicatePeelingInfo().shouldPeelLoop(fl)) { + auto loop_concrete_id = GpuLower::current()->caMap()->getConcreteMappedID( + fl->iter_domain(), IdMappingMode::LOOP); + if (!peeled_iterdomains_.count(loop_concrete_id)) { + auto peeled_loop = createPeeledLoop(fl); + registerInsertBefore(fl, peeled_loop); + + if (fl->stop()->isOneInt()) { + // This is the case for double buffer prolog, + // will just use the peeled loop as the new + // double buffer prolog. + registerRemove(fl); + } else { + // Peel off one iteration from the main + // component of the original loop. + auto new_main_loop = createMainLoop(fl); + registerReplace(fl, new_main_loop); + } + // Record peeling of this loop + peeled_iterdomains_.insert(loop_concrete_id); + } + } + } + + void handle(kir::IfThenElse* ite) final { + TORCH_INTERNAL_ASSERT( + false, "no support for inserted ite before this point"); + } + + private: + // This pass runs after double buffering pass, so + // we may encounter forloops corresponding to the + // same loop domain twice. In this case we only need + // to peel the first occurrence (the prolog). + std::unordered_set peeled_iterdomains_; +}; + +} // namespace + +std::vector PredicatePeeling::peelPredicatedLoop( + const std::vector exprs) { + return PredicatePeeledLoops::run(exprs); +} + +Val* PredicatePeeling::getSplitTileOffset(IterDomain* id, Val* tile_factor) { + // Assume X = original extent, + // L = tile factor + // Offset needs to satisfy: + // X % L == 0 : return L else return X % L + + // X + L - ceildiv(X,L)*L + auto orig_extent = id->extent(); + + // X + L + auto extent_plus_factor = + SimplifyingIrBuilder::addExpr(orig_extent, tile_factor); + + // ceildiv(X,L) + auto extent_ceildiv_factor = + SimplifyingIrBuilder::ceilDivExpr(orig_extent, tile_factor); + + // ceildiv(X,L)* L + auto extent_round_up = + SimplifyingIrBuilder::mulExpr(extent_ceildiv_factor, tile_factor); + + // X + L - ceildiv(X,L)*L + return SimplifyingIrBuilder::subExpr(extent_plus_factor, extent_round_up); +} + +Val* PredicatePeeling::getSplitTileMainOffset( + IterDomain* id, + Val* tile_factor) { + return SimplifyingIrBuilder::subExpr( + tile_factor, PredicatePeeling::getSplitTileOffset(id, tile_factor)); +} + } // namespace cuda } // namespace fuser } // namespace jit diff --git a/torch/csrc/jit/codegen/cuda/lower_predicate_peeling.h b/torch/csrc/jit/codegen/cuda/lower_predicate_peeling.h index 2904627b34f15..f43c9e1e6d991 100644 --- a/torch/csrc/jit/codegen/cuda/lower_predicate_peeling.h +++ b/torch/csrc/jit/codegen/cuda/lower_predicate_peeling.h @@ -13,15 +13,25 @@ namespace jit { namespace fuser { namespace cuda { +// Stage, Factor pair: +struct PeeledTileEntry { + PredicatePeelStage peel_stage = PredicatePeelStage::NoApplicable; + Val* inner_factor = nullptr; + kir::ForLoop* for_loop = nullptr; +}; + class PredicatePeelingInfo { public: - bool shouldPeelLoop(kir::ForLoop* forloop); + bool shouldPeelLoop(kir::ForLoop* forloop) const; - void markLoopPeeled(kir::ForLoop* forloop); + void build(Fusion* fusion); + + c10::optional getMaybePeeledTileEntry( + const std::vector& loops, + IterDomain* root_id); private: - private: - std::unordered_set concrete_loops_to_peel_; + std::unordered_set concrete_id_of_peeled_loops_; }; namespace PredicatePeeling { @@ -32,6 +42,10 @@ bool supportedPeelingLoop(IterDomain* id); std::vector peelPredicatedLoop(const std::vector exprs); +Val* getSplitTileOffset(IterDomain* id, Val* tile_factor); + +Val* getSplitTileMainOffset(IterDomain* id, Val* tile_factor); + } // namespace PredicatePeeling } // namespace cuda diff --git a/torch/csrc/jit/codegen/cuda/runtime/random_numbers.cu b/torch/csrc/jit/codegen/cuda/runtime/random_numbers.cu index e26b99c068440..4736d8ac6b176 100644 --- a/torch/csrc/jit/codegen/cuda/runtime/random_numbers.cu +++ b/torch/csrc/jit/codegen/cuda/runtime/random_numbers.cu @@ -56,10 +56,12 @@ __device__ double uniform(unsigned int x, unsigned int y) { return z * kRan2Pow53Inv + (kRan2Pow53Inv / 2.0); } -__device__ double randLike(const uint4 &rng_result, int rng_component) { - return uniform((&rng_result.x)[rng_component * 2], (&rng_result.x)[rng_component * 2 + 1]); +__device__ double randLike(const uint4& rng_result, int rng_component) { + return uniform( + (&rng_result.x)[rng_component * 2], + (&rng_result.x)[rng_component * 2 + 1]); } -__device__ float randLikef(const uint4 &rng_result, int rng_component) { +__device__ float randLikef(const uint4& rng_result, int rng_component) { return uniformf((&rng_result.x)[rng_component]); } diff --git a/torch/csrc/jit/codegen/cuda/tensor_view.cpp b/torch/csrc/jit/codegen/cuda/tensor_view.cpp index 0e81f74bcfbbb..88195b17c8497 100644 --- a/torch/csrc/jit/codegen/cuda/tensor_view.cpp +++ b/torch/csrc/jit/codegen/cuda/tensor_view.cpp @@ -11,6 +11,7 @@ #include #include #include +#include #include // Cleanup @@ -221,6 +222,9 @@ TensorView::TensorView(const TensorView* src, IrCloner* ir_cloner) for (const auto id : src->axesToSwizzle()) { axes_to_swizzle_.push_back(ir_cloner->clone(id)); } + for (const auto id : src->peeled_serial_id_) { + peeled_serial_id_.push_back(ir_cloner->clone(id)); + } } bool TensorView::hasAnyReduction() const { @@ -1182,6 +1186,13 @@ void TensorView::applyMmaSwizzle(MmaOptions options) { } } +void TensorView::peelPredicatedLoop(int axis_id) { + auto id = axis(axis_id); + TORCH_CHECK( + PredicatePeeling::supportedPeelingLoop(id), "unsupported loop peeling"); + peeled_serial_id_.push_back(id); +} + TensorViewBuilder& TensorViewBuilder::ndims(size_t ndims) { TORCH_CHECK(shape_.empty() || shape_.size() == ndims); TORCH_CHECK(contiguity_.empty() || contiguity_.size() == ndims); diff --git a/torch/csrc/jit/codegen/cuda/test/test_gpu.cpp b/torch/csrc/jit/codegen/cuda/test/test_gpu.cpp index c89edb22b8858..61605f8d026f8 100644 --- a/torch/csrc/jit/codegen/cuda/test/test_gpu.cpp +++ b/torch/csrc/jit/codegen/cuda/test/test_gpu.cpp @@ -22999,6 +22999,45 @@ TEST_F(NVFuserTest, FusionCpAsyncPredicate_CUDA) { testValidate(&fusion, cg_outputs, {t0}, {ref}, __LINE__, __FILE__); } +// Simple test case for predicate peeling use pattern +TEST_F(NVFuserTest, FusionPredicatePeeling_CUDA) { + // requires ampere+ GPU + Fusion fusion; + FusionGuard fg(&fusion); + + // Using vectorization so need to keep n multiple of 4. + int m = 33, n = 48; + + TensorView* tv0 = makeContigTensor(2); + + fusion.addInput(tv0); + auto tv1 = set(tv0); + auto tv2 = set(tv1); + fusion.addOutput(tv2); + + tv2->split(0, 16); + + tv0->computeAt(tv2, 1); + tv1->computeAt(tv2, -1); + + tv2->axis(1)->parallelize(ParallelType::TIDx); + tv2->axis(2)->parallelize(ParallelType::TIDy); + tv2->peelPredicatedLoop(0); + + fusion.printMath(); + fusion.printKernel(); + + auto options = at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0); + at::Tensor t0 = at::randn({m, n}, options); + + FusionExecutor fe; + + fe.compileFusion(&fusion, {t0}); + auto cg_outputs = fe.runFusion({t0}); + + testValidate(&fusion, cg_outputs, {t0}, {t0}, __LINE__, __FILE__); +} + // Test predicate removal on reg-to-reg expressions TEST_F(NVFuserTest, FusionPredRemovalCheck_CUDA) { Fusion fusion; diff --git a/torch/csrc/jit/codegen/cuda/type.cpp b/torch/csrc/jit/codegen/cuda/type.cpp index b14cdfeb7df3c..243842dcb8e36 100644 --- a/torch/csrc/jit/codegen/cuda/type.cpp +++ b/torch/csrc/jit/codegen/cuda/type.cpp @@ -1199,9 +1199,12 @@ TORCH_CUDA_CU_API std::ostream& operator<<( default: TORCH_INTERNAL_ASSERT(false, "unknown double buffer stage"); } + return os; } -std::ostream& operator<<(std::ostream& os, PredicatePeelStage peel_stage) { +std::ostream& operator<<( + std::ostream& os, + const PredicatePeelStage& peel_stage) { switch (peel_stage) { case PredicatePeelStage::NoApplicable: break; diff --git a/torch/csrc/jit/codegen/cuda/type.h b/torch/csrc/jit/codegen/cuda/type.h index 1fa04c519f461..d0e4d625029f7 100644 --- a/torch/csrc/jit/codegen/cuda/type.h +++ b/torch/csrc/jit/codegen/cuda/type.h @@ -334,11 +334,7 @@ enum class DoubleBufferLoopStage { LowerProlog }; -enum class PredicatePeelStage { - NoApplicable, - Prolog, - Main -} predicate_peel_stage; +enum class PredicatePeelStage { NoApplicable, Prolog, Main }; //! Supported swizzle types, //! corresponds to swizzles functions on the runtime cuda @@ -390,6 +386,9 @@ TORCH_CUDA_CU_API std::ostream& operator<<( const DoubleBufferLoopStage); TORCH_CUDA_CU_API std::ostream& operator<<(std::ostream&, const Swizzle2DType&); TORCH_CUDA_CU_API std::ostream& operator<<(std::ostream&, const SwizzleMode&); +TORCH_CUDA_CU_API std::ostream& operator<<( + std::ostream&, + const PredicatePeelStage&); std::string stringifyBooleanOp(const UnaryOpType); std::string stringifyBooleanOp(const BinaryOpType); From 109dbce6bd51232e0840c242ff3984c2c1f4fbdb Mon Sep 17 00:00:00 2001 From: shmsong Date: Sun, 21 Aug 2022 11:03:49 -0700 Subject: [PATCH 03/14] minor fix on circular buffer prolog --- torch/csrc/jit/codegen/cuda/index_compute.cpp | 13 +++++++++++-- 1 file changed, 11 insertions(+), 2 deletions(-) diff --git a/torch/csrc/jit/codegen/cuda/index_compute.cpp b/torch/csrc/jit/codegen/cuda/index_compute.cpp index 9d2102d56ca16..9f88bed7d01ed 100644 --- a/torch/csrc/jit/codegen/cuda/index_compute.cpp +++ b/torch/csrc/jit/codegen/cuda/index_compute.cpp @@ -2279,7 +2279,8 @@ std::vector Index::getNonGlobalConsumerStridedIndices( 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->index(); + db_switch_index = + db_loop->isTrivial() ? db_loop->start() : db_loop->index(); } else { // Switching index generated for main loop or epilog component. db_switch_index = SimplifyingIrBuilder::modExpr( @@ -3220,7 +3221,15 @@ std::vector Index::getReferenceRootPredicates( ->as(); } else if ( tile_entry.for_loop->doubleBufferLoopStage() == - DoubleBufferLoopStage::Main) { + DoubleBufferLoopStage::Main && + // TODO: need to unify the double buffer index/predicate calculation + // TODO: need to update this part as well in double buffer extension + // for turing/volta staging regs. + db_axis != nullptr && + GpuLower::current()->caMap()->areMapped( + db_axis, + tile_entry.for_loop->iter_domain(), + IdMappingMode::LOOP)) { auto db_index = SimplifyingIrBuilder::addExpr( tile_entry.for_loop->index(), IrBuilder::create( From 191c2b607b8e4d783fdb27a1de78f5267d69b895 Mon Sep 17 00:00:00 2001 From: shmsong Date: Sun, 21 Aug 2022 21:59:52 -0700 Subject: [PATCH 04/14] WAR on circular buffer peeled predicate. --- torch/csrc/jit/codegen/cuda/ir_iostream.cpp | 11 +++++++++-- torch/csrc/jit/codegen/cuda/lower_allocation.cpp | 7 +++++++ torch/csrc/jit/codegen/cuda/lower_predicate.cpp | 8 ++++++-- .../jit/codegen/cuda/lower_predicate_peeling.cpp | 15 +++++++++++++++ .../jit/codegen/cuda/lower_predicate_peeling.h | 2 ++ 5 files changed, 39 insertions(+), 4 deletions(-) diff --git a/torch/csrc/jit/codegen/cuda/ir_iostream.cpp b/torch/csrc/jit/codegen/cuda/ir_iostream.cpp index 6c6b8e02b2f29..c81b9c7ab8130 100644 --- a/torch/csrc/jit/codegen/cuda/ir_iostream.cpp +++ b/torch/csrc/jit/codegen/cuda/ir_iostream.cpp @@ -435,8 +435,12 @@ void IrPrinter::handle(const WelfordOp* wop) { } void IrPrinter::handle(const LoadStoreOp* ldst) { - indent() << ldst->out() << " = " << ldst->opType() << "( " << ldst->in() - << " )\n"; + indent() << ldst->out() << " = " << ldst->opType() << "( " << ldst->in(); + if (ldst->container()->isA() && ldst->predicate() != nullptr && + ldst->predicate()->hasValue()) { + os_ << ", " << ldst->predicate()->value()->toInlineString(); + } + os_ << " )\n"; } void IrPrinter::handle(const BroadcastOp* bop) { @@ -549,6 +553,9 @@ void IrPrinter::handle(const kir::Predicate* node) { } default: os_ << node->predicate_type(); + if (node->hasValue()) { + os_ << " : " << node->value()->toInlineString(); + } break; } } diff --git a/torch/csrc/jit/codegen/cuda/lower_allocation.cpp b/torch/csrc/jit/codegen/cuda/lower_allocation.cpp index 065b3bb1a9c34..67f0b7cf30d7b 100644 --- a/torch/csrc/jit/codegen/cuda/lower_allocation.cpp +++ b/torch/csrc/jit/codegen/cuda/lower_allocation.cpp @@ -449,6 +449,13 @@ class AllocationInserter : public kir::ExprMutator { auto out_tv = out->as(); auto default_val = gpu_lower->predicateElimination().getInitValue(out_tv); + if (out_tv->isCircularBuffered() && default_val == nullptr) { + if (GpuLower::current()->predicatePeelingInfo().hasPeeledId(out_tv)) { + // Always initialize cp async output if it has peeled id. + default_val = GpuLower::current()->kernel()->zeroVal(); + } + } + Val* init = nullptr; if (expr->isA() && out_tv->hasReduction()) { TORCH_INTERNAL_ASSERT( diff --git a/torch/csrc/jit/codegen/cuda/lower_predicate.cpp b/torch/csrc/jit/codegen/cuda/lower_predicate.cpp index e6b7bc332b8c9..007469c8a3712 100644 --- a/torch/csrc/jit/codegen/cuda/lower_predicate.cpp +++ b/torch/csrc/jit/codegen/cuda/lower_predicate.cpp @@ -112,8 +112,12 @@ class ConditionalFromPredicateModifier : public kir::IrVisitor { // Invert the predicate of given expr. void invertPredicateForGmemToSharedMemInitialize(Expr* expr) { auto pred = expr->predicate()->value(); - auto invert = SimplifyingIrBuilder::notExpr(pred); - expr->predicate()->setValue(invert->as()); + if (!pred->isConst()) { + // TODO: this is a WAR for the current incomplete support + // of Inline type predicate for cp.async. + auto invert = SimplifyingIrBuilder::notExpr(pred); + expr->predicate()->setValue(invert->as()); + } } // Detect if this expr is an initialization for vectorized diff --git a/torch/csrc/jit/codegen/cuda/lower_predicate_peeling.cpp b/torch/csrc/jit/codegen/cuda/lower_predicate_peeling.cpp index 8bdadb08c8efb..80f6e530b9bbe 100644 --- a/torch/csrc/jit/codegen/cuda/lower_predicate_peeling.cpp +++ b/torch/csrc/jit/codegen/cuda/lower_predicate_peeling.cpp @@ -100,6 +100,21 @@ c10::optional PredicatePeelingInfo::getMaybePeeledTileEntry( return c10::nullopt; } +bool PredicatePeelingInfo::hasPeeledId(const TensorView* tv) const { + for (auto id : concrete_id_of_peeled_loops_) { + if (std::any_of( + tv->domain()->domain().begin(), + tv->domain()->domain().end(), + [id](IterDomain* tv_id) { + return GpuLower::current()->caMap()->areMapped( + tv_id, id, IdMappingMode::LOOP); + })) { + return true; + } + } + return false; +} + bool PredicatePeelingInfo::shouldPeelLoop(kir::ForLoop* forloop) const { auto loop_concrete_id = GpuLower::current()->caMap()->getConcreteMappedID( forloop->iter_domain(), IdMappingMode::LOOP); diff --git a/torch/csrc/jit/codegen/cuda/lower_predicate_peeling.h b/torch/csrc/jit/codegen/cuda/lower_predicate_peeling.h index f43c9e1e6d991..f1ce87d2dc9b5 100644 --- a/torch/csrc/jit/codegen/cuda/lower_predicate_peeling.h +++ b/torch/csrc/jit/codegen/cuda/lower_predicate_peeling.h @@ -30,6 +30,8 @@ class PredicatePeelingInfo { const std::vector& loops, IterDomain* root_id); + bool hasPeeledId(const TensorView* tv) const; + private: std::unordered_set concrete_id_of_peeled_loops_; }; From d73ee913c11475fe3863a34fa34b3b289c66fd41 Mon Sep 17 00:00:00 2001 From: shmsong Date: Sun, 21 Aug 2022 22:07:07 -0700 Subject: [PATCH 05/14] add circular buffered test case --- torch/csrc/jit/codegen/cuda/test/test_gpu.cpp | 52 ++++++++++++++++++- 1 file changed, 51 insertions(+), 1 deletion(-) diff --git a/torch/csrc/jit/codegen/cuda/test/test_gpu.cpp b/torch/csrc/jit/codegen/cuda/test/test_gpu.cpp index 61605f8d026f8..0676cf2c99deb 100644 --- a/torch/csrc/jit/codegen/cuda/test/test_gpu.cpp +++ b/torch/csrc/jit/codegen/cuda/test/test_gpu.cpp @@ -23000,7 +23000,7 @@ TEST_F(NVFuserTest, FusionCpAsyncPredicate_CUDA) { } // Simple test case for predicate peeling use pattern -TEST_F(NVFuserTest, FusionPredicatePeeling_CUDA) { +TEST_F(NVFuserTest, FusionPredicatePeeling1_CUDA) { // requires ampere+ GPU Fusion fusion; FusionGuard fg(&fusion); @@ -23038,6 +23038,56 @@ TEST_F(NVFuserTest, FusionPredicatePeeling_CUDA) { testValidate(&fusion, cg_outputs, {t0}, {t0}, __LINE__, __FILE__); } +// A circular buffer test case for predicate peeling use pattern +TEST_F(NVFuserTest, FusionPredicatePeeling2_CUDA) { + // requires ampere+ GPU + if (!deviceMajorMinorCheck(8)) { + GTEST_SKIP() << "skipping tests on pre-AMPERE GPUs"; + return; + } + // requires ampere+ GPU + Fusion fusion; + FusionGuard fg(&fusion); + + // Using vectorization so need to keep n multiple of 4. + int m = 33, n = 45; + + TensorView* tv0 = makeContigTensor(2); + + fusion.addInput(tv0); + auto tv1 = sum(tv0, {1}); + fusion.addOutput(tv1); + + auto tv2 = tv0->cacheAfter(LoadStoreOpType::CpAsync); + + tv1->split(1, 16); + tv1->split(0, 16); + // make tile + tv1->reorder({{1, 2}, {2, 1}}); + + tv0->computeAt(tv1, 2); + + tv2->axis(-1)->parallelize(ParallelType::TIDx); + + tv1->axis(0)->parallelize(ParallelType::BIDx); + tv1->axis(2)->parallelize(ParallelType::TIDx); + tv1->peelPredicatedLoop(1); + + tv2->setMemoryType(MemoryType::Shared); + tv2->circularBuffer(3); + + auto options = at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0); + at::Tensor t0 = at::randn({m, n}, options); + + FusionExecutor fe; + + fe.compileFusion(&fusion, {t0}); + auto cg_outputs = fe.runFusion({t0}); + auto ref = t0.sum({1}); + + testValidate(&fusion, cg_outputs, {t0}, {ref}, __LINE__, __FILE__); +} + // Test predicate removal on reg-to-reg expressions TEST_F(NVFuserTest, FusionPredRemovalCheck_CUDA) { Fusion fusion; From 2c7d1dd3eb9babbe363b761ff5da33d63be35ea7 Mon Sep 17 00:00:00 2001 From: shmsong Date: Mon, 22 Aug 2022 10:53:49 -0700 Subject: [PATCH 06/14] add test back (rebase fix) --- .../codegen/cuda/test/test_gpu_tensorcore.cpp | 244 ++++++++++++++++++ 1 file changed, 244 insertions(+) diff --git a/torch/csrc/jit/codegen/cuda/test/test_gpu_tensorcore.cpp b/torch/csrc/jit/codegen/cuda/test/test_gpu_tensorcore.cpp index 2b00f3c2a6adb..54659812fb0de 100644 --- a/torch/csrc/jit/codegen/cuda/test/test_gpu_tensorcore.cpp +++ b/torch/csrc/jit/codegen/cuda/test/test_gpu_tensorcore.cpp @@ -2767,6 +2767,250 @@ TEST_F(NVFuserTest, FusionAmpereMatmulTNSwizzled_CUDA) { TORCH_CHECK(cg_outputs[0].allclose(tref, 0.0001, 0.0001)); } +// Matmul test on Ampere using ldmatrix.x4 to load operands +TEST_F(NVFuserTest, FusionAmpereMatmulLargeLoad_CUDA) { + // Keep multiples of 8 to keep vectorizable. + int M = 504, N = 136, K = 248; + for (auto layout : kAllSupportedLayout) { + Fusion fusion; + FusionGuard fg(&fusion); + auto tv0 = makeContigTensor(2, DataType::Half); + auto tv1 = makeContigTensor(2, DataType::Half); + + fusion.addInput(tv0); + fusion.addInput(tv1); + + auto tv2 = matmul(tv0, tv1, layout); + + fusion.addOutput(tv2); + + MatMulTileOptions gemm_tile; + gemm_tile.cta_tile = GemmTile(128, 128, 32); + gemm_tile.warp_tile = GemmTile(64, 64, 32); + gemm_tile.instruction_tile = GemmTile(16, 16, 16); + + auto mma_builder = + MmaBuilder(MmaOptions::MacroType::Ampere_16_16_16, gemm_tile) + .layout(layout); + + MatmulParam params(mma_builder); + params.tile_sizes = gemm_tile; + params.async_gmem_load_operands = true; + params.double_buffer_options.double_buffer_smem_write = true; + params.double_buffer_options.smem_double_buffer_stage = 4; + scheduleMatmul(tv2, tv0, tv1, params); + + at::manual_seed(0); + auto inputs = fp16MatmulAtInput(M, N, K, layout); + + FusionExecutor fe; + NVFUSER_TEST_CUDA_ARCH_COMPILE_CHECK( + 8, 0, fe.compileFusion(&fusion, {inputs.first, inputs.second})); + auto cg_outputs = fe.runFusion({inputs.first, inputs.second}); + auto tref = atMatmul( + inputs.first.to(at::kFloat), inputs.second.to(at::kFloat), layout); + TORCH_CHECK(cg_outputs[0].allclose(tref, 0.0001, 0.0001)); + } +} + +// Tile layout check for symmetric 4-warp recipes +TEST_F(NVFuserTest, FusionAmpereMatmulTileCheck4warp_CUDA) { + REQUIRE_DEVICE_SMEM_SIZE(98384, 0); + // Keep multiples of 8 to keep vectorizable. + int M = 504, N = 136, K = 248; + for (auto layout : kAllSupportedLayout) { + // Symmetric tile with 16x16x16 macro, + // supports mn_size of multiple of 32, + // and k size multiple of 16. + for (int mn_size : {32, 64, 96, 128, 160, 192}) { + for (int k_size : {32, 48, 64}) { + Fusion fusion; + FusionGuard fg(&fusion); + auto tv0 = makeContigTensor(2, DataType::Half); + auto tv1 = makeContigTensor(2, DataType::Half); + + fusion.addInput(tv0); + fusion.addInput(tv1); + + auto tv2 = matmul(tv0, tv1, layout); + + fusion.addOutput(tv2); + + MatMulTileOptions gemm_tile; + gemm_tile.cta_tile = GemmTile(mn_size, mn_size, k_size); + gemm_tile.warp_tile = GemmTile(mn_size / 2, mn_size / 2, k_size); + gemm_tile.instruction_tile = GemmTile(16, 16, 16); + + auto mma_builder = + MmaBuilder(MmaOptions::MacroType::Ampere_16_16_16, gemm_tile) + .layout(layout); + + MatmulParam params(mma_builder); + params.tile_sizes = gemm_tile; + params.async_gmem_load_operands = true; + params.double_buffer_options.double_buffer_smem_write = true; + scheduleMatmul(tv2, tv0, tv1, params); + + at::manual_seed(0); + auto inputs = fp16MatmulAtInput(M, N, K, layout); + + FusionExecutor fe; + NVFUSER_TEST_CUDA_ARCH_COMPILE_CHECK( + 8, 0, fe.compileFusion(&fusion, {inputs.first, inputs.second})); + auto cg_outputs = fe.runFusion({inputs.first, inputs.second}); + auto tref = atMatmul( + inputs.first.to(at::kFloat), inputs.second.to(at::kFloat), layout); + TORCH_CHECK(cg_outputs[0].allclose(tref, 0.0001, 0.0001)); + } + } + } +} + +TEST_F(NVFuserTest, FusionAmpereMatmulTileCheck8warp_CUDA) { + REQUIRE_DEVICE_SMEM_SIZE(98384, 0); + // Keep multiples of 8 to keep vectorizable. + int M = 504, N = 136, K = 248; + for (auto layout : kAllSupportedLayout) { + // ASymmetric tile with 16x16x16 macro, + for (int m_size : {256}) { + for (int n_size : {32, 64, 96, 128}) { + for (int k_size : {32, 48, 64}) { + Fusion fusion; + FusionGuard fg(&fusion); + auto tv0 = makeContigTensor(2, DataType::Half); + auto tv1 = makeContigTensor(2, DataType::Half); + + fusion.addInput(tv0); + fusion.addInput(tv1); + + auto tv2 = matmul(tv0, tv1, layout); + + fusion.addOutput(tv2); + + MatMulTileOptions gemm_tile; + gemm_tile.cta_tile = GemmTile(m_size, n_size, k_size); + gemm_tile.warp_tile = GemmTile(m_size / 4, n_size / 2, k_size); + gemm_tile.instruction_tile = GemmTile(16, 16, 16); + + auto mma_builder = + MmaBuilder(MmaOptions::MacroType::Ampere_16_16_16, gemm_tile) + .layout(layout); + + MatmulParam params(mma_builder); + params.tile_sizes = gemm_tile; + params.async_gmem_load_operands = true; + params.double_buffer_options.double_buffer_smem_write = true; + scheduleMatmul(tv2, tv0, tv1, params); + + at::manual_seed(0); + auto inputs = fp16MatmulAtInput(M, N, K, layout); + + FusionExecutor fe; + NVFUSER_TEST_CUDA_ARCH_COMPILE_CHECK( + 8, 0, fe.compileFusion(&fusion, {inputs.first, inputs.second})); + auto cg_outputs = fe.runFusion({inputs.first, inputs.second}); + auto tref = atMatmul( + inputs.first.to(at::kFloat), + inputs.second.to(at::kFloat), + layout); + TORCH_CHECK(cg_outputs[0].allclose(tref, 0.0001, 0.0001)); + } + } + } + } +} + +TEST_F(NVFuserTest, FusionAmpereMatmulTileCheck6warp_CUDA) { + REQUIRE_DEVICE_SMEM_SIZE(98384, 0); + // Keep multiples of 8 to keep vectorizable. + int M = 504, N = 136, K = 248; + for (auto layout : kAllSupportedLayout) { + for (int k_size : {32, 48, 64}) { + Fusion fusion; + FusionGuard fg(&fusion); + auto tv0 = makeContigTensor(2, DataType::Half); + auto tv1 = makeContigTensor(2, DataType::Half); + + fusion.addInput(tv0); + fusion.addInput(tv1); + + auto tv2 = matmul(tv0, tv1, layout); + + fusion.addOutput(tv2); + + MatMulTileOptions gemm_tile; + // 2 warp by 3 warp + gemm_tile.cta_tile = GemmTile(192, 128, k_size); + gemm_tile.warp_tile = GemmTile(64, 64, k_size); + gemm_tile.instruction_tile = GemmTile(16, 16, 16); + + auto mma_builder = + MmaBuilder(MmaOptions::MacroType::Ampere_16_16_16, gemm_tile) + .layout(layout); + + MatmulParam params(mma_builder); + params.tile_sizes = gemm_tile; + params.async_gmem_load_operands = true; + params.double_buffer_options.double_buffer_smem_write = true; + scheduleMatmul(tv2, tv0, tv1, params); + + at::manual_seed(0); + auto inputs = fp16MatmulAtInput(M, N, K, layout); + + FusionExecutor fe; + NVFUSER_TEST_CUDA_ARCH_COMPILE_CHECK( + 8, 0, fe.compileFusion(&fusion, {inputs.first, inputs.second})); + auto cg_outputs = fe.runFusion({inputs.first, inputs.second}); + auto tref = atMatmul( + inputs.first.to(at::kFloat), inputs.second.to(at::kFloat), layout); + TORCH_CHECK(cg_outputs[0].allclose(tref, 0.0001, 0.0001)); + } + } +} + +TEST_F(NVFuserTest, FusionTuringMatmulLargeLoad_CUDA) { + // Keep multiples of 8 to keep vectorizable. + int M = 504, N = 136, K = 248; + + for (auto layout : kAllSupportedLayout) { + Fusion fusion; + FusionGuard fg(&fusion); + auto tv0 = makeContigTensor(2, DataType::Half); + auto tv1 = makeContigTensor(2, DataType::Half); + + fusion.addInput(tv0); + fusion.addInput(tv1); + + auto tv2 = matmul(tv0, tv1, layout); + + fusion.addOutput(tv2); + + MatMulTileOptions gemm_tile; + gemm_tile.cta_tile = GemmTile(128, 128, 32); + gemm_tile.warp_tile = GemmTile(64, 64, 32); + gemm_tile.instruction_tile = GemmTile(16, 16, 16); + + auto mma_builder = + MmaBuilder(MmaOptions::MacroType::Turing_16_16_16, gemm_tile) + .layout(layout); + + MatmulParam params(mma_builder); + params.tile_sizes = gemm_tile; + scheduleMatmul(tv2, tv0, tv1, params); + + at::manual_seed(0); + auto inputs = fp16MatmulAtInput(M, N, K, layout); + + FusionExecutor fe; + NVFUSER_TEST_CUDA_ARCH_COMPILE_CHECK( + 7, 5, fe.compileFusion(&fusion, {inputs.first, inputs.second})); + auto cg_outputs = fe.runFusion({inputs.first, inputs.second}); + auto tref = atMatmul( + inputs.first.to(at::kFloat), inputs.second.to(at::kFloat), layout); + TORCH_CHECK(cg_outputs[0].allclose(tref, 0.0001, 0.0001)); + } +} + TEST_F(NVFuserTest, FusionSimpleSkewDoubleBuffer_CUDA) { Fusion fusion; FusionGuard fg(&fusion); From 1b36e603516369a8922a1db7c1d15c67ab06e215 Mon Sep 17 00:00:00 2001 From: shmsong Date: Mon, 22 Aug 2022 11:35:04 -0700 Subject: [PATCH 07/14] fix index lift handling --- torch/csrc/jit/codegen/cuda/index_compute.cpp | 5 ++++- torch/csrc/jit/codegen/cuda/lower_predicate_peeling.cpp | 3 ++- 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/torch/csrc/jit/codegen/cuda/index_compute.cpp b/torch/csrc/jit/codegen/cuda/index_compute.cpp index 9f88bed7d01ed..6213fce952c63 100644 --- a/torch/csrc/jit/codegen/cuda/index_compute.cpp +++ b/torch/csrc/jit/codegen/cuda/index_compute.cpp @@ -1605,7 +1605,10 @@ std::vector Index::getGlobalProducerStridedIndices( if (auto tile_entry = gpu_lower->predicatePeelingInfo().getMaybePeeledTileEntry( loops, root_dom[i])) { - if (tile_entry.value().peel_stage != PredicatePeelStage::Prolog) { + if (tile_entry.value().peel_stage != PredicatePeelStage::Prolog && + !tile_entry.value() + .for_loop->loopTransformInfo() + .is_base_index_loop) { root_ind = SimplifyingIrBuilder::subExpr( root_ind, PredicatePeeling::getSplitTileMainOffset( diff --git a/torch/csrc/jit/codegen/cuda/lower_predicate_peeling.cpp b/torch/csrc/jit/codegen/cuda/lower_predicate_peeling.cpp index 80f6e530b9bbe..5d2aeb2881887 100644 --- a/torch/csrc/jit/codegen/cuda/lower_predicate_peeling.cpp +++ b/torch/csrc/jit/codegen/cuda/lower_predicate_peeling.cpp @@ -119,7 +119,8 @@ bool PredicatePeelingInfo::shouldPeelLoop(kir::ForLoop* forloop) const { auto loop_concrete_id = GpuLower::current()->caMap()->getConcreteMappedID( forloop->iter_domain(), IdMappingMode::LOOP); - return concrete_id_of_peeled_loops_.count(loop_concrete_id); + return !forloop->loopTransformInfo().is_base_index_loop && + concrete_id_of_peeled_loops_.count(loop_concrete_id); } namespace { From 90bc134ed5c7761fdf9802531218056438770742 Mon Sep 17 00:00:00 2001 From: shmsong Date: Mon, 22 Aug 2022 11:35:15 -0700 Subject: [PATCH 08/14] integrate into matmul scheduler --- torch/csrc/jit/codegen/cuda/scheduler/matmul.cpp | 4 ++++ torch/csrc/jit/codegen/cuda/scheduler/matmul.h | 3 +++ 2 files changed, 7 insertions(+) diff --git a/torch/csrc/jit/codegen/cuda/scheduler/matmul.cpp b/torch/csrc/jit/codegen/cuda/scheduler/matmul.cpp index c4b30e8a225f3..31b30abc24b9d 100644 --- a/torch/csrc/jit/codegen/cuda/scheduler/matmul.cpp +++ b/torch/csrc/jit/codegen/cuda/scheduler/matmul.cpp @@ -544,6 +544,10 @@ void scheduleMatmul( acw_smem->liftReadAddress(); bcw_smem->liftReadAddress(); } + + if (params.peel_main_loop) { + cc->peelPredicatedLoop(2); + } } } // namespace cuda diff --git a/torch/csrc/jit/codegen/cuda/scheduler/matmul.h b/torch/csrc/jit/codegen/cuda/scheduler/matmul.h index 200343f6d6fc9..3d7fdaea8320d 100644 --- a/torch/csrc/jit/codegen/cuda/scheduler/matmul.h +++ b/torch/csrc/jit/codegen/cuda/scheduler/matmul.h @@ -48,6 +48,9 @@ class MatmulParam { // TODO: add gmem_write address for // latency bound kernels. } index_lift_options; + + //! Enables predicate peeling mainloop: + bool peel_main_loop = false; }; //! Prototype auto scheduling function. From d8496bdff18998f30ea61668bfe5701c58184f6a Mon Sep 17 00:00:00 2001 From: shmsong Date: Mon, 22 Aug 2022 13:37:19 -0700 Subject: [PATCH 09/14] lift last iteration of initialization when peeling main loop --- torch/csrc/jit/codegen/cuda/codegen.cpp | 4 +- .../csrc/jit/codegen/cuda/compute_at_map.cpp | 3 + torch/csrc/jit/codegen/cuda/index_compute.cpp | 4 +- .../jit/codegen/cuda/lower_double_buffer.cpp | 81 ++++++++++++++++++- torch/csrc/jit/codegen/cuda/type.h | 1 + 5 files changed, 89 insertions(+), 4 deletions(-) diff --git a/torch/csrc/jit/codegen/cuda/codegen.cpp b/torch/csrc/jit/codegen/cuda/codegen.cpp index 498f8d407bd38..36836748f2666 100644 --- a/torch/csrc/jit/codegen/cuda/codegen.cpp +++ b/torch/csrc/jit/codegen/cuda/codegen.cpp @@ -2296,7 +2296,9 @@ class CudaKernelGenerator : private OptOutConstDispatch { // should not apply here. if (loop->iter_domain()->isParallelized() || loop->loopTransformInfo().predicate_peel_stage == - PredicatePeelStage::Main) { + PredicatePeelStage::Main || + loop->loopTransformInfo().double_buffer_loop_stage == + DoubleBufferLoopStage::CircularInitProlog) { code_ << " = " << gen_start << "; "; } else { // Do not start at the start of the ID when not parallelized. Instead, diff --git a/torch/csrc/jit/codegen/cuda/compute_at_map.cpp b/torch/csrc/jit/codegen/cuda/compute_at_map.cpp index 34011bb92ae14..b8222ef33962e 100644 --- a/torch/csrc/jit/codegen/cuda/compute_at_map.cpp +++ b/torch/csrc/jit/codegen/cuda/compute_at_map.cpp @@ -360,6 +360,9 @@ void ComputeAtMap::allocateIndexVariables() { std::make_unique(DoubleBufferIndices( {{DoubleBufferLoopStage::Prolog, IrBuilder::create(c10::nullopt)}, + // TODO: need to add upper and lower prolog here too. + {DoubleBufferLoopStage::CircularInitProlog, + IrBuilder::create(c10::nullopt)}, {DoubleBufferLoopStage::Main, IrBuilder::create(c10::nullopt)}, {DoubleBufferLoopStage::Epilog, diff --git a/torch/csrc/jit/codegen/cuda/index_compute.cpp b/torch/csrc/jit/codegen/cuda/index_compute.cpp index 6213fce952c63..09fd40883de6b 100644 --- a/torch/csrc/jit/codegen/cuda/index_compute.cpp +++ b/torch/csrc/jit/codegen/cuda/index_compute.cpp @@ -2270,7 +2270,9 @@ std::vector Index::getNonGlobalConsumerStridedIndices( db_loop->doubleBufferLoopStage() == DoubleBufferLoopStage::UpperProlog || db_loop->doubleBufferLoopStage() == - DoubleBufferLoopStage::LowerProlog; + DoubleBufferLoopStage::LowerProlog || + db_loop->doubleBufferLoopStage() == + DoubleBufferLoopStage::CircularInitProlog; Val* db_switch_index = nullptr; diff --git a/torch/csrc/jit/codegen/cuda/lower_double_buffer.cpp b/torch/csrc/jit/codegen/cuda/lower_double_buffer.cpp index 4e394d070f4df..ba8b13ad6b17d 100644 --- a/torch/csrc/jit/codegen/cuda/lower_double_buffer.cpp +++ b/torch/csrc/jit/codegen/cuda/lower_double_buffer.cpp @@ -206,6 +206,10 @@ class DoubleBufferLoopCloner : public kir::IrVisitor { start = IrBuilder::subExpr( double_buffer_loop_->stop(), SimplifyingIrBuilder::create(stage_depth - 1)); + } else if (loop_type_ == DoubleBufferLoopStage::CircularInitProlog) { + TORCH_INTERNAL_ASSERT(start->isZeroInt()); + start = SimplifyingIrBuilder::create(stage_depth - 1); + stop = SimplifyingIrBuilder::create(stage_depth); } cloned_top_level_loop_ = IrBuilder::create( @@ -258,7 +262,9 @@ class DoubleBufferLoopCloner : public kir::IrVisitor { TORCH_INTERNAL_ASSERT(!cloned_scopes_.empty()); if (loop_type_ == DoubleBufferLoopStage::Main) { - cloned_scopes_.back()->push_back(expr); + if (!canOmitInitInMainLoop(expr, double_buffer_loop_)) { + cloned_scopes_.back()->push_back(expr); + } return; } @@ -275,6 +281,7 @@ class DoubleBufferLoopCloner : public kir::IrVisitor { TORCH_INTERNAL_ASSERT(double_buffer_tv != nullptr); return out_tv == double_buffer_tv; }); + if ((loop_type_ == DoubleBufferLoopStage::Prolog && is_double_buffer_load_expr) || (loop_type_ == DoubleBufferLoopStage::Epilog && @@ -286,9 +293,66 @@ class DoubleBufferLoopCloner : public kir::IrVisitor { } else { cloned_scopes_.back()->push_back(expr); } + } else if ( + loop_type_ == DoubleBufferLoopStage::CircularInitProlog && + is_double_buffer_load_expr) { + // Only need the init expressions in circular init prolog stage + if (ir_utils::isTensorScalarFillOp(expr)) { + cloned_scopes_.back()->push_back(expr); + } } } + //! Returns true if the expression is an initialization expr that + //! can be omitted in main loop. + bool canOmitInitInMainLoop(Expr* expr, kir::ForLoop* double_buffer_loop) { + if (!ir_utils::isCpAsyncInit(expr) || + !GpuLower::current()->predicatePeelingInfo().shouldPeelLoop( + double_buffer_loop)) { + return false; + } + + auto out_tv = ir_utils::getTvOutput(expr); + + bool db_loop_found = false; + auto& ca_map = GpuLower::current()->caMap(); + + if (!(out_tv->isDoubleBuffered() || out_tv->isCircularBuffered()) || + !ca_map->areMapped( + GpuLower::current()->doubleBufferInfo().getDoubleBufferAxis(out_tv), + double_buffer_loop->iter_domain(), + IdMappingMode::LOOP)) { + return false; + } + + for (auto id : out_tv->domain()->domain()) { + if (db_loop_found) { + auto loop_concrete_id = + ca_map->getConcreteMappedID(id, IdMappingMode::LOOP); + + if (!loop_concrete_id->isParallelized() && + !loop_concrete_id->extent()->isConstInt()) { + // If all the loops on the right of the peeled double buffer loop are + // either parallel or unrolled. Should not need to + + // TODO: also check broadcast resolution, and predicate removal here, + // for general safety. + return false; + } + } + + db_loop_found = db_loop_found || + ca_map->areMapped( + id, double_buffer_loop->iter_domain(), IdMappingMode::LOOP); + } + + TORCH_INTERNAL_ASSERT( + db_loop_found, + "Double buffer axis not on tv's tensor domain: ", + out_tv->toString()); + return true; + } + private: kir::ForLoop* double_buffer_loop_ = nullptr; const std::vector& double_buffer_load_exprs_; @@ -434,6 +498,17 @@ class DoubleBufferInserter : private kir::ExprMutator { MemoryType::Shared; }); + // If the double buffer loop is to be peeled. Will need to insert + // a circular buffer init stage to initialize the final stage of + // circular buffer space. + if (GpuLower::current()->predicatePeelingInfo().shouldPeelLoop( + double_buffer_loop) && + write_to_smem) { + auto circular_init_loop = DoubleBufferLoopCloner::clone( + double_buffer_loop, loads, DoubleBufferLoopStage::CircularInitProlog); + registerInsertBefore(double_buffer_loop, circular_init_loop); + } + // RAW sync is not inserted for double buffered tensors. The only // exception is the prologue load. bool insert_cpasync_wait = false; @@ -943,7 +1018,9 @@ kir::ForLoop* DoubleBufferInfo::getDoubleBufferLoop( (!ignore_prologue || (loop->doubleBufferLoopStage() != DoubleBufferLoopStage::Prolog && loop->doubleBufferLoopStage() != DoubleBufferLoopStage::UpperProlog && - loop->doubleBufferLoopStage() != DoubleBufferLoopStage::LowerProlog)); + loop->doubleBufferLoopStage() != DoubleBufferLoopStage::LowerProlog && + loop->doubleBufferLoopStage() != + DoubleBufferLoopStage::CircularInitProlog)); }); if (loop_it != loops.end()) { diff --git a/torch/csrc/jit/codegen/cuda/type.h b/torch/csrc/jit/codegen/cuda/type.h index d0e4d625029f7..1b2cf4f10c049 100644 --- a/torch/csrc/jit/codegen/cuda/type.h +++ b/torch/csrc/jit/codegen/cuda/type.h @@ -330,6 +330,7 @@ enum class DoubleBufferLoopStage { Prolog, Main, Epilog, + CircularInitProlog, UpperProlog, LowerProlog }; From 37e54bc71749ddd514e313f4c6d1130e34cc2fb6 Mon Sep 17 00:00:00 2001 From: shmsong Date: Mon, 22 Aug 2022 14:15:05 -0700 Subject: [PATCH 10/14] enable peeling by default --- torch/csrc/jit/codegen/cuda/scheduler/matmul.h | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/torch/csrc/jit/codegen/cuda/scheduler/matmul.h b/torch/csrc/jit/codegen/cuda/scheduler/matmul.h index 3d7fdaea8320d..354e2affeab04 100644 --- a/torch/csrc/jit/codegen/cuda/scheduler/matmul.h +++ b/torch/csrc/jit/codegen/cuda/scheduler/matmul.h @@ -50,7 +50,7 @@ class MatmulParam { } index_lift_options; //! Enables predicate peeling mainloop: - bool peel_main_loop = false; + bool peel_main_loop = true; }; //! Prototype auto scheduling function. From 5d9723a7a0e35620bd5ecba41a92638278bd2e1a Mon Sep 17 00:00:00 2001 From: shmsong Date: Sat, 17 Sep 2022 18:09:54 -0700 Subject: [PATCH 11/14] cleanup ; comment --- torch/csrc/jit/codegen/cuda/index_compute.cpp | 30 +++- .../jit/codegen/cuda/ir_interface_nodes.h | 12 +- torch/csrc/jit/codegen/cuda/kernel_ir.h | 2 + .../jit/codegen/cuda/lower_double_buffer.cpp | 26 ++-- .../csrc/jit/codegen/cuda/lower_predicate.cpp | 8 +- .../codegen/cuda/lower_predicate_peeling.cpp | 73 ++++++---- .../codegen/cuda/lower_predicate_peeling.h | 128 +++++++++++++++++- torch/csrc/jit/codegen/cuda/tensor_view.cpp | 6 +- torch/csrc/jit/codegen/cuda/test/test_gpu.cpp | 3 - 9 files changed, 226 insertions(+), 62 deletions(-) diff --git a/torch/csrc/jit/codegen/cuda/index_compute.cpp b/torch/csrc/jit/codegen/cuda/index_compute.cpp index 0d5cd6e57da8a..1a006905c8939 100644 --- a/torch/csrc/jit/codegen/cuda/index_compute.cpp +++ b/torch/csrc/jit/codegen/cuda/index_compute.cpp @@ -1621,6 +1621,8 @@ std::vector Index::getGlobalProducerStridedIndices( if (auto tile_entry = gpu_lower->predicatePeelingInfo().getMaybePeeledTileEntry( loops, root_dom[i])) { + // Add the "predicate peeling offset", see [Predicate Peeling] + // to the tensor index if this root domain is predicate peeled. if (tile_entry.value().peel_stage != PredicatePeelStage::Prolog && !tile_entry.value() .for_loop->loopTransformInfo() @@ -2129,6 +2131,8 @@ std::vector Index::getGlobalConsumerStridedIndices( if (auto tile_entry = gpu_lower->predicatePeelingInfo().getMaybePeeledTileEntry( loops, root_dom[i])) { + // Add the "predicate peeling offset", see [Predicate Peeling] + // to the tensor index if this root domain is predicate peeled. if (tile_entry.value().peel_stage != PredicatePeelStage::Prolog) { root_ind = SimplifyingIrBuilder::subExpr( root_ind, @@ -3213,6 +3217,8 @@ std::vector Index::getReferenceRootPredicates( gpu_lower->predicatePeelingInfo().getMaybePeeledTileEntry( loops, contig_id); + // Check if this predicate for this contig_id is being generated + // in predicate peeling prolog. See also [Predicate Peeling]. bool has_peeled_prolog = maybe_tiled_entry.has_value() && maybe_tiled_entry.value().peel_stage == PredicatePeelStage::Prolog; @@ -3228,29 +3234,36 @@ std::vector Index::getReferenceRootPredicates( auto stop_pred = SimplifyingIrBuilder::ltExpr( offsetted_stop_index, contig_id->extent()) ->as(); - auto maybe_tiled_entry = - gpu_lower->predicatePeelingInfo().getMaybePeeledTileEntry( - loops, contig_id); + // Modifying predicate math for predicate peeled loop: + // detailed definition see [Predicate Peeling] if (maybe_tiled_entry.has_value()) { auto tile_entry = maybe_tiled_entry.value(); if (tile_entry.peel_stage == PredicatePeelStage::Prolog) { + // In predicate peeled prolog, the stop predicate is + // stop_index < tile_residue stop_pred = SimplifyingIrBuilder::ltExpr( offsetted_stop_index, - PredicatePeeling::getSplitTileOffset( + PredicatePeeling::getPrologPredicateOffset( contig_id, tile_entry.inner_factor)) ->as(); } else if ( + // Handle the condition where the predicate peeled + // iterdomain is double/circular buffered. tile_entry.for_loop->doubleBufferLoopStage() == DoubleBufferLoopStage::Main && - // TODO: need to unify the double buffer index/predicate calculation - // TODO: need to update this part as well in double buffer extension - // for turing/volta staging regs. db_axis != nullptr && GpuLower::current()->caMap()->areMapped( db_axis, tile_entry.for_loop->iter_domain(), IdMappingMode::LOOP)) { + // When the predicate peeled loop is double buffered + // or circular buffered, the producer index is skewed + // ahead of the main loop by (stage_depth-1). So on the + // predicate peeled main loop side, instead of just removing + // the predicate for this contig_id, just re-write it to + // (loop_index + stage_depth) < loop_stop, which should + // be thread uniform and very cheap to evaluate. auto db_index = SimplifyingIrBuilder::addExpr( tile_entry.for_loop->index(), IrBuilder::create( @@ -3261,6 +3274,9 @@ std::vector Index::getReferenceRootPredicates( db_index, tile_entry.for_loop->stop()) ->as(); } else { + // If the predicate peeled loop is not double buffered + // then in the main stage of the predicate peeled loop + // this predicate can just be omitted. stop_pred = gpu_lower->kernel()->trueVal(); } } diff --git a/torch/csrc/jit/codegen/cuda/ir_interface_nodes.h b/torch/csrc/jit/codegen/cuda/ir_interface_nodes.h index 9dd6cf557ef09..c8bac72d1f390 100644 --- a/torch/csrc/jit/codegen/cuda/ir_interface_nodes.h +++ b/torch/csrc/jit/codegen/cuda/ir_interface_nodes.h @@ -549,9 +549,14 @@ class TORCH_CUDA_CU_API TensorView : public Val { friend TORCH_CUDA_CU_API void groupReductions( const std::vector&); + //! A scheduler primitive requesting predicate peeling transform + //! on the loop generated corresponding to `axis_id`. + //! See [Predicate Peeling]. void peelPredicatedLoop(int axis_id); - const auto& peeledSerialId() const { + //! Returns the iterdomain corresponding to the loop that will + //! be using the predicate peeling transform. + auto peeledSerialId() const { return peeled_serial_id_; } @@ -593,7 +598,10 @@ class TORCH_CUDA_CU_API TensorView : public Val { //! Indicates the circular buffering stage depth if applicable. unsigned int circular_buffer_stage_ = 0; - std::vector peeled_serial_id_; + //! Keeps track of the iterdomain that will use predicate + //! peeling transform on the corresponding loop. + IterDomain* peeled_serial_id_ = nullptr; + // special handling for CPU based zero-dim tensors (i.e. CPU Tensors that only // have one value). This is only used if on an input value, otherwise ignored. // This is important as special handling because these "scalars" should be diff --git a/torch/csrc/jit/codegen/cuda/kernel_ir.h b/torch/csrc/jit/codegen/cuda/kernel_ir.h index 8132f6dd9ea08..6ceb70823ad75 100644 --- a/torch/csrc/jit/codegen/cuda/kernel_ir.h +++ b/torch/csrc/jit/codegen/cuda/kernel_ir.h @@ -433,6 +433,8 @@ struct LoopTransformInfo { DoubleBufferLoopStage double_buffer_loop_stage = DoubleBufferLoopStage::NotApplicable; + //! Tracks the predicate peeling stage of this loop, + //! see [Predicate Peeling]. PredicatePeelStage predicate_peel_stage = PredicatePeelStage::NoApplicable; //! Tracks if this for loop is for base index calculation for diff --git a/torch/csrc/jit/codegen/cuda/lower_double_buffer.cpp b/torch/csrc/jit/codegen/cuda/lower_double_buffer.cpp index 7f3ef95da979a..dc0e4b9a505a8 100644 --- a/torch/csrc/jit/codegen/cuda/lower_double_buffer.cpp +++ b/torch/csrc/jit/codegen/cuda/lower_double_buffer.cpp @@ -207,6 +207,7 @@ class DoubleBufferLoopCloner : public kir::IrVisitor { double_buffer_loop_->stop(), SimplifyingIrBuilder::create(stage_depth - 1)); } else if (loop_type_ == DoubleBufferLoopStage::CircularInitProlog) { + // See [Predicate Peeling Interaction with Circular Buffering] TORCH_INTERNAL_ASSERT(start->isZeroInt()); start = SimplifyingIrBuilder::create(stage_depth - 1); stop = SimplifyingIrBuilder::create(stage_depth); @@ -305,7 +306,9 @@ class DoubleBufferLoopCloner : public kir::IrVisitor { //! Returns true if the expression is an initialization expr that //! can be omitted in main loop. + //! See [Predicate Peeling Interaction with Circular Buffering] bool canOmitInitInMainLoop(Expr* expr, kir::ForLoop* double_buffer_loop) { + // Check that this is an initialization for cp.async. if (!ir_utils::isCpAsyncInit(expr) || !GpuLower::current()->predicatePeelingInfo().shouldPeelLoop( double_buffer_loop)) { @@ -314,6 +317,9 @@ class DoubleBufferLoopCloner : public kir::IrVisitor { auto out_tv = ir_utils::getTvOutput(expr); + // Check that the double buffer loop is the main stage of + // the loop defining out_tv as there might be multiple + // loops that realize double buffers. bool db_loop_found = false; auto& ca_map = GpuLower::current()->caMap(); @@ -325,6 +331,13 @@ class DoubleBufferLoopCloner : public kir::IrVisitor { return false; } + // This optimization only applies when all the loops on the + // inner side of the double buffer main loop are either + // constant unrolled or parallel. + // TODO: + // Buffer alias and broadcast resolution might still + // break this. These are not showing in matmul kernels but + // would need to build out support for general safty usage. for (auto id : out_tv->domain()->domain()) { if (db_loop_found) { auto loop_concrete_id = @@ -332,11 +345,6 @@ class DoubleBufferLoopCloner : public kir::IrVisitor { if (!loop_concrete_id->isParallelized() && !loop_concrete_id->extent()->isConstInt()) { - // If all the loops on the right of the peeled double buffer loop are - // either parallel or unrolled. Should not need to - - // TODO: also check broadcast resolution, and predicate removal here, - // for general safety. return false; } } @@ -346,11 +354,9 @@ class DoubleBufferLoopCloner : public kir::IrVisitor { id, double_buffer_loop->iter_domain(), IdMappingMode::LOOP); } - TORCH_INTERNAL_ASSERT( - db_loop_found, - "Double buffer axis not on tv's tensor domain: ", - out_tv->toString()); - return true; + // Only when double buffer loop was found on out_tv could useful + // information have been inferred by this function. + return db_loop_found; } private: diff --git a/torch/csrc/jit/codegen/cuda/lower_predicate.cpp b/torch/csrc/jit/codegen/cuda/lower_predicate.cpp index 007469c8a3712..e6b7bc332b8c9 100644 --- a/torch/csrc/jit/codegen/cuda/lower_predicate.cpp +++ b/torch/csrc/jit/codegen/cuda/lower_predicate.cpp @@ -112,12 +112,8 @@ class ConditionalFromPredicateModifier : public kir::IrVisitor { // Invert the predicate of given expr. void invertPredicateForGmemToSharedMemInitialize(Expr* expr) { auto pred = expr->predicate()->value(); - if (!pred->isConst()) { - // TODO: this is a WAR for the current incomplete support - // of Inline type predicate for cp.async. - auto invert = SimplifyingIrBuilder::notExpr(pred); - expr->predicate()->setValue(invert->as()); - } + auto invert = SimplifyingIrBuilder::notExpr(pred); + expr->predicate()->setValue(invert->as()); } // Detect if this expr is an initialization for vectorized diff --git a/torch/csrc/jit/codegen/cuda/lower_predicate_peeling.cpp b/torch/csrc/jit/codegen/cuda/lower_predicate_peeling.cpp index 5d2aeb2881887..01e5322dd9480 100644 --- a/torch/csrc/jit/codegen/cuda/lower_predicate_peeling.cpp +++ b/torch/csrc/jit/codegen/cuda/lower_predicate_peeling.cpp @@ -35,7 +35,7 @@ bool PredicatePeeling::supportedPeelingLoop(IterDomain* id) { } // TODO: - // should extend to use separability defined in anohter PR. + // could possibly support more patterns using the separability analysis. return false; } @@ -43,26 +43,25 @@ void PredicatePeelingInfo::build(Fusion* fusion) { auto used_vals = fusion->usedMathVals(); for (auto tv : ir_utils::filterByType(used_vals)) { // Only visit tensorviews with peeling serial id info. - if (!tv->peeledSerialId().empty()) { + if (tv->peeledSerialId() != nullptr) { // Create a set of leaf ids for validation. std::unordered_set leaf_id_set{ tv->domain()->domain().begin(), tv->domain()->domain().end()}; - - for (auto peeled_id : tv->peeledSerialId()) { - // Quick check that the peeled id at schedule - // time is still leaf. - TORCH_INTERNAL_ASSERT( - leaf_id_set.count(peeled_id), - "only exisiting leaf domain supported for peeling\n", - tv->toString(), - " does not have\n", - peeled_id->toString()); - - // Insert the peeled concrete id to the recorded map. - concrete_id_of_peeled_loops_.insert( - GpuLower::current()->caMap()->getConcreteMappedID( - peeled_id, IdMappingMode::LOOP)); - } + auto peeled_id = tv->peeledSerialId(); + + // Quick check that the peeled id at schedule + // time is still a leaf domain. + TORCH_INTERNAL_ASSERT( + leaf_id_set.count(peeled_id), + "only exisiting leaf domain supported for peeling\n", + tv->toString(), + " does not have\n", + peeled_id->toString()); + + // Insert the peeled concrete id to the recorded map. + concrete_id_of_peeled_loops_.insert( + GpuLower::current()->caMap()->getConcreteMappedID( + peeled_id, IdMappingMode::LOOP)); } } } @@ -80,10 +79,11 @@ c10::optional PredicatePeelingInfo::getMaybePeeledTileEntry( } for (auto peeled_id : concrete_id_of_peeled_loops_) { - // Need to see the peeled loop to validate this check + // Need to locate the peeled loop to validate this check auto matching_loop_it = concrete_id_to_loop_map.find(peeled_id); if (matching_loop_it != concrete_id_to_loop_map.end()) { // This is the only supported case at the initial stage. + // see also [Supported Case in Predicate Peeling pass] auto split = peeled_id->definition()->as(); if (gpu_lower->caMap()->areMapped( split->in(), root_id, IdMappingMode::EXACT)) { @@ -119,11 +119,16 @@ bool PredicatePeelingInfo::shouldPeelLoop(kir::ForLoop* forloop) const { auto loop_concrete_id = GpuLower::current()->caMap()->getConcreteMappedID( forloop->iter_domain(), IdMappingMode::LOOP); - return !forloop->loopTransformInfo().is_base_index_loop && + return + // Base index calculation loop should not be involved here. + !forloop->loopTransformInfo().is_base_index_loop && concrete_id_of_peeled_loops_.count(loop_concrete_id); } namespace { + +//! A utility class that deep clones the loop body of original_fl +//! and add the cloned expressions into the loop body of new_fl. class LoopNestDeepCloner { public: static void clone(kir::ForLoop* original_fl, kir::ForLoop* new_fl) { @@ -158,6 +163,7 @@ class LoopNestDeepCloner { std::vector cloned_scopes_; }; +//! The predicate peeling transform pass implementation. class PredicatePeeledLoops : kir::ExprMutator { public: static std::vector run(const std::vector exprs) { @@ -169,9 +175,10 @@ class PredicatePeeledLoops : kir::ExprMutator { private: using kir::ExprMutator::handle; + // Create the predicate peeled prolog loop kir::ForLoop* createPeeledLoop(kir::ForLoop* fl) { // Make clone of the outermost loop, but - // only the first iteration. + // limit the loop to the first iteration. auto peeled_loop = IrBuilder::create( fl->iter_domain(), fl->index(), @@ -187,6 +194,7 @@ class PredicatePeeledLoops : kir::ExprMutator { return peeled_loop; } + // Create the predicate peeled main loop kir::ForLoop* createMainLoop(kir::ForLoop* fl) { auto start = SimplifyingIrBuilder::addExpr(fl->start(), fl->kernel()->oneVal()); @@ -208,10 +216,15 @@ class PredicatePeeledLoops : kir::ExprMutator { void handle(kir::ForLoop* fl) final { kir::ExprMutator::handle(fl); + if (GpuLower::current()->predicatePeelingInfo().shouldPeelLoop(fl)) { auto loop_concrete_id = GpuLower::current()->caMap()->getConcreteMappedID( fl->iter_domain(), IdMappingMode::LOOP); + + // Peel this loop if shouldPeel is true and it + // hasn't been processed by this pass already. if (!peeled_iterdomains_.count(loop_concrete_id)) { + // Create and insert the peeled prolog. auto peeled_loop = createPeeledLoop(fl); registerInsertBefore(fl, peeled_loop); @@ -238,10 +251,12 @@ class PredicatePeeledLoops : kir::ExprMutator { } private: - // This pass runs after double buffering pass, so - // we may encounter forloops corresponding to the - // same loop domain twice. In this case we only need - // to peel the first occurrence (the prolog). + //! Keeps track of loop concrete iterdomains that has already been + //! transformed by the loop peeling pass. + //! This pass runs after double buffering pass, so + //! we may encounter forloops corresponding to the + //! same loop domain twice. In this case we only need + //! to peel the first occurrence (the prolog). std::unordered_set peeled_iterdomains_; }; @@ -252,7 +267,9 @@ std::vector PredicatePeeling::peelPredicatedLoop( return PredicatePeeledLoops::run(exprs); } -Val* PredicatePeeling::getSplitTileOffset(IterDomain* id, Val* tile_factor) { +Val* PredicatePeeling::getPrologPredicateOffset( + IterDomain* id, + Val* tile_factor) { // Assume X = original extent, // L = tile factor // Offset needs to satisfy: @@ -280,8 +297,10 @@ Val* PredicatePeeling::getSplitTileOffset(IterDomain* id, Val* tile_factor) { Val* PredicatePeeling::getSplitTileMainOffset( IterDomain* id, Val* tile_factor) { + // This offset is to be **subtracted** from the tensor index + // on the predicate peeling main loop. return SimplifyingIrBuilder::subExpr( - tile_factor, PredicatePeeling::getSplitTileOffset(id, tile_factor)); + tile_factor, PredicatePeeling::getPrologPredicateOffset(id, tile_factor)); } } // namespace cuda diff --git a/torch/csrc/jit/codegen/cuda/lower_predicate_peeling.h b/torch/csrc/jit/codegen/cuda/lower_predicate_peeling.h index f1ce87d2dc9b5..8c6a67108b06d 100644 --- a/torch/csrc/jit/codegen/cuda/lower_predicate_peeling.h +++ b/torch/csrc/jit/codegen/cuda/lower_predicate_peeling.h @@ -13,39 +13,159 @@ namespace jit { namespace fuser { namespace cuda { -// Stage, Factor pair: +//! Note: [Predicate Peeling] +//! This is a loop transformation that attempts to eliminate predicate +//! evaluation in a serial loop. +//! +//! A simple example showing how this trick works is, say we have +//! T0 [I(T0.size[0])] -> split(32) -> T0 [Io(ceilDiv(T0.size[0],32)), +//! Ii(32)], which generates the following code: for i in +//! 0..ceilDiv(T0.size[0],32) +//! for j in 0..32: +//! // assume we need to initialize in this kernel +//! T0[i*32+j] = 0; +//! if i*32+j < T0.size[0] +//! T0[i*32+j] ... +//! The above code generates 32 predicates as the predicate is inlined in the +//! inner loop. +//! +//! The simplification trick is to convert the loop into: +//! +//! let ceilMod(a, b) = a %b == 0 ? b : a %b; +//! +//! // peeled residue prolog : (called initial evaluation in cutlass) +//! // Very similar to the original loop except the +//! // outer loop extent and the predicate extent +//! // are modified. +//! +//! for i in 0..1 +//! for j in 0..32: +//! T0[i*32+j] = 0; +//! if i*32+j < ceilMod(T0.size[0], 32) +//! T0[i*32+j] ... +//! // peeled residue main loop +//! // (called steady-state in cutlass) +//! for i in 0..ceilDiv(T0.size[0],32)-1 +//! for j in 0..32: +//! // No need to initialize as we know the predicate +//! // is all true. +//! // This significantly reduces memory instruction +//! // congestion with cp.async kernels. +//! // No longer need to predicate here as +//! // the residue part of the root iterdomain has +//! // been peeled away. +//! T0[i*32+j + ceilMod(T0.size[0],32)] ... +//! +//! Some details on the predicate peeling pass implemented here: +//! 1. The peeled loop is separate into 2 `PredicatePeelingStage`'s: +//! The first iteration is peeled and marked as +//! PredicatePeelingStage::Prolog, while +//! the rest of the iterations are PredicatePeelingStage::Main. +//! +//! 2. The predicate indexing at the (predicate peeling) prolog is modified to +//! make the access within the residue tile +//! +//! 3. The address indexing at the (predicate peeling) main loop is modified +//! by adding the residue tile as offset. +//! +//! 4. The initialization within (predicate peeling) main loop can be lifted +//! out of the main loop if there are no other not-unrolled serial loops. +//! +//! Note: [Supported Case in Predicate Peeling pass]: +//! The predicate peeling transform is a very specialized pattern used in matmul +//! and some non-trivial overhead would be involved to generalize. +//! +//! The current support for predicate peeling is for a very specific case only +//! and some consideration is needed regarding whether more complex peeling +//! pattern along this line should be pursued. +//! +//! The only supported pattern now is: +//! tile_o, tile_i = split(root_id, inner_factor); +//! where tile_o is required to be on the leaf domain and is where the loop +//! peeling primitive should be applied. +//! +//! The inner_factor is required to be a compile-time constant. +//! +// Note: [Predicate Peeling Interaction with Circular Buffering] +//! +//! 1. In the case where the original loop is double buffered, the first +//! iteration of the double buffer prolog loop is used as +//! PredicatePeelingStage::Prolog and the rest are labeled as +//! PredicatePeelingStage::Main. +//! +//! 2. If a tv is double buffered or circular buffered, the gmem load stage is +//! (stage_depth-1) iterations ahead, so would need to add an extra (simpler) +//! predicate to avoid out-of-bound access. +//! +//! 3. A circular buffer init prolog is added in the case of a predicate tiled +//! and circular buffered loop, as the circular buffer loop prolog only +//! prefetches up to iteration `stage_depth-1`, and if the initialization were +//! to be lifted out of the main loop stage, would also need to initialize for +//! iteration `stage_depth` to make sure the shared memory buffer is all zero +//! initialized. + +//! A data structure used by PredicatePeelingInfo to communicate which +//! for loop is predicate peeled along with the peeling stage and +//! original inner tiling factor +//! TODO: some info here is redundant now. struct PeeledTileEntry { + //! The peeling stage, see note above. PredicatePeelStage peel_stage = PredicatePeelStage::NoApplicable; + + //! The original splitting factor, see [Supported Case in Predicate Peeling + //! pass]. Val* inner_factor = nullptr; + + //! The actual for loop that is predicate peeled. kir::ForLoop* for_loop = nullptr; }; +//! Keeps track fo predicate peeled loops requested +//! from scheduler. class PredicatePeelingInfo { public: + //! Returns true if predicate peeling is requested by scheduler + //! for the given loop. bool shouldPeelLoop(kir::ForLoop* forloop) const; + //! Collect predicate peeling information from fusion. void build(Fusion* fusion); + //! Returns the peeled entry info if the given loop is predicate + //! peeled and the given root_id matches with the tiled root id. + //! + //! see also [Supported Case in Predicate Peeling pass]. c10::optional getMaybePeeledTileEntry( const std::vector& loops, IterDomain* root_id); + //! Returns true if any iterdomain on the given tv's tensor + //! domain corresponds to a predicate peeled loop. bool hasPeeledId(const TensorView* tv) const; private: + //! Keeps track of loop concrete iterdomains that were predicate + //! peeled. std::unordered_set concrete_id_of_peeled_loops_; }; namespace PredicatePeeling { -// User space check that makes sure the loop can -// actually be peeled to remove predicates. +//! User space check that makes sure the loop can +//! actually be peeled to remove predicates. +//! See also +//! [Supported Case in Predicate Peeling pass]: bool supportedPeelingLoop(IterDomain* id); +//! Kernel IR pass that applies the predicate peeling transformation. std::vector peelPredicatedLoop(const std::vector exprs); -Val* getSplitTileOffset(IterDomain* id, Val* tile_factor); +//! Utility to generate the residual extend used in predicate +//! peeling prolog. +Val* getPrologPredicateOffset(IterDomain* id, Val* tile_factor); +//! Utility to generate the offset applied to tensor indices +//! in predicate peeling main loop. Val* getSplitTileMainOffset(IterDomain* id, Val* tile_factor); } // namespace PredicatePeeling diff --git a/torch/csrc/jit/codegen/cuda/tensor_view.cpp b/torch/csrc/jit/codegen/cuda/tensor_view.cpp index 55e5bf43da9b5..6289eb4343cd2 100644 --- a/torch/csrc/jit/codegen/cuda/tensor_view.cpp +++ b/torch/csrc/jit/codegen/cuda/tensor_view.cpp @@ -222,8 +222,8 @@ TensorView::TensorView(const TensorView* src, IrCloner* ir_cloner) for (const auto id : src->axesToSwizzle()) { axes_to_swizzle_.push_back(ir_cloner->clone(id)); } - for (const auto id : src->peeled_serial_id_) { - peeled_serial_id_.push_back(ir_cloner->clone(id)); + if (src->peeled_serial_id_ != nullptr) { + peeled_serial_id_ = ir_cloner->clone(src->peeled_serial_id_); } } @@ -1218,7 +1218,7 @@ void TensorView::peelPredicatedLoop(int axis_id) { auto id = axis(axis_id); TORCH_CHECK( PredicatePeeling::supportedPeelingLoop(id), "unsupported loop peeling"); - peeled_serial_id_.push_back(id); + peeled_serial_id_ = id; } TensorViewBuilder& TensorViewBuilder::ndims(size_t ndims) { diff --git a/torch/csrc/jit/codegen/cuda/test/test_gpu.cpp b/torch/csrc/jit/codegen/cuda/test/test_gpu.cpp index 00c43c7f7f185..acff7115851cc 100644 --- a/torch/csrc/jit/codegen/cuda/test/test_gpu.cpp +++ b/torch/csrc/jit/codegen/cuda/test/test_gpu.cpp @@ -23024,9 +23024,6 @@ TEST_F(NVFuserTest, FusionPredicatePeeling1_CUDA) { tv2->axis(2)->parallelize(ParallelType::TIDy); tv2->peelPredicatedLoop(0); - fusion.printMath(); - fusion.printKernel(); - auto options = at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0); at::Tensor t0 = at::randn({m, n}, options); From 29d4db1a7fa9e4b8921b26f75e846146d5eac786 Mon Sep 17 00:00:00 2001 From: Xiang Gao Date: Thu, 29 Sep 2022 20:35:22 -0700 Subject: [PATCH 12/14] save --- torch/csrc/jit/codegen/cuda/lower_double_buffer.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/torch/csrc/jit/codegen/cuda/lower_double_buffer.cpp b/torch/csrc/jit/codegen/cuda/lower_double_buffer.cpp index fc1fe5b5853ed..05589386c863f 100644 --- a/torch/csrc/jit/codegen/cuda/lower_double_buffer.cpp +++ b/torch/csrc/jit/codegen/cuda/lower_double_buffer.cpp @@ -322,7 +322,7 @@ class DoubleBufferLoopCloner : public kir::IrVisitor { // the loop defining out_tv as there might be multiple // loops that realize double buffers. bool db_loop_found = false; - auto& ca_map = GpuLower::current()->caMap(); + const auto& ca_map = GpuLower::current()->caMap(); if (!(out_tv->isDoubleBuffered() || out_tv->isCircularBuffered()) || !ca_map->areMapped( From 5100949a4ba9ca7e2a4805b2b3d985964b2b28b9 Mon Sep 17 00:00:00 2001 From: Xiang Gao Date: Sun, 4 Dec 2022 21:48:37 -0800 Subject: [PATCH 13/14] fix --- torch/csrc/jit/codegen/cuda/tensor_view.cpp | 1 - 1 file changed, 1 deletion(-) diff --git a/torch/csrc/jit/codegen/cuda/tensor_view.cpp b/torch/csrc/jit/codegen/cuda/tensor_view.cpp index 4dd712455b057..a92ef481687a4 100644 --- a/torch/csrc/jit/codegen/cuda/tensor_view.cpp +++ b/torch/csrc/jit/codegen/cuda/tensor_view.cpp @@ -221,7 +221,6 @@ TensorView::TensorView(const TensorView* src, IrCloner* ir_cloner) lift_write_address_(src->lift_write_address_), skew_double_buffer_loop_(src->skew_double_buffer_loop_), peeled_serial_id_(ir_cloner->clone(src->peeled_serial_id_)), - skew_double_buffer_loop_(src->skew_double_buffer_loop_), compute_with_consumers_(ir_cloner->clone(src->compute_with_consumers_)), compute_with_pos_(src->compute_with_pos_) {} From 9b768746204fc0dfcea3143896549056fc1a07e8 Mon Sep 17 00:00:00 2001 From: Xiang Gao Date: Sun, 4 Dec 2022 21:49:34 -0800 Subject: [PATCH 14/14] fix --- torch/csrc/jit/codegen/cuda/index_compute.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/torch/csrc/jit/codegen/cuda/index_compute.cpp b/torch/csrc/jit/codegen/cuda/index_compute.cpp index 4db21b616f711..e2d64b75f8215 100644 --- a/torch/csrc/jit/codegen/cuda/index_compute.cpp +++ b/torch/csrc/jit/codegen/cuda/index_compute.cpp @@ -1680,7 +1680,7 @@ std::vector Index::getGlobalProducerStridedIndices( continue; } else { if (auto tile_entry = - gpu_lower->predicatePeelingInfo().getMaybePeeledTileEntry( + GpuLower::current()->predicatePeelingInfo().getMaybePeeledTileEntry( loops, root_dom[i])) { // Add the "predicate peeling offset", see [Predicate Peeling] // to the tensor index if this root domain is predicate peeled.