From e671341eace36e142a729455ae3126f2d79b47de Mon Sep 17 00:00:00 2001
From: Naoya Maruyama <nmaruyama@nvidia.com>
Date: Wed, 19 Apr 2023 13:34:56 -0700
Subject: [PATCH 1/2] Fix alias analysis

Set was replaced with LoadStoreOp, which was not recognized as UnaryOp,
and thus the alias analysis failed to detect safe aliasing
---
 csrc/ir_utils.cpp           |  5 +++++
 csrc/ir_utils.h             |  3 +++
 csrc/lower_alias_memory.cpp | 13 ++-----------
 3 files changed, 10 insertions(+), 11 deletions(-)

diff --git a/csrc/ir_utils.cpp b/csrc/ir_utils.cpp
index cba2c1557b8..7ba001d6578 100644
--- a/csrc/ir_utils.cpp
+++ b/csrc/ir_utils.cpp
@@ -572,6 +572,11 @@ bool isReductionTvOp(const Expr* expr) {
   return ir_utils::isTvOp(expr) && isReductionOp(expr);
 }
 
+bool isPointwiseTvOp(const Expr* expr) {
+  return isTvOp(expr) &&
+      expr->isOneOf<UnaryOp, BinaryOp, TernaryOp, LoadStoreOp>();
+}
+
 std::vector<ViewOp*> getViewOps(Fusion* fusion) {
   auto all_exprs = fusion->exprs();
 
diff --git a/csrc/ir_utils.h b/csrc/ir_utils.h
index b7f9b53eac1..74bb52803ad 100644
--- 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
index 9523d6e3ab9..58036e46dbb 100644
--- 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)) {

From 5b188cc44837aaef776c75148541e8576e9b0501 Mon Sep 17 00:00:00 2001
From: Naoya Maruyama <nmaruyama@nvidia.com>
Date: Wed, 19 Apr 2023 15:44:57 -0700
Subject: [PATCH 2/2] Do not consider transpose as pointwise

---
 csrc/ir_utils.cpp | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/csrc/ir_utils.cpp b/csrc/ir_utils.cpp
index 7ba001d6578..206e5f67192 100644
--- a/csrc/ir_utils.cpp
+++ b/csrc/ir_utils.cpp
@@ -573,8 +573,12 @@ bool isReductionTvOp(const Expr* 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, LoadStoreOp>();
+      (expr->isOneOf<UnaryOp, BinaryOp, TernaryOp>() ||
+       (expr->isA<LoadStoreOp>() &&
+        !ir_utils::getTvOutput(expr)->hasRFactor()));
 }
 
 std::vector<ViewOp*> getViewOps(Fusion* fusion) {