From e5b0594d17494cee88a38f2cf99e583464cc6bc4 Mon Sep 17 00:00:00 2001 From: root <26priya11@gmail.com> Date: Wed, 15 Jan 2025 18:45:52 -0800 Subject: [PATCH] rename to embedding_fwd --- csrc/device_lower/utils.cpp | 2 +- csrc/dispatch.h | 2 +- csrc/ir/internal_nodes.h | 6 +++--- csrc/ir/nodes.cpp | 10 +++++----- csrc/logical_domain_map.cpp | 2 +- csrc/ops/composite.cpp | 4 ++-- csrc/ops/composite.h | 2 +- csrc/python_frontend/fusion_record.h | 14 +++++++------- csrc/python_frontend/python_bindings.cpp | 6 +++--- csrc/scheduler/expr_eval_sched.cpp | 2 +- csrc/serde/fusion_cache.fbs | 2 +- csrc/serde/fusion_record.cpp | 6 +++--- tests/cpp/test_embedding_node.cpp | 4 ++-- tests/python/test_embedding.py | 4 ++-- 14 files changed, 33 insertions(+), 33 deletions(-) diff --git a/csrc/device_lower/utils.cpp b/csrc/device_lower/utils.cpp index c782b22f658..3c26d0205cd 100644 --- a/csrc/device_lower/utils.cpp +++ b/csrc/device_lower/utils.cpp @@ -159,7 +159,7 @@ bool isTvOp(const Expr* expr) { LinearOp, SdpaFwdOp, SdpaBwdOp, - EmbeddingOp, + EmbeddingFwdOp, BroadcastOp, SqueezeOp, ExpandOp, diff --git a/csrc/dispatch.h b/csrc/dispatch.h index b42772e34f3..4798d3862fa 100644 --- a/csrc/dispatch.h +++ b/csrc/dispatch.h @@ -110,7 +110,7 @@ class Val; f(LinearOp); \ f(SdpaFwdOp); \ f(SdpaBwdOp); \ - f(EmbeddingOp); \ + f(EmbeddingFwdOp); \ f(Communication); \ f(ForLoop); \ f(P2PCommunication); diff --git a/csrc/ir/internal_nodes.h b/csrc/ir/internal_nodes.h index ea8c3722d3a..af3dff6c73b 100644 --- a/csrc/ir/internal_nodes.h +++ b/csrc/ir/internal_nodes.h @@ -2679,11 +2679,11 @@ class SdpaBwdOp : public Expr { const std::vector& inputs) const override; }; -class EmbeddingOp : public Expr { +class EmbeddingFwdOp : public Expr { public: using Expr::Expr; - EmbeddingOp( + EmbeddingFwdOp( IrBuilderPasskey, TensorView* output, TensorView* input, @@ -2697,7 +2697,7 @@ class EmbeddingOp : public Expr { NVFUSER_DECLARE_CLONE_AND_CREATE const char* getOpString() const override { - return "EmbeddingOp"; + return "EmbeddingFwdOp"; } std::string toString(int indent_size = 0) const override; diff --git a/csrc/ir/nodes.cpp b/csrc/ir/nodes.cpp index 9aa6ec5cb88..546825bc45d 100644 --- a/csrc/ir/nodes.cpp +++ b/csrc/ir/nodes.cpp @@ -5257,7 +5257,7 @@ std::vector SdpaBwdOp::evaluate( slice_last_dim(grad_value)}; } -EmbeddingOp::EmbeddingOp( +EmbeddingFwdOp::EmbeddingFwdOp( IrBuilderPasskey passkey, TensorView* output, TensorView* input, @@ -5289,9 +5289,9 @@ EmbeddingOp::EmbeddingOp( } } -NVFUSER_DEFINE_CLONE_AND_CREATE(EmbeddingOp) +NVFUSER_DEFINE_CLONE_AND_CREATE(EmbeddingFwdOp) -std::string EmbeddingOp::toString(int indent_size) const { +std::string EmbeddingFwdOp::toString(int indent_size) const { std::stringstream ss; indent(ss, indent_size) << out()->toString() << ",\n"; indent(ss, indent_size + 1) @@ -5314,11 +5314,11 @@ std::string EmbeddingOp::toString(int indent_size) const { return ss.str(); } -std::string EmbeddingOp::toInlineString(int indent_size) const { +std::string EmbeddingFwdOp::toInlineString(int indent_size) const { NVF_CHECK(false, "Tensor op can not be printed inline"); } -std::vector EmbeddingOp::evaluate( +std::vector EmbeddingFwdOp::evaluate( const ExpressionEvaluator& ee, const std::vector& inputs) const { diff --git a/csrc/logical_domain_map.cpp b/csrc/logical_domain_map.cpp index e7288e73d7c..08bd787bb12 100644 --- a/csrc/logical_domain_map.cpp +++ b/csrc/logical_domain_map.cpp @@ -325,7 +325,7 @@ std::unordered_map PairwiseLogicalDomainMap::map( return dom_map; } - if (EmbeddingOp* op = dynamic_cast(consumer_tv_->definition())) { + if (EmbeddingFwdOp* op = dynamic_cast(consumer_tv_->definition())) { // Producers: // input = [*] // weight = [V, embedding_dim] diff --git a/csrc/ops/composite.cpp b/csrc/ops/composite.cpp index 148b4ce506c..dd9bffffe43 100644 --- a/csrc/ops/composite.cpp +++ b/csrc/ops/composite.cpp @@ -609,7 +609,7 @@ SdpfaBwdResult sdpfa_bwd( return {grad_query, grad_key, grad_value}; } -TensorView* embedding( +TensorView* embedding_fwd( TensorView* input, TensorView* weight, Val* padding_idx, @@ -656,7 +656,7 @@ TensorView* embedding( if (sparse == nullptr){ sparse = IrBuilder::create(false, DataType::Bool); } - IrBuilder::create( + IrBuilder::create( output, input, weight, diff --git a/csrc/ops/composite.h b/csrc/ops/composite.h index 036a14bb061..4063ca08d60 100644 --- a/csrc/ops/composite.h +++ b/csrc/ops/composite.h @@ -116,7 +116,7 @@ SdpfaBwdResult sdpfa_bwd( TensorView* philox_offset, Val* scale); -TensorView* embedding( +TensorView* embedding_fwd( TensorView* input, TensorView* weight, Val* padding_idx, diff --git a/csrc/python_frontend/fusion_record.h b/csrc/python_frontend/fusion_record.h index 536e430dd01..8b4e3cd4241 100644 --- a/csrc/python_frontend/fusion_record.h +++ b/csrc/python_frontend/fusion_record.h @@ -3052,16 +3052,16 @@ struct SdpaBwdOpRecord : RecordFunctor { } }; -struct EmbeddingOpRecord : RecordFunctor { - EmbeddingOpRecord(std::vector args, std::vector outputs) +struct EmbeddingFwdOpRecord : RecordFunctor { + EmbeddingFwdOpRecord(std::vector args, std::vector outputs) : RecordFunctor( std::move(args), std::move(outputs), - "ops.embedding", - serde::RecordType::EmbeddingOp) {} - ~EmbeddingOpRecord() override = default; + "ops.embedding_fwd", + serde::RecordType::EmbeddingFwdOp) {} + ~EmbeddingFwdOpRecord() override = default; RecordFunctor* clone() final { - return new EmbeddingOpRecord(*this); + return new EmbeddingFwdOpRecord(*this); } void operator()(FusionState& fd) final { @@ -3083,7 +3083,7 @@ struct EmbeddingOpRecord : RecordFunctor { ? fd.getFusionState(args_.at(6).index)->as() : nullptr; - auto output = embedding(input, weight, padding_idx, max_norm, norm_type, scale_grad_by_freq, sparse); + auto output = embedding_fwd(input, weight, padding_idx, max_norm, norm_type, scale_grad_by_freq, sparse); fd.setFusionState(outputs_.at(0).index, output); } }; diff --git a/csrc/python_frontend/python_bindings.cpp b/csrc/python_frontend/python_bindings.cpp index de3a3d51fb2..533d9f991e5 100644 --- a/csrc/python_frontend/python_bindings.cpp +++ b/csrc/python_frontend/python_bindings.cpp @@ -3570,7 +3570,7 @@ void initNvFuserPythonBindings(PyObject* module) { py::return_value_policy::reference); nvf_ops.def( - "embedding", + "embedding_fwd", [](FusionDefinition::Operators& self, Tensor input, Tensor weight, @@ -3579,7 +3579,7 @@ void initNvFuserPythonBindings(PyObject* module) { std::optional norm_type, std::optional scale_grad_by_freq, std::optional sparse) -> decltype(auto) { - FUSER_PERF_SCOPE("Operators.embedding"); + FUSER_PERF_SCOPE("Operators.embedding_fwd"); NVF_CHECK( self.validUse(), "Attempting to add to a completed definition!"); FusionDefinition* fd = self.fusion_definition; @@ -3602,7 +3602,7 @@ void initNvFuserPythonBindings(PyObject* module) { ? fd->recordingState(sparse.value()()) : State(/*_index=*/0, /*_stype=*/serde::StateType::None); - fd->defineRecord(new EmbeddingOpRecord( + fd->defineRecord(new EmbeddingFwdOpRecord( {fd->recordingState(input()), fd->recordingState(weight()), padding_idx_state, diff --git a/csrc/scheduler/expr_eval_sched.cpp b/csrc/scheduler/expr_eval_sched.cpp index 98f85af2788..affecd14fbd 100644 --- a/csrc/scheduler/expr_eval_sched.cpp +++ b/csrc/scheduler/expr_eval_sched.cpp @@ -29,7 +29,7 @@ bool ExprEvalScheduler::canScheduleCompileTime(Fusion* fusion) { return false; } - if (exprs.front()->isOneOf()) { + if (exprs.front()->isOneOf()) { return true; } diff --git a/csrc/serde/fusion_cache.fbs b/csrc/serde/fusion_cache.fbs index f8ec85dcaae..562f35d9e59 100644 --- a/csrc/serde/fusion_cache.fbs +++ b/csrc/serde/fusion_cache.fbs @@ -42,7 +42,7 @@ enum RecordType: int { CastTv, CastVal, CatOp, - EmbeddingOp, + EmbeddingFwdOp, End, ExpandOp, FullOp, diff --git a/csrc/serde/fusion_record.cpp b/csrc/serde/fusion_record.cpp index ded4056f264..ec585ca9ae4 100644 --- a/csrc/serde/fusion_record.cpp +++ b/csrc/serde/fusion_record.cpp @@ -322,11 +322,11 @@ void RecordFunctorFactory::registerAllParsers() { }; registerParser(RecordType::SdpaBwdOp, deserializeSdpaBwdRecord); - auto deserializeEmbeddingRecord = [&](const RecordFunctor* buffer) { - return new python_frontend::EmbeddingOpRecord( + auto deserializeEmbeddingFwdRecord = [&](const RecordFunctor* buffer) { + return new python_frontend::EmbeddingFwdOpRecord( parseStateArgs(buffer->args()), parseStateArgs(buffer->outputs())); }; - registerParser(RecordType::EmbeddingOp, deserializeEmbeddingRecord); + registerParser(RecordType::EmbeddingFwdOp, deserializeEmbeddingFwdRecord); // END OpRecord Parsers diff --git a/tests/cpp/test_embedding_node.cpp b/tests/cpp/test_embedding_node.cpp index fd77c606644..f98567c957f 100644 --- a/tests/cpp/test_embedding_node.cpp +++ b/tests/cpp/test_embedding_node.cpp @@ -20,7 +20,7 @@ using EmbeddingTest = NVFuserTest; constexpr int64_t n = 5, s = 2; -TEST_F(EmbeddingTest, EmbeddingNode) { +TEST_F(EmbeddingTest, EmbeddingFwdNode) { auto fusion = std::make_unique(); FusionGuard fg(fusion.get()); std::vector inp_shape({s}); @@ -32,7 +32,7 @@ TEST_F(EmbeddingTest, EmbeddingNode) { fusion->addInput(tv_inp); fusion->addInput(tv_weight); - auto tv_output = embedding(tv_inp, tv_weight, nullptr, nullptr, nullptr, nullptr, nullptr); + auto tv_output = embedding_fwd(tv_inp, tv_weight, nullptr, nullptr, nullptr, nullptr, nullptr); fusion->addOutput(tv_output); auto options = at::TensorOptions().device(at::kCUDA, 0); diff --git a/tests/python/test_embedding.py b/tests/python/test_embedding.py index 6fa264be983..0e2670f2cbb 100644 --- a/tests/python/test_embedding.py +++ b/tests/python/test_embedding.py @@ -46,7 +46,7 @@ def fusion_func( for idx in range(len(optional_inputs)): if has_optional_inputs[idx]: optional_inputs[idx] = fd.define_scalar(value=None, dtype=optional_inputs_dtypes[idx]) - out = fd.ops.embedding(input, weight, *optional_inputs) + out = fd.ops.embedding_fwd(input, weight, *optional_inputs) fd.add_output(out) N, S = 10, 3 @@ -73,5 +73,5 @@ def fusion_func( norm_type = 2.0 if norm_type is None else norm_type scale_grad_by_freq = False if scale_grad_by_freq is None else scale_grad_by_freq sparse = False if sparse is None else sparse - ref_out = F.embedding(input, weight, padding_idx, max_norm, norm_type, scale_grad_by_freq, sparse) + ref_out = F.embedding_fwd(input, weight, padding_idx, max_norm, norm_type, scale_grad_by_freq, sparse) torch.testing.assert_close(nvf_out[0], ref_out) \ No newline at end of file