diff --git a/csrc/contiguity.cpp b/csrc/contiguity.cpp index b04bcacb079..faf6a2857a4 100644 --- a/csrc/contiguity.cpp +++ b/csrc/contiguity.cpp @@ -441,8 +441,7 @@ NonDivisibleSplitDependencies::NonDivisibleSplitDependencies( for (auto transform : transforms) { auto inp_ids = ir_utils::filterByType(transform->inputs()); for (auto inp_id : inp_ids) { - if (std::find(alloc_domain.begin(), alloc_domain.end(), inp_id) != - alloc_domain.end()) { + if (std::ranges::find(alloc_domain, inp_id) != alloc_domain.end()) { // This generally shouldn't happen as there shouldn't be // transformations before the allocation ids, but in case for some // reason we eventually do have cases like that, we should reset the @@ -453,7 +452,7 @@ NonDivisibleSplitDependencies::NonDivisibleSplitDependencies( } bool inputs_non_divisible = - std::any_of(inp_ids.begin(), inp_ids.end(), [this](IterDomain* inp_id) { + std::ranges::any_of(inp_ids, [this](IterDomain* inp_id) { return depends_on_non_divisible_split.find(inp_id) != depends_on_non_divisible_split.end(); }); @@ -575,8 +574,7 @@ void ContigIDs::build(const std::vector& ids) { alloc_domain_id->toString()); // Index of merged reductions can always be coalesced, so considering // reduction as true contiguity. - if (alloc_contiguity.value_or(true) && - alloc_domain_id->getIterType() != IterType::GatherScatter) { + if (alloc_contiguity.value_or(true)) { contig_ids_.emplace(alloc_domain_id); is_contig_alloc_.at(alloc_domain_id) = true; within_contig_ids_[alloc_domain_id] = std::unordered_set(); diff --git a/csrc/device_lower/validation.cpp b/csrc/device_lower/validation.cpp index 29640300147..9fea234f204 100644 --- a/csrc/device_lower/validation.cpp +++ b/csrc/device_lower/validation.cpp @@ -1815,11 +1815,9 @@ void validateAndConvertIterDomainGrouping(Fusion* fusion) { // Remember if a grouped ID is found is_grouped = true; - // Grouping only makes sense for the normal iteration or gather scatter - // type + // Grouping only makes sense for the normal iteration type NVF_CHECK( - id->getIterType() == IterType::Iteration || - id->getIterType() == IterType::GatherScatter, + id->getIterType() == IterType::Iteration, "Invalid use of ParallelType::Group.", " Grouping of ", id->getIterType(), diff --git a/csrc/fusion_segmenter.cpp b/csrc/fusion_segmenter.cpp index 1855d35f65e..a200a22e13c 100644 --- a/csrc/fusion_segmenter.cpp +++ b/csrc/fusion_segmenter.cpp @@ -1474,16 +1474,6 @@ void eraseInputDistinctRootDomains(Fusion* fusion) { compare_result.dom1_has_unreachable_ids; if (has_disjoint_loop_logical) { - // NOTE: This is only the case for scatter outputs, for which loop and - // logical are disjoint. Consequently, the loop domain cannot be replayed. - // Since this scatter output is a fusion input to this segment, its loop - // domain is immaterial now and we can skip replaying it. - NVF_ERROR( - std::ranges::any_of( - tv->getLogicalDomain(), - [](IterDomain* id) { return id->isGatherScatter(); }), - "Disjoint loop and logical are only permitted for scatter outputs, ", - tv->domain()->toString(0, false)); NVF_ERROR( !isSharded(tv), "Sharding is not permitted when loop domain is disjoint from logical " diff --git a/csrc/id_model/contiguity.cpp b/csrc/id_model/contiguity.cpp index 20391471ec6..de499b6ff7b 100644 --- a/csrc/id_model/contiguity.cpp +++ b/csrc/id_model/contiguity.cpp @@ -50,8 +50,7 @@ ContigIDGroups::ContigIDGroups( auto alloc_contiguity = alloc_contiguity_.at(index_domain_i); - if (alloc_contiguity && - index_domain->getIterType() != IterType::GatherScatter) { + if (alloc_contiguity) { contig_ids_.emplace(graph_.toGroup(index_domain)); } } @@ -64,7 +63,7 @@ ContigIDGroups::ContigIDGroups( const auto outputs = direction == Direction::Forward ? graph_.outputGroups(eg) : graph_.inputGroups(eg); - if (std::any_of(inputs.begin(), inputs.end(), [&](const ValGroup& inp) { + if (std::ranges::any_of(inputs, [&](const ValGroup& inp) { return resize_deps_.count(inp) > 0; })) { for (const auto& out : outputs) { @@ -72,7 +71,7 @@ ContigIDGroups::ContigIDGroups( } } - if (std::any_of(inputs.begin(), inputs.end(), [&](const ValGroup& inp) { + if (std::ranges::any_of(inputs, [&](const ValGroup& inp) { return non_divisible_deps_.count(inp) > 0; })) { for (const auto& out : outputs) { @@ -169,10 +168,8 @@ void ContigIDGroups::handle(Split* split, Direction direction) { if (direction == Direction::Forward) { const auto& divisible_splits = GpuLower::current()->divisibleSplitSet(); const ExprGroup& split_group = graph_.toGroup(split); - bool divisible = std::any_of( - divisible_splits.begin(), - divisible_splits.end(), - [&](Split* divisible_split) -> bool { + bool divisible = std::ranges::any_of( + divisible_splits, [&](Split* divisible_split) -> bool { return split_group->has(divisible_split); }); if (!divisible) { diff --git a/csrc/ir/internal_base_nodes.cpp b/csrc/ir/internal_base_nodes.cpp index 93691cf70c3..65c80482948 100644 --- a/csrc/ir/internal_base_nodes.cpp +++ b/csrc/ir/internal_base_nodes.cpp @@ -398,12 +398,6 @@ IterDomain* IterDomain::merge( inner->getIterType() == IterType::Iteration)) { iter_type = IterType::Iteration; } - - if ((outer->isBroadcast() || inner->isBroadcast()) && - (outer->getIterType() == IterType::GatherScatter || - inner->getIterType() == IterType::GatherScatter)) { - iter_type = IterType::GatherScatter; - } } Val* expanded_extent = nullptr; @@ -700,10 +694,8 @@ void IterDomain::parallelize(ParallelType t) { if (t == ParallelType::Group) { NVF_CHECK( - getIterType() == IterType::Iteration || - getIterType() == IterType::GatherScatter, - "Grouping IterDomain of non Iteration / GatherScatter type is not " - "allowed. ", + getIterType() == IterType::Iteration, + "Grouping IterDomain of non Iteration type is not allowed. ", getIterType()); } diff --git a/csrc/ir/internal_base_nodes.h b/csrc/ir/internal_base_nodes.h index 6e2e5e943a4..e623e29ca87 100644 --- a/csrc/ir/internal_base_nodes.h +++ b/csrc/ir/internal_base_nodes.h @@ -201,10 +201,6 @@ class NVF_API IterDomain : public Val { return getIterType() == IterType::Symbolic; } - bool isGatherScatter() const { - return getIterType() == IterType::GatherScatter; - } - bool isStride() const { return getIterType() == IterType::Stride; } diff --git a/csrc/ops/indexing.cpp b/csrc/ops/indexing.cpp index b18467b9928..471f51a26a1 100644 --- a/csrc/ops/indexing.cpp +++ b/csrc/ops/indexing.cpp @@ -176,13 +176,7 @@ TensorView* gather(TensorView* inp, int64_t dim, TensorView* index) { std::vector out_domain; out_domain.reserve(idx_domain.size()); for (auto idx_domain_ptr : idx_domain) { - out_domain.push_back( - IterDomainBuilder(idx_domain_ptr) - .iter_type( - idx_domain_ptr->getIterType() == IterType::Iteration - ? IterType::GatherScatter - : idx_domain_ptr->getIterType()) - .build()); + out_domain.push_back(idx_domain_ptr->cloneWithoutRFactor()); } TensorView* out_tensor = IrBuilder::create( @@ -254,13 +248,7 @@ TensorView* scatter( // The shape of output tensor is same as self tensor. std::vector out_logical; for (const auto i : arange(self_dom.size())) { - out_logical.push_back( - IterDomainBuilder(self_dom[i]) - .iter_type( - self_dom[i]->getIterType() == IterType::Iteration - ? IterType::GatherScatter - : self_dom[i]->getIterType()) - .build()); + out_logical.push_back(self_dom[i]->cloneWithoutRFactor()); } // Create the loop domain based on the logical domain of the index diff --git a/csrc/ops/utils.cpp b/csrc/ops/utils.cpp index 8dd19be4c62..e7f5d86e27a 100644 --- a/csrc/ops/utils.cpp +++ b/csrc/ops/utils.cpp @@ -151,7 +151,6 @@ IterType promoteIterType(IterType type1, IterType type2) { // Iteration: Default // Reduction: Should not appear here // Broadcast: Propagated only if type1 and type2 are Broadcast - // GatherScatter: Converted to Iteration // Stride: Shold not appear here // VectorComponent: Converted to Iteration @@ -164,11 +163,11 @@ IterType promoteIterType(IterType type1, IterType type2) { "Invalid IterType: ", type2); - // Do not propagate GatherScatter and VectorComponent - if (type1 == IterType::VectorComponent || type1 == IterType::GatherScatter) { + // Do not propagate VectorComponent + if (type1 == IterType::VectorComponent) { type1 = IterType::Iteration; } - if (type2 == IterType::VectorComponent || type2 == IterType::GatherScatter) { + if (type2 == IterType::VectorComponent) { type2 = IterType::Iteration; } @@ -408,8 +407,6 @@ IterDomain* newOutputIterDomain( extent_val = promoteSize(extent_val, id->extent()); if (iter_type.has_value()) { iter_type = promoteIterType(iter_type.value(), id->getIterType()); - } else if (id->isGatherScatter()) { - iter_type = IterType::Iteration; } else { iter_type = id->getIterType(); } diff --git a/csrc/runtime/fusion_executor_cache.h b/csrc/runtime/fusion_executor_cache.h index ea5f972e3c4..e361404b0bf 100644 --- a/csrc/runtime/fusion_executor_cache.h +++ b/csrc/runtime/fusion_executor_cache.h @@ -7,6 +7,7 @@ // clang-format on #pragma once +#include #include #include #include @@ -28,7 +29,7 @@ class ExactLogicalDomainMap; class Fusion; class FusionKernelRuntime; class KernelArgumentHolder; -enum class PrimDataType; +enum class PrimDataType : std::uint8_t; //! [ Note -- Post-definition cache implementation ] //! diff --git a/csrc/runtime/fusion_kernel_runtime.h b/csrc/runtime/fusion_kernel_runtime.h index 6c3029b973d..37a75e15c13 100644 --- a/csrc/runtime/fusion_kernel_runtime.h +++ b/csrc/runtime/fusion_kernel_runtime.h @@ -7,6 +7,7 @@ // clang-format on #pragma once +#include #include #include @@ -23,7 +24,7 @@ namespace nvfuser { class HeuristicParamsList; -enum class PrimDataType; +enum class PrimDataType : std::uint8_t; class Fusion; class Val; diff --git a/csrc/type.cpp b/csrc/type.cpp index 7b4ba2d9d5f..ff7a3932b82 100644 --- a/csrc/type.cpp +++ b/csrc/type.cpp @@ -37,26 +37,26 @@ StructType globalTensorMetaData( StructType::FieldInfo logical_size_field; logical_size_field.name = "logical_size"; - logical_size_field.type = std::make_shared( - ArrayType{std::make_shared(DataType::Index), dim}); + logical_size_field.type = std::make_shared(ArrayType{ + .type = std::make_shared(DataType::Index), .size = dim}); logical_size_field.used_in_kernel = true; StructType::FieldInfo logical_stride_field; logical_stride_field.name = "logical_stride"; - logical_stride_field.type = std::make_shared( - ArrayType{std::make_shared(DataType::Index), dim}); + logical_stride_field.type = std::make_shared(ArrayType{ + .type = std::make_shared(DataType::Index), .size = dim}); logical_stride_field.used_in_kernel = false; StructType::FieldInfo alloc_size_field; alloc_size_field.name = "alloc_size"; - alloc_size_field.type = std::make_shared( - ArrayType{std::make_shared(DataType::Index), alloc_dim}); + alloc_size_field.type = std::make_shared(ArrayType{ + .type = std::make_shared(DataType::Index), .size = alloc_dim}); alloc_size_field.used_in_kernel = false; StructType::FieldInfo alloc_stride_field; alloc_stride_field.name = "alloc_stride"; - alloc_stride_field.type = std::make_shared( - ArrayType{std::make_shared(DataType::Index), alloc_dim}); + alloc_stride_field.type = std::make_shared(ArrayType{ + .type = std::make_shared(DataType::Index), .size = alloc_dim}); alloc_stride_field.used_in_kernel = true; return StructType::make( @@ -884,8 +884,6 @@ static const char* iter_type2string(IterType t) { return "b"; case IterType::Stride: return "s"; - case IterType::GatherScatter: - return "n"; case IterType::VectorComponent: return "v"; case IterType::Symbolic: @@ -1393,7 +1391,6 @@ at::ScalarType data_type_to_aten(const DataType& data_type) { return at::ScalarType::Float8_e8m0fnu; #if NVF_TORCH_VERSION_NO_LESS(2, 8, 0) case DataType::Float4_e2m1fn_x2: - return at::ScalarType::Float4_e2m1fn_x2; case DataType::Float4_e2m1fn: return at::ScalarType::Float4_e2m1fn_x2; #endif @@ -1434,9 +1431,7 @@ at::ScalarType data_type_to_aten(const DataType& data_type) { // there is no direct mapping, we use some data type as a proxy. // If there is a data type with the same size, we use that const int64_t size_bit = dataTypeSizeBit(data_type); - if (size_bit == 8) { - return at::ScalarType::Byte; - } else if (size_bit == 16) { + if (size_bit == 16) { return at::ScalarType::UInt16; } else if (size_bit == 32) { return at::ScalarType::UInt32; @@ -1465,7 +1460,7 @@ at::ScalarType data_type_to_aten( AdjustLastDim getLastDimAdjustment(const DataType& dtype) { if (dtype == DataType::Index) { - return AdjustLastDim{1, 1}; + return AdjustLastDim{.numerator = 1, .denominator = 1}; } const int64_t scalar_type_bit = (int64_t)c10::elementSize(data_type_to_aten(dtype)) * 8; @@ -1475,7 +1470,8 @@ AdjustLastDim getLastDimAdjustment(const DataType& dtype) { // at_size * 4 / 3 is the size of the last dimension of the corresponding // TensorView. const int64_t gcd = std::gcd(scalar_type_bit, dtype_bit); - return AdjustLastDim{scalar_type_bit / gcd, dtype_bit / gcd}; + return AdjustLastDim{ + .numerator = scalar_type_bit / gcd, .denominator = dtype_bit / gcd}; } std::ostream& operator<<(std::ostream& out, const ValType vtype) { @@ -1823,9 +1819,8 @@ int max_digits10(DataType dtype) { } else if (dtype == DataType::Float8_e4m3fn) { return 3; } else if ( - dtype == DataType::Float8_e5m2 || dtype == DataType::Float8_e8m0fnu) { - return 2; - } else if (dtype == DataType::Float4_e2m1fn) { + dtype == DataType::Float8_e5m2 || dtype == DataType::Float8_e8m0fnu || + dtype == DataType::Float4_e2m1fn) { return 2; } else { NVF_CHECK( diff --git a/csrc/type.h b/csrc/type.h index ea40db96d09..9726291c4c6 100644 --- a/csrc/type.h +++ b/csrc/type.h @@ -42,7 +42,7 @@ namespace nvfuser { // Order of strength -enum class ValType { +enum class ValType : std::uint8_t { TensorDomain, IterDomain, RaggedIterDomain, @@ -62,7 +62,7 @@ enum class ValType { // ElectSync - Select a single thread to launch asynchronous operations. // OneDimTmaLoadExpectArrive - Predicate for expect arrive bytes and 1D TMA // load. OneDimTmaWaitParity - Predicate for wait parity for 1D TMA load. -enum class PredicateType { +enum class PredicateType : std::uint8_t { Manual, Inline, Unswitch, @@ -79,7 +79,7 @@ enum class PredicateType { // type might be. This allows us to prevent assuming the welford count must be // int64_t which is relatively heavy to carry around. Index will be resolved // at compile time with KernelIndexMode. -enum class PrimDataType { +enum class PrimDataType : std::uint8_t { // Floating point types Double, Float, @@ -145,7 +145,7 @@ struct StructType { template static StructType make(std::vector fields, std::string name = "") { static_assert( - std::is_base_of::value, + std::is_base_of_v, "StructType::make only accepts Struct types"); return StructType{ .name = std::move(name), @@ -271,7 +271,7 @@ class Val; //! Get the type of a Val's metadata, currently only supporting tensors NVF_API DataType metaDataTypeOf(const Val* tv); -enum class KernelIndexMode { INT32, INT64 }; +enum class KernelIndexMode : std::uint8_t { INT32, INT64 }; PrimDataType indexModeToDtype(KernelIndexMode index_mode); KernelIndexMode indexTypeToMode(DataType index_type); @@ -529,7 +529,7 @@ inline bool hasCompatibleDataType( //! binary->text->binary round-trip. For exact types, this function returns 0. int max_digits10(DataType dtype); -enum class UnaryOpType { +enum class UnaryOpType : std::uint8_t { Cast, BitCast, RefCast, @@ -603,7 +603,7 @@ enum class UnaryOpType { // TODO: Order of this list is important as it affects type promotion. it's not // in the right order now. -enum class BinaryOpType { +enum class BinaryOpType : std::uint8_t { // Math Ops Add, Atan2, @@ -652,7 +652,7 @@ enum class BinaryOpType { Complex }; -enum class RNGOpType { +enum class RNGOpType : std::uint8_t { Uniform, // Uniform in [0, 1) UniformRange, // Uniform in [low, high] NormalStandard, // Normal with mean 0, std 1 @@ -666,9 +666,15 @@ bool isIntegerOp(const BinaryOpType bopt); // Return if output of operator should be a boolean bool isLogicalOp(const BinaryOpType bopt); -enum class TernaryOpType { Clamp, Lerp, Threshold, Where, Philox }; +enum class TernaryOpType : std::uint8_t { + Clamp, + Lerp, + Threshold, + Where, + Philox +}; -enum class ParallelType { +enum class ParallelType : std::uint8_t { DIDx = 0, DIDy, DIDz, @@ -721,21 +727,26 @@ static constexpr std::array kParallelTypeDIDs = { ParallelType::DIDy, ParallelType::DIDz}; -enum class MemoryType { Local, Shared, Global, Tensor, Symmetric }; +enum class MemoryType : std::uint8_t { + Local, + Shared, + Global, + Tensor, + Symmetric +}; // Symbolic: Undetermined between Iteration or Broadcast -enum class IterType { +enum class IterType : std::uint8_t { Iteration, Reduction, Broadcast, Stride, - GatherScatter, VectorComponent, Symbolic }; // Used for Iteration Domain mapping modes in ComputeAtMap -enum class IdMappingMode { +enum class IdMappingMode : std::uint8_t { EXACT, ALMOSTEXACT, BROADCAST, @@ -760,7 +771,7 @@ static constexpr std::array kIdMappingModes = { // https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#cache-operators // for what each option means. Will also consider .L1::no_allocate because .cs // still pollutes cache to some extent. -enum class CacheOp { +enum class CacheOp : std::uint8_t { Unspecified, // Opt in for the default cache operator or when the LoadStoreOp // doesn't take a cache operator. AllLevels, @@ -773,7 +784,7 @@ enum class CacheOp { //! //! SegmenterSet here is used to hint segmenter to break kernel on the output //! of the node -enum class LoadStoreOpType { +enum class LoadStoreOpType : std::uint8_t { Set, SegmenterSet, LdMatrix, @@ -787,7 +798,7 @@ enum class LoadStoreOpType { // Used to label what part of the circular buffered iterdomain // a for loop is materializing. -enum class CircularBufferLoopStage { +enum class CircularBufferLoopStage : std::uint8_t { Prolog = 0, Main, Epilog, @@ -833,10 +844,10 @@ inline bool mayHaveWarHazard(CircularBufferLoopStage stage) { //! //! TODO: unify with existing swizzle logic, currently //! doesn't have the same type. -enum class SwizzleType { NoSwizzle = 0, XOR, CyclicShift }; +enum class SwizzleType : std::uint8_t { NoSwizzle = 0, XOR, CyclicShift }; //! Modes of swizzle, see [Note on swizzle mode]. -enum class SwizzleMode { NoSwizzle = 0, Data, Loop }; +enum class SwizzleMode : std::uint8_t { NoSwizzle = 0, Data, Loop }; // Returns if function needs an f suffix on the operator when operating on a // float value i.e. sin->sinf @@ -1100,12 +1111,12 @@ constexpr inline size_t primDataTypeSizeBit(PrimDataType type) { } constexpr inline size_t primDataTypeSizeByte(PrimDataType type) { - int64_t bits = primDataTypeSizeBit(type); + size_t bits = primDataTypeSizeBit(type); NVF_CHECK(bits % 8 == 0, "Size is not a multiple of 8 bits."); return bits / 8; } -enum class LaunchConfigType { +enum class LaunchConfigType : std::uint8_t { Compatible, SharedMemory, BIDz, @@ -1157,7 +1168,7 @@ constexpr auto toUnderlying(E e) noexcept { return static_cast>(e); } -enum class AsyncOpType { NotAsync, CpAsync, CpAsyncBulk, WgMma }; +enum class AsyncOpType : std::uint8_t { NotAsync, CpAsync, CpAsyncBulk, WgMma }; // Data path between TMem and register file. Tensor memory is not a general // byte-addressable memory like other memory types. The register <-> TMem @@ -1165,7 +1176,7 @@ enum class AsyncOpType { NotAsync, CpAsync, CpAsyncBulk, WgMma }; // well-defined specification about which thread's which register access to // which part of TMem. See: // https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#tcgen05-memory-layout -enum class TMemRegisterDataPath { +enum class TMemRegisterDataPath : std::uint8_t { Path32x32b, Path16x64b, Path16x128b, @@ -1180,7 +1191,7 @@ std::ostream& operator<<(std::ostream&, cudaDriverEntryPointQueryResult); // Layout for block scaling factor used by mx-format with narrow precision, this // indicates how to index into block scaling factor. see: // https://docs.nvidia.com/cutlass/media/docs/cpp/blackwell_functionality.html#scale-factor-layouts -enum class BlockScalingFactorLayout { +enum class BlockScalingFactorLayout : std::uint8_t { Block128x4, }; diff --git a/tests/cpp/test_scatter.cpp b/tests/cpp/test_scatter.cpp index 5e75131f92f..cad1c42db35 100644 --- a/tests/cpp/test_scatter.cpp +++ b/tests/cpp/test_scatter.cpp @@ -23,31 +23,6 @@ namespace nvfuser { -TEST_F(NVFuserTest, GatherScatterIterType) { - auto fusion_ptr = std::make_unique(); - Fusion& fusion = *fusion_ptr.get(); - FusionGuard fg(&fusion); - - auto tv0 = makeContigTensor(1, DataType::Int); - fusion.addInput(tv0); - auto tv1 = makeContigTensor(1, DataType::Int); - fusion.addInput(tv1); - auto tv2 = - scatter(tv0, 0, tv1, fusion.oneVal(DataType::Int), BinaryOpType::Add); - auto tv3 = add(tv2, fusion.oneVal()); - fusion.addOutput(tv3); - - for (auto tv : fusion.allTvs()) { - if (tv->definition() != nullptr && tv->definition()->isA()) { - continue; - } - EXPECT_TRUE(std::none_of( - tv->getLoopDomain().begin(), - tv->getLoopDomain().end(), - [](const IterDomain* id) { return id->isGatherScatter(); })); - } -} - using ScatterTestConfig = bool; // manual_scheduling class ScatterTest : public NVFuserFixtureParamTest {