minor review

NVIDIA · Dec 23, 2024 · 741202b · 741202b
1 parent ed4440a
commit 741202b
Show file tree

Hide file tree

Showing 4 changed files with 5 additions and 17 deletions.
diff --git a/csrc/host_ir/executor.cpp b/csrc/host_ir/executor.cpp
@@ -188,7 +188,8 @@ HostIrEvaluator::HostIrEvaluator(
     HostIrEvaluatorParams params)
     : container_(std::move(container)),
       communicator_(communicator),
-      params_(params) {
+      params_(params),
+      my_device_index_(communicator_ ? communicator_->deviceId() : 0) {
   const DeviceIdxType device_index =
       (communicator_ != nullptr && communicator_->is_available())
       ? communicator_->deviceId()
@@ -274,11 +275,10 @@ void HostIrEvaluator::handle(SetCurrentStream* set_current_stream) {
 }
 
 void HostIrEvaluator::handle(GetCurrentStream* get_current_stream) {
-  auto my_device_index = static_cast<c10::DeviceIndex>(
-      communicator_ ? communicator_->deviceId() : 0);
   streams_.insert(
       {get_current_stream->stream(),
-       c10::cuda::getCurrentCUDAStream(my_device_index)});
+       c10::cuda::getCurrentCUDAStream(
+           static_cast<c10::DeviceIndex>(my_device_index_))});
 }
 
 void HostIrEvaluator::handle(Synchronize* synchronize) {

diff --git a/csrc/host_ir/executor.h b/csrc/host_ir/executor.h
@@ -139,6 +139,7 @@ class HostIrEvaluator final : public OptOutDispatch {
   using StreamKey = std::variant<int64_t, Stream*>;
   std::unordered_map<StreamKey, c10::cuda::CUDAStream> streams_;
   std::unordered_map<Expr*, c10::intrusive_ptr<c10d::Work>> works_;
+  const int64_t my_device_index_;
 };
 
 } // namespace hir

diff --git a/csrc/host_ir/host_ir.cpp b/csrc/host_ir/host_ir.cpp
@@ -195,16 +195,6 @@ std::string GetCurrentStream::toString(int indent_size) const {
   return ss.str();
 }
 
-// TODO: implement better ?
-std::string GetCurrentStream::toInlineString(int indent_size) const {
-  NVF_CHECK(false, "Cannot be printed inline");
-}
-
-// TODO: implement
-bool GetCurrentStream::sameAs(const Statement* other) const {
-  return false;
-}
-
 Wait::Wait(IrBuilderPasskey passkey, Expr* expr)
     : Expr(passkey, {}, {}, {expr}) {
   NVF_ERROR(passkey.ir_container_ != nullptr);

diff --git a/csrc/host_ir/host_ir.h b/csrc/host_ir/host_ir.h
@@ -174,13 +174,10 @@ class GetCurrentStream : public Expr {
   NVFUSER_DECLARE_CLONE_AND_CREATE
 
   std::string toString(int indent_size = 0) const override;
-  std::string toInlineString(int indent_size = 0) const override;
   const char* getOpString() const override {
     return "hir::GetCurrentStream";
   }
 
-  bool sameAs(const Statement* other) const override;
-
   Stream* stream() const {
     return attributes_.at(0)->as<Stream>();
   }