pytorch/aten/src/ATen/DTensorState.h at main · gnosyslambda/pytorch · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
#pragma once

#include <c10/macros/Macros.h>

namespace at {

TORCH_API bool get_dtensor_allow_implicit_replication();
TORCH_API void set_dtensor_allow_implicit_replication(bool enabled);

struct DTensorAllowImplicitReplication {
  DTensorAllowImplicitReplication()
      : prev_dtensor_allow_implicit_replication_(
            get_dtensor_allow_implicit_replication()) {
    set_dtensor_allow_implicit_replication(true);
  }

  DTensorAllowImplicitReplication(const DTensorAllowImplicitReplication&) =
      delete;
  DTensorAllowImplicitReplication& operator=(
      const DTensorAllowImplicitReplication&) = delete;
  DTensorAllowImplicitReplication(DTensorAllowImplicitReplication&&) = delete;
  DTensorAllowImplicitReplication& operator=(
      DTensorAllowImplicitReplication&&) = delete;

  ~DTensorAllowImplicitReplication() {
    set_dtensor_allow_implicit_replication(
        prev_dtensor_allow_implicit_replication_);
  }

 private:
  bool prev_dtensor_allow_implicit_replication_;
};

} // namespace at