Fix alias analysis (#185)

NVIDIA · Apr 20, 2023 · d68da6b · d68da6b
1 parent 255af2b
commit d68da6b
Show file tree

Hide file tree

Showing 3 changed files with 14 additions and 11 deletions.
diff --git a/csrc/ir_utils.cpp b/csrc/ir_utils.cpp
@@ -572,6 +572,15 @@ bool isReductionTvOp(const Expr* expr) {
   return ir_utils::isTvOp(expr) && isReductionOp(expr);
 }
 
+bool isPointwiseTvOp(const Expr* expr) {
+  // LoadStoreOp with rfactor domain means transpose, which is not
+  // considered pointwise
+  return isTvOp(expr) &&
+      (expr->isOneOf<UnaryOp, BinaryOp, TernaryOp>() ||
+       (expr->isA<LoadStoreOp>() &&
+        !ir_utils::getTvOutput(expr)->hasRFactor()));
+}
+
 std::vector<ViewOp*> getViewOps(Fusion* fusion) {
   auto all_exprs = fusion->exprs();
 

diff --git a/csrc/ir_utils.h b/csrc/ir_utils.h
@@ -338,6 +338,9 @@ TORCH_CUDA_CU_API bool isReductionOp(const Expr*);
 // Returns if Expr is a reduction op with TensorView or TensorIndex
 TORCH_CUDA_CU_API bool isReductionTvOp(const Expr*);
 
+// Returns if Expr is a pointwise op op with TensorView or TensorIndex
+TORCH_CUDA_CU_API bool isPointwiseTvOp(const Expr* expr);
+
 // Returns all non-trivial view operations. We shouldn't have trivial view
 // operations but this function is to simply make sure if we ever do we don't
 // pull them in.

diff --git a/csrc/lower_alias_memory.cpp b/csrc/lower_alias_memory.cpp
@@ -1225,7 +1225,8 @@ class ReusableAllocationFinder : private kir::IrVisitor {
         if (!tv_def) {
           continue;
         }
-        if (!isPointwiseTvOp(tv_def) && !ir_utils::isReductionTvOp(tv_def)) {
+        if (!ir_utils::isPointwiseTvOp(tv_def) &&
+            !ir_utils::isReductionTvOp(tv_def)) {
           if (isBroadcastTvOp(tv_def)) {
             info.has_broadcast_between = true;
           } else {
@@ -1266,16 +1267,6 @@ class ReusableAllocationFinder : private kir::IrVisitor {
     }
   }
 
-  // Do we have a true pointwise op?
-  // (ie. a TV op, excluding direct assignments and reductions)
-  bool isPointwiseTvOp(const Expr* expr) {
-    if (ir_utils::isTvOp(expr)) {
-      return expr->isA<UnaryOp>() || expr->isA<BinaryOp>() ||
-          expr->isA<TernaryOp>();
-    }
-    return false;
-  }
-
   // Utility to capture broadcast ops
   bool isBroadcastTvOp(const Expr* expr) {
     if (!ir_utils::isTvOp(expr)) {