Fix negative position in InlinePropagator (csarofeen#1813)

shmsong · Jul 11, 2022 · de6b7ca · de6b7ca
1 parent 10a996c
commit de6b7ca
Show file tree

Hide file tree

Showing 3 changed files with 7 additions and 7 deletions.
diff --git a/torch/csrc/jit/codegen/cuda/compute_at.cpp b/torch/csrc/jit/codegen/cuda/compute_at.cpp
@@ -163,7 +163,7 @@ std::unordered_set<TensorView*> getPropagationSubgraph(
 void ComputeAt::runAt(
     TensorView* producer,
     TensorView* consumer,
-    unsigned int consumer_position,
+    int64_t consumer_position,
     ComputeAtMode mode) {
   FUSER_PERF_SCOPE("ComputeAt::runAt");
 
@@ -176,7 +176,7 @@ void ComputeAt::runAt(
       " are not in the same fusion.");
 
   if (mode == ComputeAtMode::MostInlined) {
-    consumer_position = consumer->nDims();
+    consumer_position = -1;
   }
 
   FusionGuard fg(producer->fusion());
@@ -205,7 +205,7 @@ void ComputeAt::runAt(
 void ComputeAt::runWith(
     TensorView* producer,
     TensorView* consumer,
-    unsigned int producer_position,
+    int64_t producer_position,
     ComputeAtMode mode) {
   FUSER_PERF_SCOPE("ComputeAt::runWith");
 
@@ -218,7 +218,7 @@ void ComputeAt::runWith(
       " are not in the same fusion.");
 
   if (mode == ComputeAtMode::MostInlined) {
-    producer_position = producer->nDims();
+    producer_position = -1;
   }
 
   FusionGuard fg(producer->fusion());

diff --git a/torch/csrc/jit/codegen/cuda/compute_at.h b/torch/csrc/jit/codegen/cuda/compute_at.h
@@ -27,15 +27,15 @@ struct ComputeAt {
   static void runAt(
       TensorView* producer,
       TensorView* consumer,
-      unsigned int consumer_position,
+      int64_t consumer_position,
       ComputeAtMode mode = ComputeAtMode::Standard);
 
   // Runs the compute with pass making consumer look like producer, computing
   // producer relative to consumer
   static void runWith(
       TensorView* producer,
       TensorView* consumer,
-      unsigned int producer_position,
+      int64_t producer_position,
       ComputeAtMode mode = ComputeAtMode::Standard);
 };
 

diff --git a/torch/csrc/jit/codegen/cuda/inline_propagator.cpp b/torch/csrc/jit/codegen/cuda/inline_propagator.cpp
@@ -178,7 +178,6 @@ InlinePropagator::InlinePropagator(
     : max_pos_calc(mode),
       selected_(std::move(selected)),
       reference_(reference),
-      reference_pos_(reference_pos),
       mode_(mode) {
   if (reference_pos < 0) {
     reference_pos += int64_t(reference->nDims()) + 1;
@@ -192,6 +191,7 @@ InlinePropagator::InlinePropagator(
       " and <= ",
       reference->nDims(),
       ".");
+  reference_pos_ = reference_pos;
 }
 
 void InlinePropagator::propagateC2P(TensorView* from, TensorView* to) {