1 #pragma once 2 #include <c10/core/ScalarType.h> 3 #include <torch/csrc/Export.h> 4 #include <vector> 5 6 #include <torch/csrc/jit/tensorexpr/ir_mutator.h> 7 8 namespace torch { 9 namespace jit { 10 namespace tensorexpr { 11 12 class TORCH_API IRCloner : public IRMutator { 13 public: 14 ~IRCloner() override = default; 15 ExprPtr mutate(const AddPtr& v) override; 16 ExprPtr mutate(const SubPtr& v) override; 17 ExprPtr mutate(const MulPtr& v) override; 18 ExprPtr mutate(const DivPtr& v) override; 19 ExprPtr mutate(const ModPtr& v) override; 20 ExprPtr mutate(const MaxPtr& v) override; 21 ExprPtr mutate(const MinPtr& v) override; 22 ExprPtr mutate(const AndPtr& v) override; 23 ExprPtr mutate(const OrPtr& v) override; 24 ExprPtr mutate(const XorPtr& v) override; 25 ExprPtr mutate(const LshiftPtr& v) override; 26 ExprPtr mutate(const RshiftPtr& v) override; 27 ExprPtr mutate(const CompareSelectPtr& v) override; 28 #define IMM_MUTATE_DECLARE(Type, Name) \ 29 ExprPtr mutate(const Name##ImmPtr& v) override; 30 AT_FORALL_SCALAR_TYPES_AND3(Bool, Half, BFloat16, IMM_MUTATE_DECLARE); 31 #undef IMM_MUTATE_DECLARE 32 ExprPtr mutate(const CastPtr& v) override; 33 ExprPtr mutate(const BitCastPtr& v) override; 34 ExprPtr mutate(const VarPtr& v) override; 35 ExprPtr mutate(const BufPtr& v) override; 36 ExprPtr mutate(const RampPtr& v) override; 37 ExprPtr mutate(const LoadPtr& v) override; 38 ExprPtr mutate(const BroadcastPtr& v) override; 39 ExprPtr mutate(const IfThenElsePtr& v) override; 40 ExprPtr mutate(const IntrinsicsPtr& v) override; 41 42 ExprPtr mutate(const TermPtr& v) override; 43 ExprPtr mutate(const PolynomialPtr& v) override; 44 ExprPtr mutate(const RoundOffPtr& v) override; 45 ExprPtr mutate(const MaxTermPtr& v) override; 46 ExprPtr mutate(const MinTermPtr& v) override; 47 48 ExprPtr mutate(const ReduceOpPtr& v) override; 49 50 StmtPtr mutate(const ForPtr& v) override; 51 StmtPtr mutate(const BlockPtr& v) override; 52 StmtPtr mutate(const StorePtr& v) override; 53 StmtPtr mutate(const AtomicAddPtr& v) override; 54 StmtPtr mutate(const SyncThreadsPtr& v) override; 55 StmtPtr mutate(const ExternalCallPtr& v) override; 56 StmtPtr mutate(const ExternalCallWithAllocPtr& v) override; 57 58 StmtPtr mutate(const AllocatePtr& v) override; 59 StmtPtr mutate(const FreePtr& v) override; 60 StmtPtr mutate(const LetPtr& v) override; 61 StmtPtr mutate(const CondPtr& v) override; 62 }; 63 64 } // namespace tensorexpr 65 } // namespace jit 66 } // namespace torch 67