Repository: NVIDIA/Fuser
Pull request #3968: Remove trivial G->G broadcasts at lowering
Author: login: jacobhinkle
Created: 2025-02-26T00:33:58Z | Merged: 2025-03-05T14:11:05Z

In this PR I am introducing new behavior in lowering which skips `BroadcastOp`, `SqueezeOp`, or `LoadStoreOp` when
1. both the inputs and outputs lie in global memory
2. both the inputs and outputs are contiguous
3. the output is not a Fusion output
4. their allocation domains "match": i.e. when all non-broadcast non-reduction IDs are exact mapped and in the same stride order.

The intention behind these conditions is to isolate ops for which it is safe to avoid moving data into a temporary global buffer when no computation is performed and all elements are in the same relative linear positions in that intermediate buffer. Instead, we will use the producer buffer as an alias to avoid a needless G->G copy. **Importantly, this allows us to schedule G->G operations in order to _view_ Fusion inputs as broadcast/squeezed/permuted logically** without needing to codegen new operations and incur avoidable data movements.

When these operations are skipped, no `Allocate` nodes are created meaning we do not actually allocate the intermediate global tensor, and we do not generate any code for initializing or copying data into those global tensors. We generate a linear index for the _unaliased_ tensor whenever it is used, but then we replace the `view` of the `TensorIndex` to point to the alias target.

This PR is motivated by #1628 and #3671. It is an alternative to #3366 and #3372, and would allow us to roll back #3391, #3414, #3416. In a follow-up, I will create an alternative to #3440 which translates `MatmulPattern` on Hopper, which I believe will let us fix #3671 and enable persistent matmul with `MatmulOp` and `LinearOp`.

Linked issue #3671: Inlining error in Hopper matmul with AxisMapping and grid swizzling
The inlining logic for `MmaOp` with `AxisMapping` checks that unmapped dimensions are `Broadcast`. We expect to have something like this
```
t0 {
  logical: [ iS0{M} iS1{K} ]
  loop: [ iS2{ ceilDiv(M, 256) }  bS5{1}  iS3{256} bS6{256} ... ]
    Split iS0 by 256 -> iS2, iS3
    Split bS4 by 256 -> bS5, bS6
      ...
  additional ids: bS4{1}
}
t1 {
  logical: [ iS10{N} iS11{K} ]
  loop: [ bS13{1} iS15{ ceilDiv(N, 256) } bS14{256} iS16{256} ... ]
    Split bS12 by 256 -> bS13, bS14
    Split iS10 by 256 -> iS15, iS16
      ...
  additional ids: bS12{1}
}
```
In this case, we are able to inline the mma operation that consumes these two tensors, but we check that the unmapped IDs 5, 6, 13, and 14 are `Broadcast` and that the operation is an `MmaOp`.

In the case of grid swizzling by a factor 4, we will do some further scheduling here. For example we will have
```
t0 {
  logical: [ iS0{M} iS1{K} ]
  loop: [ iS21{ ceilDiv( ceilDiv(M, 256), 4) } iS22{4} iS3{256} bS6{256} ... ]
    Split iS0 by 256 -> iS2, iS3
    Split bS4 by 256 -> bS5, bS6
    Split iS2 by 4 -> iS20, iS21{4}
    Merge bS5 with iS21{4} -> iS22{4}
      ...
  additional ids: bS4{1}
}
```
Now we have mixed the first two outer dimensions with this swizzle and what used to be a simple split of a loop broadcast (`bS5`) is now an iteration ID `iS22{4}` resulting from the merge.

I am not sure yet how to address this. I don't think we can just inline here without some other changes since when I disable this check I get errors in expression sorting.
Linked issue comments: user: login: jacobhinkle
created_at: 2025-02-03T19:01:42Z
body: This error also appears whenever we try and do persistent kernel scheduling when there is a translated `MatmulOp` or `LinearOp` node. In that case, commenting out the `isBroadcast` check leads to an error in SyncMap
> INTERNAL ASSERT FAILED at "/opt/pytorch/nvfuser/csrc/device_lower/analysis/sync_information.cpp":798, please report a bug with repro script to NVFuser at https://github.com/NVIDIA/Fuser/issues. Inconsistent parallelization found between TV3 (T3_s___bfloat[iS27{( ceilDiv(( ceilDiv(( (( (( getMetaData(T0) )).logical_size ))[0] ), 128) ), 132) )}, iblockIdx.x28{132}, iS20{( ceilDiv(( (( (( getMetaData(T0) )).logical_size ))[1] ), 16) )}, bS23{256}, iS29{1}, iB31{16}, iB37{2}, iB34{4}, iB38{2}, iB36{8}] ca_pos( 3 )) and TV6(T6_l_float[iS65{( ceilDiv(( ( ceilDiv(( (( (( getMetaData(T1) )).logical_size ))[0] ), 256) ) * ( ceilDiv(( (( (( getMetaData(T0) )).logical_size ))[0] ), 128) ) ), 132) )}, iblockIdx.x66{132}, rS58{( ceilDiv(( (( (( getMetaData(T0) )).logical_size ))[1] ), 16) )}, ithreadIdx.y77{2}, iS69{1}, iS73{1}, rS75{1}, iMMA70{64}, iMMA74{256}, rMMA76{16}] ca_pos( 2 ) produce_pos( 3 )). Producer is required to be in Global Memory based on parallelization strategy. RAW flags: (blockIdx.x threadIdx.y) 

 I think the problem is in the `AxisMapping` approach itself; it feels brittle to me. We have already had to carve out special cases for MmaOp in inlining and predicate elimination, and now we see that that is not enough since the missing axes are sometimes merged with `Iteration` domains. The swizzle example above shows this issue. Ideally this case would behave just like the `fusedMultiplySum` case in which we have `Broadcast` axes in the fusion input. In that case, when we merge e.g. a `Broadcast` N axis with an `Iteration` M axis, we get an `Iteration` ID that permissive maps to the `Iteration` N/`Broadcast` M axis in the other operand and to the corresponding ID in the result.

I am considering now whether #3372 was a mistake and if we should revisit something like #3366 instead. What do you think @naoyam ?
author_association: COLLABORATOR
user: login: naoyam
created_at: 2025-02-03T22:48:42Z
body: So, the issue is we can't inline `iS22`. What are we trying to inline it to?
author_association: COLLABORATOR
user: login: jacobhinkle
created_at: 2025-02-04T15:29:02Z
body: Here is a diagram of grid swizzling for matmul:

![Image](https://github.com/user-attachments/assets/4f58f385-5519-4470-a6e7-d9315fa37235)

Here we need to inline `iS22`, `iS13`, and `iS38`, i.e. the outermost dimension in a scheduled `MmaOp`. The broadcast dimensions `bS4` and `bS2` (red) are `additional_ids_` in this case. If they were logical IDs, then the permissive map would map all these IDs to one another. 

Inlining is one thing: we can pretty much repeat the logic of permissive map to create an inlining check that allows this. However, when I bypass that check currently, I hit an error in `SyncMap` (inconsistent parallelization). That error is avoided when the IDs are Permissive mapped. Creating another carveout for `MmaOp` there should be possible just like for inlining and predicate elimination.

Another option (besides #3366) would be to try and make the additional IDs broadcast mapped with the corresponding consumer dimension, which would make these IDs all permissive mapped. For example we could add that as an option to `TensorView::broadcast` like `int64_t map_to_consumer_root_pos`. Inside pairwise logical domain map, we would check for that flag, check that the corresponding position in the consumer is not mapped to any producer ID already, and broadcast map it to the new additional_id.
author_association: COLLABORATOR
user: login: naoyam
created_at: 2025-02-04T16:52:21Z
body: Thanks. I think I understand what's happening and what the problem is.

I wonder if we could just get rid of the offending BroadcastOp during lowering. #3366 may work. I'd also try `scheduleLoopDomainsLike`, but it seems the easiest and quickest way is to keep the BroadcastOps there in the Fusion IR and just remove them once they are no longer necessary. Would that work?

author_association: COLLABORATOR
user: login: jacobhinkle
created_at: 2025-02-04T17:03:55Z
body: > I wonder if we could just get rid of the offending BroadcastOp during lowering. [#3366](https://github.com/NVIDIA/Fuser/issues/3366) may work. I'd also try `scheduleLoopDomainsLike`, but it seems the easiest and quickest way is to keep the BroadcastOps there in the Fusion IR and just remove them once they are no longer necessary. Would that work?

If we left the BroadcastOps in and scheduled it similar to how we do on Ampere, then we'd have something like this:
```
tv0_g[ i0, i2 ]
tv1_g[ i1, i2 ]
tv2_s[ i0, i2 ] = set(tv0_g)
tv3_s[ i1, i2 ] = set(tv1_g)
tv4_l[ i0, 1, i2 ] = broadcast(tv2_s)
tv5_l[ 1, i1, i2 ] = broadcast(tv3_s)
tv6_l[ i0, i1, r{i2} ] = mma(tv4_l, tv5_l)
```
If I understand correctly, you're saying we could potentially remove those broadcasts at lowering (say in the indexing pass) and replace `tv4_l` and `tv5_l` with `tv2_s` and `tv3_s` in the mma instruction there. Yes that might work! I will give it a shot.

author_association: COLLABORATOR
user: login: naoyam
created_at: 2025-02-04T17:08:28Z
body: Some of the localized KIR changes are also done after indexing, e.g., https://github.com/NVIDIA/Fuser/blob/main/csrc/device_lower/lower2device.cpp#L283

Removing the ops in KIR may feel ad-hoc, but I'd suspect #3366 would be a significant change.
author_association: COLLABORATOR
user: login: jacobhinkle
created_at: 2025-03-05T17:30:36Z
body: This issue is still open pending #3986
author_association: COLLABORATOR
user: login: jacobhinkle
created_at: 2025-04-15T16:32:40Z
body: Fixed by #4072 
author_association: COLLABORATOR

Diff:
diff --git a/CMakeLists.txt b/CMakeLists.txt
index 8157e6fdcb9..00a7b3e9155 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -99,6 +99,7 @@ list(APPEND NVFUSER_SRCS
   ${NVFUSER_SRCS_DIR}/device_lower/analysis/predicate_elimination.cpp
   ${NVFUSER_SRCS_DIR}/device_lower/analysis/sync_information.cpp
   ${NVFUSER_SRCS_DIR}/device_lower/analysis/tensor_memory.cpp
+  ${NVFUSER_SRCS_DIR}/device_lower/analysis/tensor_producer_aliases.cpp
   ${NVFUSER_SRCS_DIR}/device_lower/analysis/thread_predicate.cpp
   ${NVFUSER_SRCS_DIR}/device_lower/analysis/tma.cpp
   ${NVFUSER_SRCS_DIR}/device_lower/analysis/trivial_broadcast.cpp
@@ -595,6 +596,7 @@ list(APPEND JIT_TEST_SRCS
   ${NVFUSER_ROOT}/tests/cpp/test_predicate_elimination.cpp
   ${NVFUSER_ROOT}/tests/cpp/test_preseg_passes.cpp
   ${NVFUSER_ROOT}/tests/cpp/test_remove_bcast_squeeze.cpp
+  ${NVFUSER_ROOT}/tests/cpp/test_remove_trivial_ops.cpp
   ${NVFUSER_ROOT}/tests/cpp/test_replay.cpp
   ${NVFUSER_ROOT}/tests/cpp/test_resharding.cpp
   ${NVFUSER_ROOT}/tests/cpp/test_resize.cpp
diff --git a/csrc/device_lower/analysis/tensor_producer_aliases.cpp b/csrc/device_lower/analysis/tensor_producer_aliases.cpp
new file mode 100644
index 00000000000..9df3a74f54a
--- /dev/null
+++ b/csrc/device_lower/analysis/tensor_producer_aliases.cpp
@@ -0,0 +1,145 @@
+// clang-format off
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2023-present NVIDIA CORPORATION & AFFILIATES.
+ * All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ */
+// clang-format on
+#include <device_lower/analysis/tensor_producer_aliases.h>
+#include <device_lower/lower2device.h>
+#include <device_lower/utils.h>
+#include <exceptions.h>
+#include <ir/all_nodes.h>
+#include <ir/utils.h>
+#include <kernel_ir_dispatch.h>
+#include <type.h>
+#include <val_graph.h>
+
+#include <unordered_set>
+#include <vector>
+
+namespace nvfuser {
+
+namespace {
+
+bool isTrivialExpr(Expr* expr) {
+  TensorView* in = ir_utils::getTvInput(expr);
+  TensorView* out = ir_utils::getTvOutput(expr);
+  if (in == nullptr || out == nullptr ||
+      in->getMemoryType() != MemoryType::Global ||
+      out->getMemoryType() != MemoryType::Global || out->isFusionOutput() ||
+      !expr->isOneOf<BroadcastOp, LoadStoreOp, SqueezeOp>()) {
+    return false;
+  }
+  // This is a tensor op that does no computation. However, it may still be
+  // non-trivial if the allocation domains are not equivalent, as that would
+  // indicate data movement. We check that any non-Broadcast and non-Reduction
+  // domains in the two allocation domains are exact mapped and appear in the
+  // same order. We also check that the tensors are fully contiguous so that we
+  // can re-use the exact same linear index.
+  // TODO: support discontiguous inputs. The output should be an intermediate
+  // tensor which we would always assume to be contiguous, but the input might
+  // not be. This would necessitate using a different linear index.
+  const std::vector<IterDomain*>& in_alloc = TensorDomain::noReductions(
+      TensorDomain::noBroadcasts(in->getMaybeAllocationDomain()));
+  const std::vector<IterDomain*>& out_alloc =
+      TensorDomain::noBroadcasts(out->getMaybeAllocationDomain());
+
+  if (in_alloc.size() != out_alloc.size()) {
+    // Non-trivial allocation domains cannot be in bijective correspondence if
+    // there are different numbers of them
+    return false;
+  }
+
+  std::vector<bool> in_contig;
+  std::vector<bool> out_contig;
+  for (const std::optional<bool>& c : in->getContiguity()) {
+    if (c.has_value()) {
+      in_contig.push_back(c.value());
+    }
+  }
+  for (const std::optional<bool>& c : out->getContiguity()) {
+    if (c.has_value()) {
+      out_contig.push_back(c.value());
+    }
+  }
+
+  const ValGraph& exact_graph =
+      GpuLower::current()->idModel().idGraph(IdMappingMode::EXACT);
+
+  for (size_t pos : c10::irange(in_alloc.size())) {
+    // At this point in_pos and out_pos are both in range and point to
+    // non-broadcast IDs
+    IterDomain* in_id = in_alloc.at(pos);
+    IterDomain* out_id = out_alloc.at(pos);
+
+    // If this allocation ID is parallelized such that its loop index is not
+    // used, then we can ignore it for this analysis.
+    const auto id_is_indexed = [](TensorView* tv, IterDomain* id) {
+      // isMemoryPartitionedAcross is true if the dimension is not allocated
+      return !ir_utils::isMemoryPartitionedAcross(
+          tv->getMemoryType(), id->getParallelType());
+    };
+    if (!id_is_indexed(in, in_id) || !id_is_indexed(out, out_id)) {
+      continue;
+    }
+
+    NVF_ERROR(
+        out_contig.at(pos),
+        "Found discontiguous intermediate global tensor ",
+        out->toString());
+    if (in_contig.at(pos) != out_contig.at(pos)) {
+      return false;
+    }
+
+    if (exact_graph.toGroup(in_id) != exact_graph.toGroup(out_id)) {
+      // Unmapped pair
+      return false;
+    }
+  }
+
+  return true;
+}
+
+} // namespace
+
+void findTensorProducerAliases(Fusion* fusion) {
+  // First, find any inputs that have been aliased as an output. Since we don't
+  // currently guarantee that all reads of the aliased global intermediate
+  // tensors are performed before the aliased output is written, we must exclude
+  // aliases to these inputs.
+  std::unordered_set<TensorView*> inputs_aliased_by_outputs;
+  for (Val* v : fusion->outputs()) {
+    if (auto* io_alias_target =
+            dynamic_cast<TensorView*>(fusion->getOutputAlias(v).aliased_io)) {
+      inputs_aliased_by_outputs.insert(io_alias_target);
+    }
+  }
+  for (Expr* expr : fusion->exprs()) {
+    TensorView* in = ir_utils::getTvInput(expr);
+    if (in != nullptr && inputs_aliased_by_outputs.count(in) == 0 &&
+        isTrivialExpr(expr)) {
+      GpuLower::current()->aliasTensorProducer(ir_utils::getTvOutput(expr), in);
+    }
+  }
+}
+
+std::vector<Expr*> removeTensorProducerAliases(
+    const std::vector<Expr*>& exprs) {
+  // There are no ForLoops or IfThenElse exprs yet since this pass runs before
+  // LoopNestGenerator
+  std::vector<Expr*> filtered_exprs;
+  for (Expr* expr : exprs) {
+    // Here we check whether this expression's output is a TensorView aliased
+    // to one of its producer tensors. If so, that indicates that this
+    // expression is trivial and should not be considered for codegen.
+    if (TensorView* tv_out = ir_utils::getTvOutput(expr); tv_out &&
+        GpuLower::current()->getTensorProducerAlias(tv_out) != nullptr) {
+      continue;
+    }
+    filtered_exprs.push_back(expr);
+  }
+  return filtered_exprs;
+}
+
+} // namespace nvfuser
diff --git a/csrc/device_lower/analysis/tensor_producer_aliases.h b/csrc/device_lower/analysis/tensor_producer_aliases.h
new file mode 100644
index 00000000000..99faa9ec3ca
--- /dev/null
+++ b/csrc/device_lower/analysis/tensor_producer_aliases.h
@@ -0,0 +1,57 @@
+// clang-format off
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2023-present NVIDIA CORPORATION & AFFILIATES.
+ * All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ */
+// clang-format on
+#pragma once
+
+#include <fusion.h>
+#include <ir/base_nodes.h>
+
+#include <vector>
+
+namespace nvfuser {
+
+//! This pass finds trivial computations that can be avoided.
+//!
+//! For example, suppose we have the following fusion:
+//!
+//!   T0_g[ iS0{i1} ]
+//!   T1_g[ iS1{i1}, bS2{1} ] = broadcast(T0_g)
+//!   T2_l[ iS3{i1}, bS4{1} ] = set(T1_g)
+//!
+//! In a case like this, we would ordinarily allocate an intermediate global
+//! tensor to hold T1, and it would be realized. However, under some conditions
+//! we can avoid this allocation and copying by re-using T0_g in any uses of
+//! T1_g. In the example above, it is clear that in the definition of T2_l, we
+//! could simply find the appropriate element of T0_g and load that instead of
+//! loading from T1_g. When we do this we call it a tensor producer alias to
+//! distinguish it from other uses of the term "alias" in nvfuser such as smem
+//! buffer aliasing and fusion input/output buffer reuse. Note that it is not
+//! necessarily the case that a tensor producer alias points to a _direct_
+//! producer of that tensor, since chains of trivial ops may be skipped by this
+//! pass.
+//!
+//! Currently, we only allow aliasing across Exprs that satisfy these
+//! conditions:
+//!   1) The expr is a LoadStoreOp, BroadcastOp, or SqueezeOp
+//!   2) It has input and output whose allocation IterDomains are exact mapped,
+//!      ignoring Broadcast and Reduction axes.
+//!   3) It has both input and output that are fully contiguous.
+//! The first condition guarantees that we do not skip any important
+//! computation, while the second and third condition allow us to re-use the
+//! same linear index that we would have used without the alias but merely
+//! changing the base address of the buffer. In the future, we may allow
+//! relaxing the third condition, in which case the linear index would be
+//! modified by replacing the intermediate buffer's strides with those of the
+//! producer.
+void findTensorProducerAliases(Fusion* fusion);
+
+//! This is a lowering pass that actually removes top-level exprs whose outputs
+//! have been marked as tensor producer aliases in the findTensorProducerAliases
+//! analysis pass.
+std::vector<Expr*> removeTensorProducerAliases(const std::vector<Expr*>& exprs);
+
+} // namespace nvfuser
diff --git a/csrc/device_lower/lower2device.cpp b/csrc/device_lower/lower2device.cpp
index 7a2faebb717..79ae28c91a0 100644
--- a/csrc/device_lower/lower2device.cpp
+++ b/csrc/device_lower/lower2device.cpp
@@ -11,6 +11,7 @@
 #include <debug.h>
 #include <device_lower/analysis/device_version.h>
 #include <device_lower/analysis/divisible_split.h>
+#include <device_lower/analysis/tensor_producer_aliases.h>
 #include <device_lower/pass/alias_memory.h>
 #include <device_lower/pass/allocation.h>
 #include <device_lower/pass/circular_buffer.h>
@@ -265,7 +266,8 @@ GpuLower::GpuLower(Fusion* fusion, const CompileParams& cparams)
           // Each pass is a pair of (name, function), where the name will be
           // printed in verbose mode of lowering. The function must take a
           // const std::vector<Expr*>& and return a std::vector<Expr*>.
-          {{"LoopNestGenerator", LoopNestGenerator::loweredExprs},
+          {{"removeTensorProducerAliases", removeTensorProducerAliases},
+           {"LoopNestGenerator", LoopNestGenerator::loweredExprs},
            {"loadStoreOpInserter", loadStoreOpInserter},
            {"insertGridSerializationSyncs", insertGridSerializationSyncs},
            {"insertAllocations", insertAllocations},
@@ -567,6 +569,11 @@ void GpuLower::analysis(Fusion* fusion) {
   validateLookupTV(fusion_);
   dumpExprsIfEnabled(fusion_->exprs(), "validateLookupTV");
 
+  // Find trivial global to global broadcast, squeeze, and set operations and
+  // mark their outputs as aliases of their inputs.
+  findTensorProducerAliases(fusion_);
+  dumpExprsIfEnabled(fusion_->exprs(), "findTensorProducerAliases");
+
   // Depends on thread_pred_map_, validates parallelization collects which
   // tensor views need WAR or RAW syncs
   sync_map_ = std::make_shared<const SyncMap>(fusion_);
@@ -665,4 +672,28 @@ Val* GpuLower::getLoopIndexVariable(
   }
 }
 
+void GpuLower::aliasTensorProducer(TensorView* consumer, TensorView* producer) {
+  if (TensorView* producer_alias = getTensorProducerAlias(producer)) {
+    // Chase reference. If producer itself is an alias, then get the tensor it
+    // is aliased to.
+    NVF_ERROR(
+        getTensorProducerAlias(producer_alias) == nullptr,
+        "Found unsimplified alias from ",
+        producer->toString(),
+        " to ",
+        producer_alias->toString(),
+        " which is then aliased to ",
+        getTensorProducerAlias(producer_alias)->toString());
+    producer = producer_alias;
+  }
+  tensor_producer_alias_map_[consumer] = producer;
+  for (auto& [c, p] : tensor_producer_alias_map_) {
+    // If anything was previously aliased _to_ consumer, update those links to
+    // point to producer
+    if (p == consumer) {
+      p = producer;
+    }
+  }
+}
+
 } // namespace nvfuser
diff --git a/csrc/device_lower/lower2device.h b/csrc/device_lower/lower2device.h
index 4d101824fad..6e148d8de13 100644
--- a/csrc/device_lower/lower2device.h
+++ b/csrc/device_lower/lower2device.h
@@ -315,6 +315,46 @@ class GpuLower : public NonCopyable {
     return id_model_options_;
   }
 
+  //! Define an alias for consumer as producer.
+  //!
+  //! If producer is already aliased, we chase the alias. If there are tensors
+  //! aliased to consumer, their aliases are updated to point to the new
+  //! producer. This guarantees that any aliases are to producers that get
+  //! codegened.
+  //!
+  //! If there is a chain of trivial ops that should be skipped, then all of the
+  //! intermediate tensors should be aliased to the common producer:
+  //!
+  //!   b = broadcast(a)
+  //!   c = permute(b)
+  //!   d = squeeze(c)
+  //!
+  //! In this example, if all four of a, b, c, and d share the same memory type
+  //! and would use the same index, then we don't need any of these three
+  //! expressions and can simply replace the TensorIndex for d with that for a
+  //! in codegen'd expressions. So we should set up the following aliases:
+  //!
+  //!   d -> a
+  //!   c -> a
+  //!   b -> a
+  //!
+  //! Omitting one of these aliases might cause errors since that tensor's
+  //! definition might get codegen'd without an allocation.
+  void aliasTensorProducer(TensorView* consumer, TensorView* producer);
+
+  //! Return producer that this tensor should be aliased to. Returns nullptr if
+  //! no alias exists, i.e. that we should codegen tv's definition.
+  TensorView* getTensorProducerAlias(TensorView* tv) const {
+    auto it = tensor_producer_alias_map_.find(tv);
+    return it != tensor_producer_alias_map_.end() ? it->second : nullptr;
+  }
+
+  //! Return producer alias for tv or tv itself if it is unaliased
+  TensorView* getMaybeTensorProducerAlias(TensorView* tv) const {
+    TensorView* alias_tv = getTensorProducerAlias(tv);
+    return alias_tv == nullptr ? tv : alias_tv;
+  }
+
  private:
   void analysis(Fusion* fusion);
 
@@ -377,6 +417,12 @@ class GpuLower : public NonCopyable {
   // Information about tensor memory usage
   TensorMemoryInfo tmem_info_;
 
+  // Track TensorViews that will be aliased to their producers because of
+  // trivial ops and scheduling such that the same index is used. Note that the
+  // alias does not need to be a direct producer in case there is a chain of
+  // trivial ops like permute->bcast->set.
+  std::unordered_map<TensorView*, TensorView*> tensor_producer_alias_map_;
+
   // Keep track of validations needed at runtime. For example, a pair of
   //! "extent mod split_factor == 0" and an error message for divisibility check
   //! for vectorization.
diff --git a/csrc/device_lower/pass/index.cpp b/csrc/device_lower/pass/index.cpp
index e3059aecfcc..583e7c29fde 100644
--- a/csrc/device_lower/pass/index.cpp
+++ b/csrc/device_lower/pass/index.cpp
@@ -13,6 +13,7 @@
 #include <index_compute.h>
 #include <ir/iostream.h>
 #include <ir/utils.h>
+#include <kernel_ir.h>
 #include <ops/arith.h>
 #include <options.h>
 #include <predicate_compute.h>
@@ -45,7 +46,7 @@ Val* IndexLowering::lowerSrcIndex(
     DataType as_type) const {
   if (auto tv = dynamic_cast<TensorView*>(src)) {
     NVF_ERROR(dst->isA<TensorView>());
-    return Index::getProducerIndex(
+    kir::TensorIndex* tind = Index::getProducerIndex(
         tv,
         dst->as<TensorView>(),
         for_loops_,
@@ -53,6 +54,13 @@ Val* IndexLowering::lowerSrcIndex(
         override_index,
         generate_pointer,
         as_type);
+    if (TensorView* aliased_producer =
+            GpuLower::current()->getTensorProducerAlias(tv)) {
+      return IrBuilder::create<kir::TensorIndex>(
+          aliased_producer, tind->index());
+    } else {
+      return tind;
+    }
   } else {
     return src;
   }
diff --git a/csrc/device_lower/pass/insert_syncs.cpp b/csrc/device_lower/pass/insert_syncs.cpp
index 8121ae0c851..45d9d0ddf19 100644
--- a/csrc/device_lower/pass/insert_syncs.cpp
+++ b/csrc/device_lower/pass/insert_syncs.cpp
@@ -433,6 +433,9 @@ class ReadAfterWriteSyncs : public kir::ExprMutator {
     // TODO: unify the handle of cp.async
     std::unordered_map<AsyncOpType, std::unordered_set<Expr*>> input_async_ops;
     for (auto inp : expr->inputs()) {
+      if (auto* inp_tv = dynamic_cast<TensorView*>(inp)) {
+        inp = GpuLower::current()->getMaybeTensorProducerAlias(inp_tv);
+      }
       auto def = inp->definition();
       auto async_type = ir_utils::getAsyncOpType(def);
 
@@ -657,6 +660,8 @@ class ReadAfterWriteSyncs : public kir::ExprMutator {
       const std::vector<Val*>& tvs) const {
     std::unordered_set<Expr*> last_writes;
     for (auto tv : ir_utils::filterByType<TensorView>(tvs)) {
+      tv = GpuLower::current()->getMaybeTensorProducerAlias(tv);
+
       if (GpuLower::current()->syncMap()->needsRawSync(tv).none()) {
         continue;
       }
diff --git a/csrc/id_model/indexing.cpp b/csrc/id_model/indexing.cpp
index 66c30b8680d..75585db7f72 100644
--- a/csrc/id_model/indexing.cpp
+++ b/csrc/id_model/indexing.cpp
@@ -118,10 +118,13 @@ class AllocationDomainSetup : private kir::IrVisitor {
         continue;
       }
 
-      // Not yet set. This must be an input tensor.
+      // Not yet set. This must be an input tensor or it must be aliased via
+      // aliasTensorProducer, in which case it will not be allocated.
       NVF_ERROR(
-          producer_tv->isFusionInput(),
-          "Expected a fusion input: ",
+          producer_tv->isFusionInput() ||
+              GpuLower::current()->getTensorProducerAlias(producer_tv) !=
+                  nullptr,
+          "Expected a fusion input or aliased tensor but found: ",
           producer_tv->toString());
 
       // For fusion input, we can just use getMaybeAllocationDomain.
diff --git a/tests/cpp/test_gpu2.cpp b/tests/cpp/test_gpu2.cpp
index 1e327fc4c73..abbb6098ae6 100644
--- a/tests/cpp/test_gpu2.cpp
+++ b/tests/cpp/test_gpu2.cpp
@@ -4527,7 +4527,7 @@ TEST_F(NVFuserTest, FusionValidateParallelize6_CUDA) {
   ASSERT_ANY_THROW(fusion.printKernel());
 }
 
-// Repro of #2046
+// Repro of https://github.com/csarofeen/pytorch/issues/2046
 TEST_F(NVFuserTest, FusionValidateParallelize7_CUDA) {
   Fusion fusion;
   FusionGuard fg(&fusion);
@@ -4535,7 +4535,10 @@ TEST_F(NVFuserTest, FusionValidateParallelize7_CUDA) {
   auto tv0 = makeSymbolicTensor(2);
   fusion.addInput(tv0);
 
-  auto tv1 = set(tv0);
+  // The original repro used set instead of exp. However, that is now aliased as
+  // a tensor producer alias so it gets skipped when tv1 is scheduled on Global
+  // memory.
+  auto tv1 = exp(tv0);
   auto tv2 = set(tv1);
   auto tv3 = set(tv1);
   fusion.addOutput(tv2);
diff --git a/tests/cpp/test_matmul.cpp b/tests/cpp/test_matmul.cpp
index e5b281aaea8..510e3bf36be 100644
--- a/tests/cpp/test_matmul.cpp
+++ b/tests/cpp/test_matmul.cpp
@@ -4913,4 +4913,206 @@ TEST_F(HopperMatmulTest, IndexTypeValidation) {
   }
 }
 
+TEST_F(HopperMatmulTest, HSH_NT_128BSwizzle_BroadcastOp) {
+  Fusion fusion;
+  FusionGuard fg(&fusion);
+
+  constexpr int64_t M = 2048, N = 2048, K = 8192;
+  constexpr auto macro = MmaMacro::Hopper_64_256_16;
+  constexpr auto layout = MmaLayout::NT; // [K, M] x [K, N] -> [M, N]
+  constexpr auto swizzle = MmaInputSmemSwizzle::B128;
+  const auto dtype = DataType::Half;
+
+  constexpr bool use_smem_epilogue = false;
+  constexpr bool use_warp_specialization = false;
+
+  constexpr int64_t stages = 4;
+  constexpr int64_t prefetch = 3;
+  const int64_t cta_m = 2 * getM(macro);
+  const int64_t cta_n = 1 * getN(macro);
+
+  constexpr std::tuple<int64_t, int64_t, int64_t> cluster_dims{2, 1, 1};
+
+  // auto tv0 = makeContigConcreteTensor({-1, -1, 1}, dtype);
+  // auto tv1 = makeContigConcreteTensor({-1, 1, -1}, dtype);
+  auto tv0 = makeContigConcreteTensor({-1, -1}, dtype);
+  auto tv1 = makeContigConcreteTensor({-1, -1}, dtype);
+  fusion.addInput(tv0);
+  fusion.addInput(tv1);
+
+  auto tv0b = broadcast(tv0, {false, false, true});
+  auto tv1b = broadcast(tv1, {false, true, false});
+
+  auto tv2 = fusedMultiplySum(tv0b, tv1b, {0});
+
+  // Reorder the accumulator as [M, N, K]
+  // [K, M, N] -> [M, N, K]
+  tv2->reorder({{-3, -1}});
+  tv2->commitLeafToLogical();
+
+  auto tv3 = castOp(DataType::Half, tv2);
+  fusion.addOutput(tv3);
+
+  if constexpr (
+      cluster_dims != std::tuple<int64_t, int64_t, int64_t>{1, 1, 1}) {
+    fusion.manage("cluster_dims", cluster_dims);
+  }
+
+  auto mma_ops = ir_utils::getOpsOfType<MmaOp>(&fusion);
+  NVF_CHECK(
+      1 == mma_ops.size(),
+      "Invalid number of MmaOp instances in fusion definition, expected 1, got ",
+      mma_ops.size());
+  mma_ops.front()->setMacro(macro);
+
+  // gmem [K, M, 1] x gmem [K, 1, N] -mma-> register [M, N, rK]
+  // register [M, N, rK] -cast-> gmem [M, N]
+
+  auto tv0c = tv0b->cacheAfter(LoadStoreOpType::CpAsyncBulkTensorTile);
+  tv0c->setMemoryType(MemoryType::Shared);
+  auto tv1c = tv1b->cacheAfter(LoadStoreOpType::CpAsyncBulkTensorTile);
+  tv1c->setMemoryType(MemoryType::Shared);
+
+  tv0b->setMemoryType(MemoryType::Global);
+  tv1b->setMemoryType(MemoryType::Global);
+
+  TensorView *tv3c = nullptr, *tv3_shmem = nullptr;
+  if (use_smem_epilogue) {
+    tv3_shmem = tv3->cacheBefore();
+    tv3c = tv3_shmem->cacheBefore();
+    tv3_shmem->setMemoryType(MemoryType::Shared);
+    tv3c->setMemoryType(MemoryType::Local);
+    tv3_shmem->definition()->as<LoadStoreOp>()->setOpType(
+        LoadStoreOpType::StMatrix);
+    tv3->definition()->as<LoadStoreOp>()->setOpType(
+        LoadStoreOpType::CpAsyncBulkTensorTile);
+  } else {
+    tv3c = tv3->cacheBefore();
+    tv3c->setMemoryType(MemoryType::Local);
+  }
+
+  // gmem [K, M, 1] -TMA-> smem [K, M, 1]
+  // gmem [K, 1, N] -TMA-> smem [K, 1, N]
+  // smem [K, M, 1] x smem [K, 1, N] -mma-> register [M, N, rK]
+  // register [M, N, rK] -cast-> register [M, N] -set-> gmem [M, N]
+
+  // Create tiles
+  tv2->split(-3, cta_m);
+  tv2->split(-2, cta_n);
+  tv2->split(-1, getK(macro));
+  // [Mo, Mi, No, Ni, Ko, Ki] -> [Mo, No, Ko, Mi, Ni, Ki]
+  tv2->reorder({{-5, -3}, {-3, -2}});
+  tv2->axis(0)->parallelize(ParallelType::BIDy);
+  tv2->axis(1)->parallelize(ParallelType::BIDx);
+
+  TransformPropagator propagator(tv2);
+  MaxLogicalDomainInfoSpanningTree(tv2).traverse(&propagator);
+  scheduler_utils::parallelizeAllLike(tv2);
+
+  // [..., Mi, Ki] -> [..., Ki, Mi]
+  tv0c->reorder({{-3, -1}});
+  tv0c->applyMmaSwizzleForTMALoad(swizzle);
+  tv0c->axis(-6)->parallelize(ParallelType::Unroll);
+  // [..., Ni, Ki] -> [..., Ki, Ni]
+  tv1c->reorder({{-1, -2}});
+  tv1c->applyMmaSwizzleForTMALoad(swizzle);
+  tv1c->axis(-6)->parallelize(ParallelType::Unroll);
+
+  // Strip ParallelType::Bulk from the broadcast tensors, since its definition
+  // is not a TMA
+  for (TensorView* tv : {tv0b, tv1b}) {
+    for (IterDomain* id : tv->getLoopDomain()) {
+      if (id->isBulk()) {
+        id->parallelize(ParallelType::Serial);
+      }
+    }
+  }
+
+  {
+    tv2->split(-3, getM(macro));
+    tv2->split(-2, getN(macro));
+    // [Mo, No, Ko, Mio, Mii, Nio, Nii, Ki]
+    // -> [Mo, No, Ko, Mio, Nio, Mii, Nii, Ki]
+    tv2->reorder({{-4, -3}});
+    tv2->merge(-5);
+    tv2->axis(-4)->parallelize(ParallelType::TIDy);
+    scheduler_utils::BoundedDirectionalTransformPropagator::forward(
+        tv2,
+        -1,
+        {tv3},
+        scheduler_utils::BoundedDirectionalTransformPropagator::Options()
+            .propagateParallelType()
+            .propagateToBoundary());
+  }
+
+  {
+    auto s = mma_utils::MmaSwizzler::scheduleMmaOutputAllocation(
+        tv2->getLoopDomain());
+    tv2->setAllocationDomain(s.as<IterDomain*>(), true);
+    tv2->axis(-1)->parallelize(ParallelType::Mma);
+    tv2->axis(-2)->parallelize(ParallelType::Mma);
+    tv2->axis(-3)->parallelize(ParallelType::Mma);
+  }
+
+  if (!use_smem_epilogue) {
+    for (auto tv : {tv3c, tv3}) {
+      auto s = mma_utils::MmaSwizzler::scheduleMmaOutputAllocation(
+          tv->getLoopDomain());
+      tv->setLoopDomain(s.as<IterDomain*>());
+    }
+    tv3->axis(-1)->parallelize(ParallelType::Vectorize);
+  } else {
+    auto s = mma_utils::MmaSwizzler::scheduleMmaOutputAllocation(
+        tv3c->getLoopDomain());
+    tv3c->setLoopDomain(s.as<IterDomain*>());
+    tv3c->setAllocationDomain(s.as<IterDomain*>(), true);
+
+    constexpr int64_t stmatrix_tile_m = 16;
+    constexpr int64_t stmatrix_tile_n = 16;
+    fusion.manage("st_matrix_m_tile", stmatrix_tile_m);
+    fusion.manage("st_matrix_n_tile", stmatrix_tile_n);
+    fusion.manage("st_matrix_m", getM(macro));
+    fusion.manage("st_matrix_n", getN(macro));
+
+    MmaInputSmemSwizzle store_swizzle =
+        mma_utils::tmaSwizzleSharedMemory(tv3_shmem);
+
+    // This internally calls
+    // Schedule shared memory cache; Output from StMatrix
+    mma_utils::scheduleStMatrixForMmaOutput(
+        tv3_shmem, stmatrix_tile_m, stmatrix_tile_n);
+
+    // Schedule global memory output; Output from TMA Store
+    mma_utils::scheduleTMAStoreForMmaOutput(tv3, store_swizzle);
+  }
+
+  inlineMost(ir_utils::allTvsExcept(&fusion, {tv0, tv1, tv0b, tv1b}));
+
+  if (use_warp_specialization) {
+    tv0c->circularBuffer(stages, prefetch, WarpSpecialized(ParallelType::TIDy));
+    tv1c->circularBuffer(stages, prefetch, WarpSpecialized(ParallelType::TIDy));
+  } else {
+    tv0c->circularBuffer(stages, prefetch);
+    tv1c->circularBuffer(stages, prefetch);
+  }
+
+  auto inputs =
+      matmulAtInput3DHopperSS(M, N, K, layout, data_type_to_aten(dtype));
+  inputs.first = inputs.first.squeeze();
+  inputs.second = inputs.second.squeeze();
+
+  KernelExecutor ke;
+  ke.compile(
+      &fusion, {inputs.first, inputs.second}, LaunchParams(), matmul_cparams);
+  auto cg_outputs = ke.run({inputs.first, inputs.second});
+  auto tref = atMatmul(inputs.first.squeeze(), inputs.second.squeeze(), layout);
+  EXPECT_TRUE(at::allclose(cg_outputs[0].as<at::Tensor>(), tref, 1e-5, 1e-5));
+
+  // The following check fails if the BroadcastOps are not removed at lowering
+  // time, resulting in two intermediate global allocations.
+  const kir::KernelSummary& summary = ke.compiledKernel()->kernel()->summary();
+  EXPECT_EQ(summary.global_allocations.size(), 0)
+      << "Expected to have no intermediate global allocations";
+}
+
 } // namespace nvfuser
diff --git a/tests/cpp/test_remove_trivial_ops.cpp b/tests/cpp/test_remove_trivial_ops.cpp
new file mode 100644
index 00000000000..456ce575546
--- /dev/null
+++ b/tests/cpp/test_remove_trivial_ops.cpp
@@ -0,0 +1,465 @@
+// clang-format off
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2023-present NVIDIA CORPORATION & AFFILIATES.
+ * All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ */
+// clang-format on
+#include <csrc/exceptions.h>
+#include <fusion.h>
+#include <kernel_ir_dispatch.h>
+#include <ops/all_ops.h>
+#include <runtime/executor.h>
+#include <tests/cpp/utils.h>
+#include <tests/cpp/validator.h>
+
+#include <gtest/gtest.h>
+
+namespace nvfuser {
+
+using RemoveTrivialOpsTest = NVFuserTest;
+
+template <typename OpType>
+class KernelContainsExpr : kir::IrVisitor {
+ public:
+  static bool check(const kir::Kernel* kernel) {
+    KernelContainsExpr checker;
+    checker.handle(kernel->topLevelExprs());
+    return checker.has_op_;
+  }
+
+ private:
+  using kir::IrVisitor::handle;
+
+  void handle(OpType* op) final {
+    has_op_ = true;
+  }
+
+ private:
+  bool has_op_ = false;
+};
+
+// Test that we remove a trivial gmem->gmem set at lowering unless it is an
+// alias.
+TEST_F(RemoveTrivialOpsTest, Set) {
+  Fusion fusion;
+  FusionGuard fg(&fusion);
+
+  auto tv0 = makeContigConcreteTensor({2, 3});
+  fusion.addInput(tv0);
+  auto tv1 = set(tv0);
+  auto tv2 = neg(tv1);
+  fusion.addOutput(tv2);
+
+  auto options = at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0);
+  auto t0 = at::randn({2, 3}, options);
+  std::vector<c10::IValue> inputs{t0};
+
+  {
+    // In this case we do not remove the squeeze since it is G->L not G->G
+    KernelExecutor ke;
+    ke.compile(&fusion, inputs);
+    const kir::Kernel* kernel = ke.compiledKernel()->kernel();
+    EXPECT_TRUE(KernelContainsExpr<LoadStoreOp>::check(kernel));
+    auto outputs = ke.run(inputs);
+    testValidate(&fusion, outputs, inputs, __LINE__, __FILE__);
+  }
+
+  {
+    // Setting tv1 to Global means the LoadStoreOp is now G->G. This means that
+    // when we lower it, we will be able to safely remove it
+    tv1->setMemoryType(MemoryType::Global);
+
+    KernelExecutor ke;
+    ke.compile(&fusion, inputs);
+    const kir::Kernel* kernel = ke.compiledKernel()->kernel();
+    EXPECT_FALSE(KernelContainsExpr<LoadStoreOp>::check(kernel));
+    EXPECT_EQ(kernel->summary().global_allocations.size(), 0)
+        << "Expected to have no intermediate global allocations";
+    auto outputs = ke.run(inputs);
+    testValidate(&fusion, outputs, inputs, __LINE__, __FILE__);
+  }
+
+  {
+    NVF_ERROR(tv1->getMemoryType() == MemoryType::Global);
+    // Transpose the input's allocation domain
+    tv0->setAllocationDomain(
+        {
+            tv0->getLogicalDomain().at(1),
+            tv0->getLogicalDomain().at(0),
+        },
+        true);
+
+    KernelExecutor ke;
+
+    std::vector<c10::IValue> transposed_inputs{at::randn({3, 2}, options).t()};
+    ke.compile(&fusion, transposed_inputs);
+    // We do not remove the set because the allocation domains do not match
+    const kir::Kernel* kernel = ke.compiledKernel()->kernel();
+    EXPECT_TRUE(KernelContainsExpr<LoadStoreOp>::check(kernel));
+    EXPECT_EQ(kernel->summary().global_allocations.size(), 1)
+        << "Expected to have one intermediate global allocation";
+    auto outputs = ke.run(transposed_inputs);
+    testValidate(&fusion, outputs, transposed_inputs, __LINE__, __FILE__);
+  }
+
+  {
+    // Reset tv0's allocation domain but add a new output which reuses the
+    // buffer of the input tv0
+    tv0->setAllocationDomain(tv0->getLogicalDomain(), true);
+    auto* tv3 = exp(tv2);
+    fusion.addOutput(tv3);
+    fusion.aliasOutputToInput(tv3, tv0, AllocationType::ReuseBuffer);
+
+    KernelExecutor ke;
+    ke.compile(&fusion, inputs);
+    // We do not remove the squeeze because the input is target of an io alias
+    const kir::Kernel* kernel = ke.compiledKernel()->kernel();
+    EXPECT_TRUE(KernelContainsExpr<LoadStoreOp>::check(kernel));
+    // Run with copy of inputs since we will overwrite the original inputs,
+    // complicating validation
+    std::vector<c10::IValue> inputs_copy{t0.clone()};
+    auto outputs = ke.run(inputs_copy);
+    testValidate(&fusion, outputs, inputs, __LINE__, __FILE__);
+  }
+}
+
+// Test that we remove a trivial gmem->gmem broadcast at lowering
+TEST_F(RemoveTrivialOpsTest, Broadcast) {
+  Fusion fusion;
+  FusionGuard fg(&fusion);
+
+  auto tv0 = makeContigConcreteTensor({2, 3});
+  fusion.addInput(tv0);
+  auto tv1 = broadcast(tv0, {false, true, false});
+  auto tv2 = neg(tv1);
+  fusion.addOutput(tv2);
+
+  auto options = at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0);
+  auto t0 = at::randn({2, 3}, options);
+  std::vector<c10::IValue> inputs{t0};
+
+  {
+    // In this case we do not remove the broadcast since it is G->L not G->G
+    KernelExecutor ke;
+    ke.compile(&fusion, inputs);
+    const kir::Kernel* kernel = ke.compiledKernel()->kernel();
+    EXPECT_TRUE(KernelContainsExpr<BroadcastOp>::check(kernel));
+    auto outputs = ke.run(inputs);
+    testValidate(&fusion, outputs, inputs, __LINE__, __FILE__);
+  }
+
+  {
+    // Setting tv1 to Global means the BroadcastOp is now G->G. This means that
+    // when we lower it, we will be able to safely remove it
+    tv1->setMemoryType(MemoryType::Global);
+
+    KernelExecutor ke;
+    ke.compile(&fusion, inputs);
+    const kir::Kernel* kernel = ke.compiledKernel()->kernel();
+    EXPECT_FALSE(KernelContainsExpr<BroadcastOp>::check(kernel));
+    EXPECT_EQ(kernel->summary().global_allocations.size(), 0)
+        << "Expected to have no intermediate global allocations";
+    auto outputs = ke.run(inputs);
+    testValidate(&fusion, outputs, inputs, __LINE__, __FILE__);
+  }
+
+  {
+    NVF_ERROR(tv1->getMemoryType() == MemoryType::Global);
+    // Transpose the input's allocation domain
+    tv0->setAllocationDomain(
+        {
+            tv0->getLogicalDomain().at(1),
+            tv0->getLogicalDomain().at(0),
+        },
+        true);
+
+    KernelExecutor ke;
+
+    std::vector<c10::IValue> transposed_inputs{at::randn({3, 2}, options).t()};
+    ke.compile(&fusion, transposed_inputs);
+    // We do not remove the broadcast because the allocation domains do not
+    // match
+    const kir::Kernel* kernel = ke.compiledKernel()->kernel();
+    EXPECT_TRUE(KernelContainsExpr<BroadcastOp>::check(kernel));
+    EXPECT_EQ(kernel->summary().global_allocations.size(), 1)
+        << "Expected to have one intermediate global allocation";
+    auto outputs = ke.run(transposed_inputs);
+    testValidate(&fusion, outputs, transposed_inputs, __LINE__, __FILE__);
+  }
+
+  {
+    // Reset tv0's allocation domain but set tv1 as output
+    tv0->setAllocationDomain(tv0->getLogicalDomain(), true);
+    fusion.addOutput(tv1);
+
+    KernelExecutor ke;
+    ke.compile(&fusion, inputs);
+    // We do not remove the broadcast because the output is a fusion output
+    const kir::Kernel* kernel = ke.compiledKernel()->kernel();
+    EXPECT_TRUE(KernelContainsExpr<BroadcastOp>::check(kernel));
+    auto outputs = ke.run(inputs);
+    testValidate(&fusion, outputs, inputs, __LINE__, __FILE__);
+  }
+}
+
+// Test that we remove a trivial gmem->gmem squeeze at lowering
+TEST_F(RemoveTrivialOpsTest, Squeeze) {
+  Fusion fusion;
+  FusionGuard fg(&fusion);
+
+  auto tv0 = makeContigConcreteTensor({2, 1, 3});
+  fusion.addInput(tv0);
+  auto tv1 = squeeze(tv0, {1});
+  auto tv2 = neg(tv1);
+  fusion.addOutput(tv2);
+
+  auto options = at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0);
+  auto t0 = at::randn({2, 1, 3}, options);
+  std::vector<c10::IValue> inputs{t0};
+
+  {
+    // In this case we do not remove the squeeze since it is G->L not G->G
+    KernelExecutor ke;
+    ke.compile(&fusion, inputs);
+    const kir::Kernel* kernel = ke.compiledKernel()->kernel();
+    // NOTE: SqueezeOp is converted to LoadStoreOp by the loadStoreOpInserter
+    // lowering pass
+    EXPECT_TRUE(KernelContainsExpr<LoadStoreOp>::check(kernel));
+    auto outputs = ke.run(inputs);
+    testValidate(&fusion, outputs, inputs, __LINE__, __FILE__);
+  }
+
+  {
+    // Setting tv1 to Global means the SqueezeOp is now G->G. This means that
+    // when we lower it, we will be able to safely remove it
+    tv1->setMemoryType(MemoryType::Global);
+
+    KernelExecutor ke;
+    ke.compile(&fusion, inputs);
+    const kir::Kernel* kernel = ke.compiledKernel()->kernel();
+    EXPECT_FALSE(KernelContainsExpr<LoadStoreOp>::check(kernel));
+    EXPECT_EQ(kernel->summary().global_allocations.size(), 0)
+        << "Expected to have no intermediate global allocations";
+    auto outputs = ke.run(inputs);
+    testValidate(&fusion, outputs, inputs, __LINE__, __FILE__);
+  }
+
+  {
+    NVF_ERROR(tv1->getMemoryType() == MemoryType::Global);
+    // Transpose the input's allocation domain
+    tv0->setAllocationDomain(
+        {
+            tv0->getLogicalDomain().at(2),
+            tv0->getLogicalDomain().at(1),
+            tv0->getLogicalDomain().at(0),
+        },
+        true);
+
+    KernelExecutor ke;
+
+    std::vector<c10::IValue> transposed_inputs{
+        at::randn({3, 2}, options).t().unsqueeze(1)};
+    ke.compile(&fusion, transposed_inputs);
+    // We do not remove the broadcast because the allocation domains do not
+    // match
+    const kir::Kernel* kernel = ke.compiledKernel()->kernel();
+    EXPECT_TRUE(KernelContainsExpr<LoadStoreOp>::check(kernel));
+    EXPECT_EQ(kernel->summary().global_allocations.size(), 1)
+        << "Expected to have one intermediate global allocation";
+    auto outputs = ke.run(transposed_inputs);
+    testValidate(&fusion, outputs, transposed_inputs, __LINE__, __FILE__);
+  }
+
+  {
+    // Reset tv0's allocation domain but set tv1 as output
+    tv0->setAllocationDomain(tv0->getLogicalDomain(), true);
+    fusion.addOutput(tv1);
+
+    KernelExecutor ke;
+    ke.compile(&fusion, inputs);
+    // We do not remove the squeeze because the output is a fusion output
+    const kir::Kernel* kernel = ke.compiledKernel()->kernel();
+    EXPECT_TRUE(KernelContainsExpr<LoadStoreOp>::check(kernel));
+    auto outputs = ke.run(inputs);
+    testValidate(&fusion, outputs, inputs, __LINE__, __FILE__);
+  }
+}
+
+// Test that we remove a gmem->gmem permute whose allocation domain is
+// unpermuted at lowering
+TEST_F(RemoveTrivialOpsTest, Permute) {
+  Fusion fusion;
+  FusionGuard fg(&fusion);
+
+  auto tv0 = makeContigConcreteTensor({2, 3});
+  fusion.addInput(tv0);
+  // This is a logical transpose. It implies data shuffling unless the
+  // allocation domain is also transposed.
+  auto tv1 = transpose(tv0);
+  auto tv2 = neg(tv1);
+  fusion.addOutput(tv2);
+
+  auto options = at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0);
+  auto t0 = at::randn({2, 3}, options);
+  std::vector<c10::IValue> inputs{t0};
+
+  tv1->setAllocationDomain(
+      {
+          tv1->getLogicalDomain().at(1),
+          tv1->getLogicalDomain().at(0),
+      },
+      true);
+
+  {
+    // In this case we do not remove the permute since it is G->L not G->G
+    KernelExecutor ke;
+    ke.compile(&fusion, inputs);
+    const kir::Kernel* kernel = ke.compiledKernel()->kernel();
+    EXPECT_TRUE(KernelContainsExpr<LoadStoreOp>::check(kernel));
+    auto outputs = ke.run(inputs);
+    testValidate(&fusion, outputs, inputs, __LINE__, __FILE__);
+  }
+
+  {
+    // Setting tv1 to Global means the LoadStoreOp is now G->G. This means that
+    // when we lower it, we will be able to safely remove it
+    tv1->setMemoryType(MemoryType::Global);
+
+    KernelExecutor ke;
+    ke.compile(&fusion, inputs);
+    const kir::Kernel* kernel = ke.compiledKernel()->kernel();
+    EXPECT_FALSE(KernelContainsExpr<LoadStoreOp>::check(kernel));
+    EXPECT_EQ(kernel->summary().global_allocations.size(), 0)
+        << "Expected to have no intermediate global allocations";
+    auto outputs = ke.run(inputs);
+    testValidate(&fusion, outputs, inputs, __LINE__, __FILE__);
+  }
+
+  {
+    NVF_ERROR(tv1->getMemoryType() == MemoryType::Global);
+    tv1->setAllocationDomain(tv1->getLogicalDomain(), true);
+
+    KernelExecutor ke;
+
+    ke.compile(&fusion, inputs);
+    // We do not remove the transpose because the allocation domains do not
+    // match
+    const kir::Kernel* kernel = ke.compiledKernel()->kernel();
+    EXPECT_TRUE(KernelContainsExpr<LoadStoreOp>::check(kernel));
+    EXPECT_EQ(kernel->summary().global_allocations.size(), 1)
+        << "Expected to have one intermediate global allocation";
+    auto outputs = ke.run(inputs);
+    testValidate(&fusion, outputs, inputs, __LINE__, __FILE__);
+  }
+
+  {
+    // Reset tv1's allocation domain but also set it as an output
+    tv1->setAllocationDomain(
+        {
+            tv1->getLogicalDomain().at(1),
+            tv1->getLogicalDomain().at(0),
+        },
+        true);
+    fusion.addOutput(tv1);
+
+    KernelExecutor ke;
+    ke.compile(&fusion, inputs);
+    // We do not remove the transpose because the output is a fusion output
+    const kir::Kernel* kernel = ke.compiledKernel()->kernel();
+    EXPECT_TRUE(KernelContainsExpr<LoadStoreOp>::check(kernel));
+    auto outputs = ke.run(inputs);
+    testValidate(&fusion, outputs, inputs, __LINE__, __FILE__);
+  }
+}
+
+// Test that we remove a gmem->gmem->gmem broadcast+squeeze at lowering
+TEST_F(RemoveTrivialOpsTest, BroadcastSqueeze) {
+  Fusion fusion;
+  FusionGuard fg(&fusion);
+
+  auto tv0 = makeContigConcreteTensor({2, 3});
+  fusion.addInput(tv0);
+  auto tv1 = broadcast(tv0, {false, true, false});
+  auto tv2 = squeeze(tv1, {1});
+  auto tv3 = neg(tv2);
+  fusion.addOutput(tv3);
+
+  auto options = at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0);
+  auto t0 = at::randn({2, 3}, options);
+  std::vector<c10::IValue> inputs{t0};
+
+  {
+    // In this case we do not remove the broadcast since it is G->L not G->G,
+    // and the squeeze is L->L
+    KernelExecutor ke;
+    ke.compile(&fusion, inputs);
+    const kir::Kernel* kernel = ke.compiledKernel()->kernel();
+    // NOTE: SqueezeOp is converted to LoadStoreOp by the loadStoreOpInserter
+    // lowering pass
+    EXPECT_TRUE(KernelContainsExpr<BroadcastOp>::check(kernel));
+    EXPECT_TRUE(KernelContainsExpr<LoadStoreOp>::check(kernel));
+    auto outputs = ke.run(inputs);
+    testValidate(&fusion, outputs, inputs, __LINE__, __FILE__);
+  }
+
+  {
+    // Set G->G->G so that both ops get removed
+    tv1->setMemoryType(MemoryType::Global);
+    tv2->setMemoryType(MemoryType::Global);
+
+    KernelExecutor ke;
+    ke.compile(&fusion, inputs);
+    const kir::Kernel* kernel = ke.compiledKernel()->kernel();
+    EXPECT_FALSE(KernelContainsExpr<BroadcastOp>::check(kernel));
+    EXPECT_FALSE(KernelContainsExpr<LoadStoreOp>::check(kernel));
+    EXPECT_EQ(kernel->summary().global_allocations.size(), 0)
+        << "Expected to have no intermediate global allocations";
+    auto outputs = ke.run(inputs);
+    testValidate(&fusion, outputs, inputs, __LINE__, __FILE__);
+  }
+
+  {
+    NVF_ERROR(tv1->getMemoryType() == MemoryType::Global);
+    NVF_ERROR(tv2->getMemoryType() == MemoryType::Global);
+    // Transpose the input's allocation domain
+    tv0->setAllocationDomain(
+        {
+            tv0->getLogicalDomain().at(1),
+            tv0->getLogicalDomain().at(0),
+        },
+        true);
+
+    KernelExecutor ke;
+
+    std::vector<c10::IValue> transposed_inputs{at::randn({3, 2}, options).t()};
+    ke.compile(&fusion, transposed_inputs);
+    // We do not remove the broadcast because the allocation domains do not
+    // match
+    const kir::Kernel* kernel = ke.compiledKernel()->kernel();
+    EXPECT_TRUE(KernelContainsExpr<BroadcastOp>::check(kernel));
+    EXPECT_FALSE(KernelContainsExpr<LoadStoreOp>::check(kernel));
+    EXPECT_EQ(kernel->summary().global_allocations.size(), 1)
+        << "Expected to have one intermediate global allocation";
+    auto outputs = ke.run(transposed_inputs);
+    testValidate(&fusion, outputs, transposed_inputs, __LINE__, __FILE__);
+  }
+
+  {
+    // Reset tv0's allocation domain but set tv1 as output
+    tv0->setAllocationDomain(tv0->getLogicalDomain(), true);
+    fusion.addOutput(tv1);
+
+    KernelExecutor ke;
+    ke.compile(&fusion, inputs);
+    // We do not remove the broadcast because the output is a fusion output, but
+    // we _do_ remove the squeeze
+    const kir::Kernel* kernel = ke.compiledKernel()->kernel();
+    EXPECT_TRUE(KernelContainsExpr<BroadcastOp>::check(kernel));
+    EXPECT_FALSE(KernelContainsExpr<LoadStoreOp>::check(kernel));
+    auto outputs = ke.run(inputs);
+    testValidate(&fusion, outputs, inputs, __LINE__, __FILE__);
+  }
+}
+
+} // namespace nvfuser


Reviews:
user: login: naoyam
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-02-28T21:24:24Z
user: login: zasdfgbnm
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-02-28T22:16:29Z
user: login: jacobhinkle
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-02-28T23:51:31Z
user: login: jacobhinkle
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-02-28T23:53:02Z
user: login: jacobhinkle
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-01T18:25:38Z
user: login: jacobhinkle
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-01T18:30:13Z
user: login: jacobhinkle
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-03T14:31:15Z
user: login: jacobhinkle
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-03T15:11:26Z
user: login: naoyam
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-03T18:54:08Z
user: login: jacobhinkle
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-03T21:47:09Z
user: login: jacobhinkle
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-03T21:50:04Z
user: login: naoyam
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-04T02:25:00Z
user: login: naoyam
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-04T05:38:17Z
user: login: naoyam
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-04T05:45:53Z
user: login: naoyam
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-04T05:52:04Z
user: login: naoyam
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-04T06:01:33Z
user: login: jacobhinkle
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-04T10:09:55Z
user: login: jacobhinkle
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-04T15:06:35Z
user: login: jacobhinkle
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-04T15:08:12Z
user: login: jacobhinkle
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-04T15:25:08Z
user: login: jacobhinkle
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-04T15:26:20Z
user: login: naoyam
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-04T16:51:57Z
user: login: naoyam
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-04T17:04:42Z
user: login: jacobhinkle
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-04T17:22:48Z
user: login: jacobhinkle
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-04T20:46:08Z
user: login: naoyam
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-04T22:54:47Z
user: login: naoyam
body: LGTM
state: APPROVED
author_association: COLLABORATOR
submitted_at: 2025-03-04T22:55:12Z

Review comments:
diff_hunk: @@ -120,7 +120,9 @@ class AllocationDomainSetup : private kir::IrVisitor {
 
       // Not yet set. This must be an input tensor.
path: csrc/id_model/indexing.cpp
user: login: naoyam
body: Please update the comment and the error message
created_at: 2025-02-28T21:02:18Z
author_association: COLLABORATOR
diff_hunk: @@ -1715,6 +1715,11 @@ std::vector<Expr*> ExprSegmentationSorter::getExprs() const {
       if (!lower_utils::isScalarExpr(expr)) {
         active_exprs = &remaining_exprs;
       }
+      if (expr->outputs().size() == 1L && expr->output(0)->isA<TensorView>() &&
path: csrc/device_lower/pass/expr_sort.cpp
user: login: naoyam
body: Why is this? Could you add comments?
created_at: 2025-02-28T21:09:12Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,16 @@
+// clang-format off
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2023-present NVIDIA CORPORATION & AFFILIATES.
+ * All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ */
+// clang-format on
+#pragma once
+
+#include <fusion.h>
+
+namespace nvfuser {
+
+void findTensorProducerAliases(Fusion* fusion);
path: csrc/device_lower/analysis/tensor_producer_aliases.h
user: login: naoyam
body: Please add comments
created_at: 2025-02-28T21:10:37Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,115 @@
+// clang-format off
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2023-present NVIDIA CORPORATION & AFFILIATES.
+ * All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ */
+// clang-format on
+#include <device_lower/analysis/tensor_producer_aliases.h>
+#include <device_lower/lower2device.h>
+#include <kernel_ir_dispatch.h>
+
+#include <unordered_set>
+
+namespace nvfuser {
+
+namespace {
+
+bool isTrivialExpr(Expr* expr) {
+  if (expr->outputs().size() != 1 || expr->inputs().size() != 1 ||
+      expr->output(0)->isFusionOutput() ||
+      !expr->isOneOf<BroadcastOp, LoadStoreOp, SqueezeOp>() ||
+      !expr->input(0)->isA<TensorView>() ||
+      !expr->output(0)->isA<TensorView>() ||
+      expr->input(0)->as<TensorView>()->getMemoryType() != MemoryType::Global ||
+      expr->output(0)->as<TensorView>()->getMemoryType() !=
+          MemoryType::Global) {
+    return false;
+  }
+  // This is a tensor op that does no computation. However, it may still be
+  // non-trivial if the allocation domains are not equivalent, as that would
+  // indicate data movement. We check that any non-Broadcast and non-Reduction
+  // domains in the two allocation domains are exact mapped and appear in the
+  // same order. We also check that the tensors are fully contiguous so that we
+  // can re-use the exact same linear index.
+  // TODO: support discontiguous inputs. The output should be an intermediate
+  // tensor which we would always assume to be contiguous, but the input might
+  // not be. This would necessitate using a different linear index.
+  TensorView* in = expr->input(0)->as<TensorView>();
+  TensorView* out = expr->output(0)->as<TensorView>();
+  size_t in_pos = 0;
+  size_t out_pos = 0;
+  const std::vector<IterDomain*>& in_alloc = in->getMaybeAllocationDomain();
+  const std::vector<IterDomain*>& out_alloc = out->getMaybeAllocationDomain();
+  const std::vector<std::optional<bool>>& in_contig = in->getContiguity();
+  const std::vector<std::optional<bool>>& out_contig = out->getContiguity();
+
+  // TODO: Cache this IdModel build it only once, if needed, to prevent building
+  // it for every G->G op in the Fusion.
+  IdModel id_model(in->fusion(), /*build_graphs=*/false);
path: csrc/device_lower/analysis/tensor_producer_aliases.cpp
user: login: naoyam
body: `GpuLower` has an IdModel by default.
created_at: 2025-02-28T21:20:17Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,115 @@
+// clang-format off
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2023-present NVIDIA CORPORATION & AFFILIATES.
+ * All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ */
+// clang-format on
+#include <device_lower/analysis/tensor_producer_aliases.h>
+#include <device_lower/lower2device.h>
+#include <kernel_ir_dispatch.h>
+
+#include <unordered_set>
+
+namespace nvfuser {
+
+namespace {
+
+bool isTrivialExpr(Expr* expr) {
+  if (expr->outputs().size() != 1 || expr->inputs().size() != 1 ||
+      expr->output(0)->isFusionOutput() ||
+      !expr->isOneOf<BroadcastOp, LoadStoreOp, SqueezeOp>() ||
+      !expr->input(0)->isA<TensorView>() ||
+      !expr->output(0)->isA<TensorView>() ||
+      expr->input(0)->as<TensorView>()->getMemoryType() != MemoryType::Global ||
+      expr->output(0)->as<TensorView>()->getMemoryType() !=
+          MemoryType::Global) {
+    return false;
+  }
+  // This is a tensor op that does no computation. However, it may still be
+  // non-trivial if the allocation domains are not equivalent, as that would
+  // indicate data movement. We check that any non-Broadcast and non-Reduction
+  // domains in the two allocation domains are exact mapped and appear in the
+  // same order. We also check that the tensors are fully contiguous so that we
+  // can re-use the exact same linear index.
+  // TODO: support discontiguous inputs. The output should be an intermediate
+  // tensor which we would always assume to be contiguous, but the input might
+  // not be. This would necessitate using a different linear index.
+  TensorView* in = expr->input(0)->as<TensorView>();
+  TensorView* out = expr->output(0)->as<TensorView>();
+  size_t in_pos = 0;
+  size_t out_pos = 0;
+  const std::vector<IterDomain*>& in_alloc = in->getMaybeAllocationDomain();
+  const std::vector<IterDomain*>& out_alloc = out->getMaybeAllocationDomain();
+  const std::vector<std::optional<bool>>& in_contig = in->getContiguity();
+  const std::vector<std::optional<bool>>& out_contig = out->getContiguity();
+
+  // TODO: Cache this IdModel build it only once, if needed, to prevent building
+  // it for every G->G op in the Fusion.
+  IdModel id_model(in->fusion(), /*build_graphs=*/false);
+  id_model.buildExactGraph();
+  const ValGraph& exact_graph = id_model.idGraph(IdMappingMode::EXACT);
+
+  while (true) {
+    while (in_pos < in_alloc.size() &&
+           (in_alloc.at(in_pos)->isBroadcast() ||
+            in_alloc.at(in_pos)->isReduction())) {
+      in_pos++;
+    }
+    while (out_pos < out_alloc.size() && out_alloc.at(out_pos)->isBroadcast()) {
+      out_pos++;
+    }
+    if (in_pos >= in_alloc.size()) {
+      NVF_ERROR(
+          out_pos >= out_alloc.size(),
+          "Found non-broadcast output domain ",
+          out_alloc.at(out_pos)->toString(),
+          " which has no corresponding input allocation domain");
+      break;
+    }
+    if (out_pos >= out_alloc.size()) {
+      NVF_ERROR(
+          in_pos >= in_alloc.size(),
+          "Found non-broadcast non-reduction input domain ",
+          in_alloc.at(in_pos)->toString(),
+          " which has no corresponding output allocation domain");
+      break;
+    }
path: csrc/device_lower/analysis/tensor_producer_aliases.cpp
user: login: naoyam
body: Are both of these two conditions necessary? Looks like just one of them is sufficient.
created_at: 2025-02-28T21:23:32Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,115 @@
+// clang-format off
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2023-present NVIDIA CORPORATION & AFFILIATES.
+ * All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ */
+// clang-format on
+#include <device_lower/analysis/tensor_producer_aliases.h>
+#include <device_lower/lower2device.h>
+#include <kernel_ir_dispatch.h>
+
+#include <unordered_set>
+
+namespace nvfuser {
+
+namespace {
+
+bool isTrivialExpr(Expr* expr) {
+  if (expr->outputs().size() != 1 || expr->inputs().size() != 1 ||
+      expr->output(0)->isFusionOutput() ||
+      !expr->isOneOf<BroadcastOp, LoadStoreOp, SqueezeOp>() ||
path: csrc/device_lower/analysis/tensor_producer_aliases.cpp
user: login: zasdfgbnm
body: Wouldn't this automatically guarantee the input and output size check?
created_at: 2025-02-28T21:59:56Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,115 @@
+// clang-format off
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2023-present NVIDIA CORPORATION & AFFILIATES.
+ * All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ */
+// clang-format on
+#include <device_lower/analysis/tensor_producer_aliases.h>
+#include <device_lower/lower2device.h>
+#include <kernel_ir_dispatch.h>
+
+#include <unordered_set>
+
+namespace nvfuser {
+
+namespace {
+
+bool isTrivialExpr(Expr* expr) {
+  if (expr->outputs().size() != 1 || expr->inputs().size() != 1 ||
+      expr->output(0)->isFusionOutput() ||
+      !expr->isOneOf<BroadcastOp, LoadStoreOp, SqueezeOp>() ||
+      !expr->input(0)->isA<TensorView>() ||
+      !expr->output(0)->isA<TensorView>() ||
+      expr->input(0)->as<TensorView>()->getMemoryType() != MemoryType::Global ||
+      expr->output(0)->as<TensorView>()->getMemoryType() !=
+          MemoryType::Global) {
+    return false;
+  }
+  // This is a tensor op that does no computation. However, it may still be
+  // non-trivial if the allocation domains are not equivalent, as that would
+  // indicate data movement. We check that any non-Broadcast and non-Reduction
+  // domains in the two allocation domains are exact mapped and appear in the
+  // same order. We also check that the tensors are fully contiguous so that we
+  // can re-use the exact same linear index.
path: csrc/device_lower/analysis/tensor_producer_aliases.cpp
user: login: zasdfgbnm
body: Do we need to do that? IIUC the lowering decides how we allocate gmem intermediate buffer, so the lowering decide whether it should be allocated as contiguous or not. So here, we can just decide now that we want to allocate the intermediate tensor the same as the input.
created_at: 2025-02-28T22:02:42Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,16 @@
+// clang-format off
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2023-present NVIDIA CORPORATION & AFFILIATES.
+ * All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ */
+// clang-format on
+#pragma once
+
+#include <fusion.h>
+
+namespace nvfuser {
+
+void findTensorProducerAliases(Fusion* fusion);
path: csrc/device_lower/analysis/tensor_producer_aliases.h
user: login: zasdfgbnm
body: It would nice to describe the high level goal of this aliasing mechanism. What problem does it solve? Why do we need it?
created_at: 2025-02-28T22:06:12Z
author_association: COLLABORATOR
diff_hunk: @@ -377,6 +417,10 @@ class GpuLower : public NonCopyable {
   // Information about tensor memory usage
   TensorMemoryInfo tmem_info_;
 
+  // Track TensorViews that will be aliased to their producers because of
+  // trivial ops and scheduling such that the same index is used.
+  std::unordered_map<TensorView*, TensorView*> tensor_producer_alias_map_;
path: csrc/device_lower/lower2device.h
user: login: zasdfgbnm
body: Could you add in your comment describing that this map maps a TV to a fusion input, instead of its direct producer?
created_at: 2025-02-28T22:13:10Z
author_association: COLLABORATOR
diff_hunk: @@ -45,14 +46,21 @@ Val* IndexLowering::lowerSrcIndex(
     DataType as_type) const {
   if (auto tv = dynamic_cast<TensorView*>(src)) {
     NVF_ERROR(dst->isA<TensorView>());
-    return Index::getProducerIndex(
+    kir::TensorIndex* tind = Index::getProducerIndex(
         tv,
path: csrc/device_lower/pass/index.cpp
user: login: zasdfgbnm
body: Should we use `GpuLower::current()->getTensorProducerAlias(tv))` instead of `tv` for `getProducerIndex`?
created_at: 2025-02-28T22:15:49Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,115 @@
+// clang-format off
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2023-present NVIDIA CORPORATION & AFFILIATES.
+ * All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ */
+// clang-format on
+#include <device_lower/analysis/tensor_producer_aliases.h>
+#include <device_lower/lower2device.h>
+#include <kernel_ir_dispatch.h>
+
+#include <unordered_set>
+
+namespace nvfuser {
+
+namespace {
+
+bool isTrivialExpr(Expr* expr) {
+  if (expr->outputs().size() != 1 || expr->inputs().size() != 1 ||
+      expr->output(0)->isFusionOutput() ||
+      !expr->isOneOf<BroadcastOp, LoadStoreOp, SqueezeOp>() ||
+      !expr->input(0)->isA<TensorView>() ||
+      !expr->output(0)->isA<TensorView>() ||
+      expr->input(0)->as<TensorView>()->getMemoryType() != MemoryType::Global ||
+      expr->output(0)->as<TensorView>()->getMemoryType() !=
+          MemoryType::Global) {
+    return false;
+  }
+  // This is a tensor op that does no computation. However, it may still be
+  // non-trivial if the allocation domains are not equivalent, as that would
+  // indicate data movement. We check that any non-Broadcast and non-Reduction
+  // domains in the two allocation domains are exact mapped and appear in the
+  // same order. We also check that the tensors are fully contiguous so that we
+  // can re-use the exact same linear index.
path: csrc/device_lower/analysis/tensor_producer_aliases.cpp
user: login: jacobhinkle
body: Good point. I was worried about generating an index that included strides of the intermediate tensor since those will not be available later. But we should just do as you say, then respect the alias when grabbing the stride during indexing.
created_at: 2025-02-28T23:51:31Z
author_association: COLLABORATOR
diff_hunk: @@ -45,14 +46,21 @@ Val* IndexLowering::lowerSrcIndex(
     DataType as_type) const {
   if (auto tv = dynamic_cast<TensorView*>(src)) {
     NVF_ERROR(dst->isA<TensorView>());
-    return Index::getProducerIndex(
+    kir::TensorIndex* tind = Index::getProducerIndex(
         tv,
path: csrc/device_lower/pass/index.cpp
user: login: jacobhinkle
body: I think the issue with that is that sometimes we use a pairwise map to determine the index, in which case producer needs to be a direct producer.
created_at: 2025-02-28T23:53:01Z
author_association: COLLABORATOR
diff_hunk: @@ -377,6 +417,10 @@ class GpuLower : public NonCopyable {
   // Information about tensor memory usage
   TensorMemoryInfo tmem_info_;
 
+  // Track TensorViews that will be aliased to their producers because of
+  // trivial ops and scheduling such that the same index is used.
+  std::unordered_map<TensorView*, TensorView*> tensor_producer_alias_map_;
path: csrc/device_lower/lower2device.h
user: login: jacobhinkle
body: I will improve this comment. As you say, it seems to indicate that they are aliases to _direct_ producers which is not what I meant. Note that currently the target of the alias doesn't need to be a fusion input, although we could impose that restriction easily.
created_at: 2025-03-01T18:25:38Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,115 @@
+// clang-format off
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2023-present NVIDIA CORPORATION & AFFILIATES.
+ * All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ */
+// clang-format on
+#include <device_lower/analysis/tensor_producer_aliases.h>
+#include <device_lower/lower2device.h>
+#include <kernel_ir_dispatch.h>
+
+#include <unordered_set>
+
+namespace nvfuser {
+
+namespace {
+
+bool isTrivialExpr(Expr* expr) {
+  if (expr->outputs().size() != 1 || expr->inputs().size() != 1 ||
+      expr->output(0)->isFusionOutput() ||
+      !expr->isOneOf<BroadcastOp, LoadStoreOp, SqueezeOp>() ||
+      !expr->input(0)->isA<TensorView>() ||
+      !expr->output(0)->isA<TensorView>() ||
+      expr->input(0)->as<TensorView>()->getMemoryType() != MemoryType::Global ||
+      expr->output(0)->as<TensorView>()->getMemoryType() !=
+          MemoryType::Global) {
+    return false;
+  }
+  // This is a tensor op that does no computation. However, it may still be
+  // non-trivial if the allocation domains are not equivalent, as that would
+  // indicate data movement. We check that any non-Broadcast and non-Reduction
+  // domains in the two allocation domains are exact mapped and appear in the
+  // same order. We also check that the tensors are fully contiguous so that we
+  // can re-use the exact same linear index.
+  // TODO: support discontiguous inputs. The output should be an intermediate
+  // tensor which we would always assume to be contiguous, but the input might
+  // not be. This would necessitate using a different linear index.
+  TensorView* in = expr->input(0)->as<TensorView>();
+  TensorView* out = expr->output(0)->as<TensorView>();
+  size_t in_pos = 0;
+  size_t out_pos = 0;
+  const std::vector<IterDomain*>& in_alloc = in->getMaybeAllocationDomain();
+  const std::vector<IterDomain*>& out_alloc = out->getMaybeAllocationDomain();
+  const std::vector<std::optional<bool>>& in_contig = in->getContiguity();
+  const std::vector<std::optional<bool>>& out_contig = out->getContiguity();
+
+  // TODO: Cache this IdModel build it only once, if needed, to prevent building
+  // it for every G->G op in the Fusion.
+  IdModel id_model(in->fusion(), /*build_graphs=*/false);
+  id_model.buildExactGraph();
+  const ValGraph& exact_graph = id_model.idGraph(IdMappingMode::EXACT);
+
+  while (true) {
+    while (in_pos < in_alloc.size() &&
+           (in_alloc.at(in_pos)->isBroadcast() ||
+            in_alloc.at(in_pos)->isReduction())) {
+      in_pos++;
+    }
+    while (out_pos < out_alloc.size() && out_alloc.at(out_pos)->isBroadcast()) {
+      out_pos++;
+    }
+    if (in_pos >= in_alloc.size()) {
+      NVF_ERROR(
+          out_pos >= out_alloc.size(),
+          "Found non-broadcast output domain ",
+          out_alloc.at(out_pos)->toString(),
+          " which has no corresponding input allocation domain");
+      break;
+    }
+    if (out_pos >= out_alloc.size()) {
+      NVF_ERROR(
+          in_pos >= in_alloc.size(),
+          "Found non-broadcast non-reduction input domain ",
+          in_alloc.at(in_pos)->toString(),
+          " which has no corresponding output allocation domain");
+      break;
+    }
path: csrc/device_lower/analysis/tensor_producer_aliases.cpp
user: login: jacobhinkle
body: They are both just checking that the number of non-ignored IDs in each allocation domain match, so one checks a<=b and the other checks a>=b. Of course, probably neither one is necessary because that is the kind of thing that's already checked by validateDomainEquivalence.
created_at: 2025-03-01T18:30:13Z
author_association: COLLABORATOR
diff_hunk: @@ -45,14 +46,21 @@ Val* IndexLowering::lowerSrcIndex(
     DataType as_type) const {
   if (auto tv = dynamic_cast<TensorView*>(src)) {
     NVF_ERROR(dst->isA<TensorView>());
-    return Index::getProducerIndex(
+    kir::TensorIndex* tind = Index::getProducerIndex(
         tv,
path: csrc/device_lower/pass/index.cpp
user: login: jacobhinkle
body: In particular, I hit an error here when I try doing as you say: https://github.com/NVIDIA/Fuser/blob/27258f5849fc2b6525d65d098e26731c64f348e0/csrc/index_compute.cpp#L1774-L1775
created_at: 2025-03-03T14:31:15Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,16 @@
+// clang-format off
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2023-present NVIDIA CORPORATION & AFFILIATES.
+ * All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ */
+// clang-format on
+#pragma once
+
+#include <fusion.h>
+
+namespace nvfuser {
+
+void findTensorProducerAliases(Fusion* fusion);
path: csrc/device_lower/analysis/tensor_producer_aliases.h
user: login: jacobhinkle
body: Added a comment in latest push
created_at: 2025-03-03T15:11:26Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,122 @@
+// clang-format off
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2023-present NVIDIA CORPORATION & AFFILIATES.
+ * All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ */
+// clang-format on
+#include <device_lower/analysis/tensor_producer_aliases.h>
+#include <device_lower/lower2device.h>
+#include <ir/utils.h>
+#include <kernel_ir_dispatch.h>
+#include <type.h>
+
+#include <unordered_set>
+
+namespace nvfuser {
+
+namespace {
+
+bool isTrivialExpr(Expr* expr) {
+  TensorView* in = ir_utils::getTvInput(expr);
+  TensorView* out = ir_utils::getTvOutput(expr);
+  if (in == nullptr || out == nullptr ||
+      in->getMemoryType() != MemoryType::Global ||
+      out->getMemoryType() != MemoryType::Global || out->isFusionOutput() ||
+      !expr->isOneOf<BroadcastOp, LoadStoreOp, SqueezeOp>()) {
+    return false;
+  }
+  // This is a tensor op that does no computation. However, it may still be
+  // non-trivial if the allocation domains are not equivalent, as that would
+  // indicate data movement. We check that any non-Broadcast and non-Reduction
+  // domains in the two allocation domains are exact mapped and appear in the
+  // same order. We also check that the tensors are fully contiguous so that we
+  // can re-use the exact same linear index.
+  // TODO: support discontiguous inputs. The output should be an intermediate
+  // tensor which we would always assume to be contiguous, but the input might
+  // not be. This would necessitate using a different linear index.
+  size_t in_pos = 0;
+  size_t out_pos = 0;
+  const std::vector<IterDomain*>& in_alloc = in->getMaybeAllocationDomain();
+  const std::vector<IterDomain*>& out_alloc = out->getMaybeAllocationDomain();
+  const std::vector<std::optional<bool>>& in_contig = in->getContiguity();
+  const std::vector<std::optional<bool>>& out_contig = out->getContiguity();
+
+  const ValGraph& exact_graph =
+      GpuLower::current()->idModel().idGraph(IdMappingMode::EXACT);
+
+  while (true) {
path: csrc/device_lower/analysis/tensor_producer_aliases.cpp
user: login: naoyam
body: nit: this could be simplified by using `TensorDomain::noReductions` and `TensorDomain::noBroadcast`.
created_at: 2025-03-03T18:44:59Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,122 @@
+// clang-format off
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2023-present NVIDIA CORPORATION & AFFILIATES.
+ * All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ */
+// clang-format on
+#include <device_lower/analysis/tensor_producer_aliases.h>
+#include <device_lower/lower2device.h>
+#include <ir/utils.h>
+#include <kernel_ir_dispatch.h>
+#include <type.h>
+
+#include <unordered_set>
+
+namespace nvfuser {
+
+namespace {
+
+bool isTrivialExpr(Expr* expr) {
+  TensorView* in = ir_utils::getTvInput(expr);
+  TensorView* out = ir_utils::getTvOutput(expr);
+  if (in == nullptr || out == nullptr ||
+      in->getMemoryType() != MemoryType::Global ||
+      out->getMemoryType() != MemoryType::Global || out->isFusionOutput() ||
+      !expr->isOneOf<BroadcastOp, LoadStoreOp, SqueezeOp>()) {
+    return false;
+  }
+  // This is a tensor op that does no computation. However, it may still be
path: csrc/device_lower/analysis/tensor_producer_aliases.cpp
user: login: naoyam
body: We should probably also check DID parallelization as those domains are not actually allocated. `ir_utils::isMemoryPartitionedAcross` or `ir_utils::isMemorySharedAcross` should be used.
created_at: 2025-03-03T18:47:51Z
author_association: COLLABORATOR
diff_hunk: @@ -1715,6 +1715,13 @@ std::vector<Expr*> ExprSegmentationSorter::getExprs() const {
       if (!lower_utils::isScalarExpr(expr)) {
         active_exprs = &remaining_exprs;
       }
+      // Here we check whether this expression's output is a TensorView aliased
+      // to one of its producer tensors. If so, that indicates that this
+      // expression is trivial and should not be considered for codegen.
+      if (TensorView* tv_out = ir_utils::getTvOutput(expr); tv_out &&
+          GpuLower::current()->getTensorProducerAlias(tv_out) != nullptr) {
+        continue;
path: csrc/device_lower/pass/expr_sort.cpp
user: login: naoyam
body: Am I understanding correctly that this is where the removal of the trivial ops is done?
created_at: 2025-03-03T18:54:04Z
author_association: COLLABORATOR
diff_hunk: @@ -1715,6 +1715,13 @@ std::vector<Expr*> ExprSegmentationSorter::getExprs() const {
       if (!lower_utils::isScalarExpr(expr)) {
         active_exprs = &remaining_exprs;
       }
+      // Here we check whether this expression's output is a TensorView aliased
+      // to one of its producer tensors. If so, that indicates that this
+      // expression is trivial and should not be considered for codegen.
+      if (TensorView* tv_out = ir_utils::getTvOutput(expr); tv_out &&
+          GpuLower::current()->getTensorProducerAlias(tv_out) != nullptr) {
+        continue;
path: csrc/device_lower/pass/expr_sort.cpp
user: login: jacobhinkle
body: Yes, this causes the expression to not be included before the loop nests are generated for initialization and definition. Later changes are needed to reference to the producer alias in the actual expression instead of the aliased consumer but this guarantees we skip those definitions and their allocations.
created_at: 2025-03-03T21:47:09Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,122 @@
+// clang-format off
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2023-present NVIDIA CORPORATION & AFFILIATES.
+ * All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ */
+// clang-format on
+#include <device_lower/analysis/tensor_producer_aliases.h>
+#include <device_lower/lower2device.h>
+#include <ir/utils.h>
+#include <kernel_ir_dispatch.h>
+#include <type.h>
+
+#include <unordered_set>
+
+namespace nvfuser {
+
+namespace {
+
+bool isTrivialExpr(Expr* expr) {
+  TensorView* in = ir_utils::getTvInput(expr);
+  TensorView* out = ir_utils::getTvOutput(expr);
+  if (in == nullptr || out == nullptr ||
+      in->getMemoryType() != MemoryType::Global ||
+      out->getMemoryType() != MemoryType::Global || out->isFusionOutput() ||
+      !expr->isOneOf<BroadcastOp, LoadStoreOp, SqueezeOp>()) {
+    return false;
+  }
+  // This is a tensor op that does no computation. However, it may still be
path: csrc/device_lower/analysis/tensor_producer_aliases.cpp
user: login: jacobhinkle
body: Good point. I will check that for the promoted IDs.
created_at: 2025-03-03T21:50:04Z
author_association: COLLABORATOR
diff_hunk: @@ -1715,6 +1715,13 @@ std::vector<Expr*> ExprSegmentationSorter::getExprs() const {
       if (!lower_utils::isScalarExpr(expr)) {
         active_exprs = &remaining_exprs;
       }
+      // Here we check whether this expression's output is a TensorView aliased
+      // to one of its producer tensors. If so, that indicates that this
+      // expression is trivial and should not be considered for codegen.
+      if (TensorView* tv_out = ir_utils::getTvOutput(expr); tv_out &&
+          GpuLower::current()->getTensorProducerAlias(tv_out) != nullptr) {
+        continue;
path: csrc/device_lower/pass/expr_sort.cpp
user: login: naoyam
body: I'm a bit feeling uneasy that sorting expressions can mean dropping some expressions. Do you think this is the best place to do that?
created_at: 2025-03-04T02:25:00Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,131 @@
+// clang-format off
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2023-present NVIDIA CORPORATION & AFFILIATES.
+ * All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ */
+// clang-format on
+#include <device_lower/analysis/tensor_producer_aliases.h>
+#include <device_lower/lower2device.h>
+#include <device_lower/utils.h>
+#include <ir/utils.h>
+#include <kernel_ir_dispatch.h>
+#include <type.h>
+
+#include <unordered_set>
+
+namespace nvfuser {
+
+namespace {
+
+bool isTrivialExpr(Expr* expr) {
+  TensorView* in = ir_utils::getTvInput(expr);
+  TensorView* out = ir_utils::getTvOutput(expr);
+  if (in == nullptr || out == nullptr ||
+      in->getMemoryType() != MemoryType::Global ||
+      out->getMemoryType() != MemoryType::Global || out->isFusionOutput() ||
+      !expr->isOneOf<BroadcastOp, LoadStoreOp, SqueezeOp>()) {
+    return false;
+  }
+  // This is a tensor op that does no computation. However, it may still be
+  // non-trivial if the allocation domains are not equivalent, as that would
+  // indicate data movement. We check that any non-Broadcast and non-Reduction
+  // domains in the two allocation domains are exact mapped and appear in the
+  // same order. We also check that the tensors are fully contiguous so that we
+  // can re-use the exact same linear index.
+  // TODO: support discontiguous inputs. The output should be an intermediate
+  // tensor which we would always assume to be contiguous, but the input might
+  // not be. This would necessitate using a different linear index.
+  const std::vector<IterDomain*>& in_alloc = TensorDomain::noReductions(
+      TensorDomain::noBroadcasts(in->getMaybeAllocationDomain()));
+  const std::vector<IterDomain*>& out_alloc =
+      TensorDomain::noBroadcasts(out->getMaybeAllocationDomain());
+
+  if (in_alloc.size() != out_alloc.size()) {
+    // Non-trivial allocation domains cannot be in bijective correspondence if
+    // there are different numbers of them
+    return false;
+  }
+
+  std::vector<bool> in_contig;
+  std::vector<bool> out_contig;
+  for (const std::optional<bool>& c : in->getContiguity()) {
+    if (c.has_value()) {
+      in_contig.push_back(c.value());
+    }
+  }
+  for (const std::optional<bool>& c : out->getContiguity()) {
+    if (c.has_value()) {
+      out_contig.push_back(c.value());
+    }
+  }
+
+  const ValGraph& exact_graph =
+      GpuLower::current()->idModel().idGraph(IdMappingMode::EXACT);
+
+  for (size_t pos : c10::irange(in_alloc.size())) {
+    // At this point in_pos and out_pos are both in range and point to
+    // non-broadcast IDs
+    IterDomain* in_id = in_alloc.at(pos);
+    IterDomain* out_id = out_alloc.at(pos);
+
+    // If this allocation ID is parallelized such that its loop index is not
+    // used, then we can ignore it for this analysis.
+    const auto id_is_indexed = [](TensorView* tv, IterDomain* id) {
+      IterDomain* loop_id = lower_utils::getConcreteLoopID(id);
+      if (!loop_id->isParallelized()) {
+        return true;
+      }
+      if (loop_id->isDeviceDim()) {
+        return false;
+      }
+      return !(
+          loop_id->isThread() &&
path: csrc/device_lower/analysis/tensor_producer_aliases.cpp
user: login: naoyam
body: I think `loop_id->isThread()` is redundant.
created_at: 2025-03-04T05:38:17Z
author_association: COLLABORATOR
diff_hunk: @@ -433,6 +433,9 @@ class ReadAfterWriteSyncs : public kir::ExprMutator {
     // TODO: unify the handle of cp.async
     std::unordered_map<AsyncOpType, std::unordered_set<Expr*>> input_async_ops;
     for (auto inp : expr->inputs()) {
+      if (auto* inp_tv = dynamic_cast<TensorView*>(inp)) {
+        inp = GpuLower::current()->getMaybeTensorProducerAlias(inp_tv);
path: csrc/device_lower/pass/insert_syncs.cpp
user: login: naoyam
body: Why is this necessary? Is it because there's something that only the aliased tensor has but not the aliasing one?
created_at: 2025-03-04T05:45:53Z
author_association: COLLABORATOR
diff_hunk: @@ -97,6 +97,10 @@ void LoopNestGenerator::handle(Expr* expr) {
 
   TensorView* out_tv = expr->output(0)->as<TensorView>();
 
+  if (GpuLower::current()->getTensorProducerAlias(out_tv) != nullptr) {
path: csrc/device_lower/pass/loops.cpp
user: login: naoyam
body: Shouldn't this be already filtered out at the initial expression sorting?
created_at: 2025-03-04T05:52:04Z
author_association: COLLABORATOR
diff_hunk: @@ -757,6 +757,17 @@ class AllocationInfoMap : private kir::IrVisitor {
   using kir::IrVisitor::handle;
 
   void dispatch(Expr* expr) final {
+    if (expr->outputs().size() == 1) {
+      // Look for forced aliases
+      auto* out_tv = dynamic_cast<TensorView*>(expr->output(0));
+      if (out_tv) {
+        auto* alias_tv = GpuLower::current()->getTensorProducerAlias(out_tv);
+        if (alias_tv != nullptr) {
path: csrc/device_lower/pass/alias_memory.cpp
user: login: naoyam
body: Same as the loop pass case, shouldn't this be already filtered out?
created_at: 2025-03-04T06:01:33Z
author_association: COLLABORATOR
diff_hunk: @@ -1715,6 +1715,13 @@ std::vector<Expr*> ExprSegmentationSorter::getExprs() const {
       if (!lower_utils::isScalarExpr(expr)) {
         active_exprs = &remaining_exprs;
       }
+      // Here we check whether this expression's output is a TensorView aliased
+      // to one of its producer tensors. If so, that indicates that this
+      // expression is trivial and should not be considered for codegen.
+      if (TensorView* tv_out = ir_utils::getTvOutput(expr); tv_out &&
+          GpuLower::current()->getTensorProducerAlias(tv_out) != nullptr) {
+        continue;
path: csrc/device_lower/pass/expr_sort.cpp
user: login: jacobhinkle
body: Good point maybe we should do it as a separate pass just after sorting so that it's more explicit.
created_at: 2025-03-04T10:09:55Z
author_association: COLLABORATOR
diff_hunk: @@ -433,6 +433,9 @@ class ReadAfterWriteSyncs : public kir::ExprMutator {
     // TODO: unify the handle of cp.async
     std::unordered_map<AsyncOpType, std::unordered_set<Expr*>> input_async_ops;
     for (auto inp : expr->inputs()) {
+      if (auto* inp_tv = dynamic_cast<TensorView*>(inp)) {
+        inp = GpuLower::current()->getMaybeTensorProducerAlias(inp_tv);
path: csrc/device_lower/pass/insert_syncs.cpp
user: login: jacobhinkle
body: I will add a comment here. Example:
```
A
B = broadcast(A)  # aliased to A
C = foo(B)
```
C is actually going to be a read of A instead of B, so we don't need to sync between B and C.
created_at: 2025-03-04T15:06:35Z
author_association: COLLABORATOR
diff_hunk: @@ -97,6 +97,10 @@ void LoopNestGenerator::handle(Expr* expr) {
 
   TensorView* out_tv = expr->output(0)->as<TensorView>();
 
+  if (GpuLower::current()->getTensorProducerAlias(out_tv) != nullptr) {
path: csrc/device_lower/pass/loops.cpp
user: login: jacobhinkle
body: You are correct, I think this is no longer needed. 
created_at: 2025-03-04T15:08:12Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,131 @@
+// clang-format off
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2023-present NVIDIA CORPORATION & AFFILIATES.
+ * All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ */
+// clang-format on
+#include <device_lower/analysis/tensor_producer_aliases.h>
+#include <device_lower/lower2device.h>
+#include <device_lower/utils.h>
+#include <ir/utils.h>
+#include <kernel_ir_dispatch.h>
+#include <type.h>
+
+#include <unordered_set>
+
+namespace nvfuser {
+
+namespace {
+
+bool isTrivialExpr(Expr* expr) {
+  TensorView* in = ir_utils::getTvInput(expr);
+  TensorView* out = ir_utils::getTvOutput(expr);
+  if (in == nullptr || out == nullptr ||
+      in->getMemoryType() != MemoryType::Global ||
+      out->getMemoryType() != MemoryType::Global || out->isFusionOutput() ||
+      !expr->isOneOf<BroadcastOp, LoadStoreOp, SqueezeOp>()) {
+    return false;
+  }
+  // This is a tensor op that does no computation. However, it may still be
+  // non-trivial if the allocation domains are not equivalent, as that would
+  // indicate data movement. We check that any non-Broadcast and non-Reduction
+  // domains in the two allocation domains are exact mapped and appear in the
+  // same order. We also check that the tensors are fully contiguous so that we
+  // can re-use the exact same linear index.
+  // TODO: support discontiguous inputs. The output should be an intermediate
+  // tensor which we would always assume to be contiguous, but the input might
+  // not be. This would necessitate using a different linear index.
+  const std::vector<IterDomain*>& in_alloc = TensorDomain::noReductions(
+      TensorDomain::noBroadcasts(in->getMaybeAllocationDomain()));
+  const std::vector<IterDomain*>& out_alloc =
+      TensorDomain::noBroadcasts(out->getMaybeAllocationDomain());
+
+  if (in_alloc.size() != out_alloc.size()) {
+    // Non-trivial allocation domains cannot be in bijective correspondence if
+    // there are different numbers of them
+    return false;
+  }
+
+  std::vector<bool> in_contig;
+  std::vector<bool> out_contig;
+  for (const std::optional<bool>& c : in->getContiguity()) {
+    if (c.has_value()) {
+      in_contig.push_back(c.value());
+    }
+  }
+  for (const std::optional<bool>& c : out->getContiguity()) {
+    if (c.has_value()) {
+      out_contig.push_back(c.value());
+    }
+  }
+
+  const ValGraph& exact_graph =
+      GpuLower::current()->idModel().idGraph(IdMappingMode::EXACT);
+
+  for (size_t pos : c10::irange(in_alloc.size())) {
+    // At this point in_pos and out_pos are both in range and point to
+    // non-broadcast IDs
+    IterDomain* in_id = in_alloc.at(pos);
+    IterDomain* out_id = out_alloc.at(pos);
+
+    // If this allocation ID is parallelized such that its loop index is not
+    // used, then we can ignore it for this analysis.
+    const auto id_is_indexed = [](TensorView* tv, IterDomain* id) {
+      IterDomain* loop_id = lower_utils::getConcreteLoopID(id);
+      if (!loop_id->isParallelized()) {
+        return true;
+      }
+      if (loop_id->isDeviceDim()) {
+        return false;
+      }
+      return !(
+          loop_id->isThread() &&
path: csrc/device_lower/analysis/tensor_producer_aliases.cpp
user: login: jacobhinkle
body: You're right. It's a little confusing. What I really want is "is this ID actually allocated". The `isMemorySharedAcross` check gives me that when the ID is parallelized but not for other ParallelTypes other than DID, BID, and TID. That is why I have the `!loop_id->isParallelized()` line above. I'll remove the isThread check.
created_at: 2025-03-04T15:25:08Z
author_association: COLLABORATOR
diff_hunk: @@ -1715,6 +1715,13 @@ std::vector<Expr*> ExprSegmentationSorter::getExprs() const {
       if (!lower_utils::isScalarExpr(expr)) {
         active_exprs = &remaining_exprs;
       }
+      // Here we check whether this expression's output is a TensorView aliased
+      // to one of its producer tensors. If so, that indicates that this
+      // expression is trivial and should not be considered for codegen.
+      if (TensorView* tv_out = ir_utils::getTvOutput(expr); tv_out &&
+          GpuLower::current()->getTensorProducerAlias(tv_out) != nullptr) {
+        continue;
path: csrc/device_lower/pass/expr_sort.cpp
user: login: jacobhinkle
body: I made it a pass. It is now the first pass run (right before `LoopnestGenerator`).
created_at: 2025-03-04T15:26:20Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,131 @@
+// clang-format off
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2023-present NVIDIA CORPORATION & AFFILIATES.
+ * All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ */
+// clang-format on
+#include <device_lower/analysis/tensor_producer_aliases.h>
+#include <device_lower/lower2device.h>
+#include <device_lower/utils.h>
+#include <ir/utils.h>
+#include <kernel_ir_dispatch.h>
+#include <type.h>
+
+#include <unordered_set>
+
+namespace nvfuser {
+
+namespace {
+
+bool isTrivialExpr(Expr* expr) {
+  TensorView* in = ir_utils::getTvInput(expr);
+  TensorView* out = ir_utils::getTvOutput(expr);
+  if (in == nullptr || out == nullptr ||
+      in->getMemoryType() != MemoryType::Global ||
+      out->getMemoryType() != MemoryType::Global || out->isFusionOutput() ||
+      !expr->isOneOf<BroadcastOp, LoadStoreOp, SqueezeOp>()) {
+    return false;
+  }
+  // This is a tensor op that does no computation. However, it may still be
+  // non-trivial if the allocation domains are not equivalent, as that would
+  // indicate data movement. We check that any non-Broadcast and non-Reduction
+  // domains in the two allocation domains are exact mapped and appear in the
+  // same order. We also check that the tensors are fully contiguous so that we
+  // can re-use the exact same linear index.
+  // TODO: support discontiguous inputs. The output should be an intermediate
+  // tensor which we would always assume to be contiguous, but the input might
+  // not be. This would necessitate using a different linear index.
+  const std::vector<IterDomain*>& in_alloc = TensorDomain::noReductions(
+      TensorDomain::noBroadcasts(in->getMaybeAllocationDomain()));
+  const std::vector<IterDomain*>& out_alloc =
+      TensorDomain::noBroadcasts(out->getMaybeAllocationDomain());
+
+  if (in_alloc.size() != out_alloc.size()) {
+    // Non-trivial allocation domains cannot be in bijective correspondence if
+    // there are different numbers of them
+    return false;
+  }
+
+  std::vector<bool> in_contig;
+  std::vector<bool> out_contig;
+  for (const std::optional<bool>& c : in->getContiguity()) {
+    if (c.has_value()) {
+      in_contig.push_back(c.value());
+    }
+  }
+  for (const std::optional<bool>& c : out->getContiguity()) {
+    if (c.has_value()) {
+      out_contig.push_back(c.value());
+    }
+  }
+
+  const ValGraph& exact_graph =
+      GpuLower::current()->idModel().idGraph(IdMappingMode::EXACT);
+
+  for (size_t pos : c10::irange(in_alloc.size())) {
+    // At this point in_pos and out_pos are both in range and point to
+    // non-broadcast IDs
+    IterDomain* in_id = in_alloc.at(pos);
+    IterDomain* out_id = out_alloc.at(pos);
+
+    // If this allocation ID is parallelized such that its loop index is not
+    // used, then we can ignore it for this analysis.
+    const auto id_is_indexed = [](TensorView* tv, IterDomain* id) {
+      IterDomain* loop_id = lower_utils::getConcreteLoopID(id);
+      if (!loop_id->isParallelized()) {
+        return true;
+      }
+      if (loop_id->isDeviceDim()) {
+        return false;
+      }
+      return !(
+          loop_id->isThread() &&
path: csrc/device_lower/analysis/tensor_producer_aliases.cpp
user: login: naoyam
body: Then,`isMemoryPartitionedAcross` may make more sense. Partitioning means no allocation.
created_at: 2025-03-04T16:51:57Z
author_association: COLLABORATOR
diff_hunk: @@ -433,6 +433,9 @@ class ReadAfterWriteSyncs : public kir::ExprMutator {
     // TODO: unify the handle of cp.async
     std::unordered_map<AsyncOpType, std::unordered_set<Expr*>> input_async_ops;
     for (auto inp : expr->inputs()) {
+      if (auto* inp_tv = dynamic_cast<TensorView*>(inp)) {
+        inp = GpuLower::current()->getMaybeTensorProducerAlias(inp_tv);
path: csrc/device_lower/pass/insert_syncs.cpp
user: login: naoyam
body: Let's say we don't do this and end up requiring a sync between B and C. Would the lowered exprs look like below?

```
A
sync
C
```

Here, I assume `A` is an input, so this sync should not be necessary. Am I following correctly?

I'm asking this because this seems like a yet another cognitive burden since we would need to remember sometimes we would need to follow the alias connection. 

I'm wondering if somehow `expr->inputs()` return aliased tensors, things might be easier to deal with?
created_at: 2025-03-04T17:04:42Z
author_association: COLLABORATOR
diff_hunk: @@ -433,6 +433,9 @@ class ReadAfterWriteSyncs : public kir::ExprMutator {
     // TODO: unify the handle of cp.async
     std::unordered_map<AsyncOpType, std::unordered_set<Expr*>> input_async_ops;
     for (auto inp : expr->inputs()) {
+      if (auto* inp_tv = dynamic_cast<TensorView*>(inp)) {
+        inp = GpuLower::current()->getMaybeTensorProducerAlias(inp_tv);
path: csrc/device_lower/pass/insert_syncs.cpp
user: login: jacobhinkle
body: Yes that is what it would look like. This actually is checked in one of our tests: `FusionValidateParallelize7_CUDA`. The fusion is scheduled like this:
```
Inputs:
  T0_g_float[iS0{i0}, iS1{i2}]
Outputs:
  T2_g_float[iS4{i0}, ithreadIdx.y5{i2}]
  T3_g_float[iblockIdx.x6{i0}, iS7{i2}]

%kernel {
T1_g_float[iblockIdx.x2{i0}, ithreadIdx.x3{i2}]
   = Set( T0_g_float[iS0{i0}, iS1{i2}], cache_op=Streaming )
T2_g_float[iS4{i0}, ithreadIdx.y5{i2}]
   = Set( T1_g_float[iblockIdx.x2{i0}, ithreadIdx.x3{i2}], cache_op=Streaming )
T3_g_float[iblockIdx.x6{i0}, iS7{i2}]
   = Set( T1_g_float[iblockIdx.x2{i0}, ithreadIdx.x3{i2}], cache_op=Streaming )
```
The test in that case checks that there is a GridSync in the fusion, but now we alias `T1` to `T0` so no syncs are necessary.

> I'm wondering if somehow expr->inputs() return aliased tensors, things might be easier to deal with?

I did try something like this previously, and it might be workable. One issue with that is that we will get errors after such replacement if PairwiseLogicalDomainMap is used, since the producer will have different number of logical dims than the consumer.

I don't think there's a major problem with just removing this change to the RAW sync insertion pass to make the change a little less invasive. In practice I don't expect us to be doing this very often so it should not be a major concern when writing lowering passes.
created_at: 2025-03-04T17:22:48Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,131 @@
+// clang-format off
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2023-present NVIDIA CORPORATION & AFFILIATES.
+ * All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ */
+// clang-format on
+#include <device_lower/analysis/tensor_producer_aliases.h>
+#include <device_lower/lower2device.h>
+#include <device_lower/utils.h>
+#include <ir/utils.h>
+#include <kernel_ir_dispatch.h>
+#include <type.h>
+
+#include <unordered_set>
+
+namespace nvfuser {
+
+namespace {
+
+bool isTrivialExpr(Expr* expr) {
+  TensorView* in = ir_utils::getTvInput(expr);
+  TensorView* out = ir_utils::getTvOutput(expr);
+  if (in == nullptr || out == nullptr ||
+      in->getMemoryType() != MemoryType::Global ||
+      out->getMemoryType() != MemoryType::Global || out->isFusionOutput() ||
+      !expr->isOneOf<BroadcastOp, LoadStoreOp, SqueezeOp>()) {
+    return false;
+  }
+  // This is a tensor op that does no computation. However, it may still be
+  // non-trivial if the allocation domains are not equivalent, as that would
+  // indicate data movement. We check that any non-Broadcast and non-Reduction
+  // domains in the two allocation domains are exact mapped and appear in the
+  // same order. We also check that the tensors are fully contiguous so that we
+  // can re-use the exact same linear index.
+  // TODO: support discontiguous inputs. The output should be an intermediate
+  // tensor which we would always assume to be contiguous, but the input might
+  // not be. This would necessitate using a different linear index.
+  const std::vector<IterDomain*>& in_alloc = TensorDomain::noReductions(
+      TensorDomain::noBroadcasts(in->getMaybeAllocationDomain()));
+  const std::vector<IterDomain*>& out_alloc =
+      TensorDomain::noBroadcasts(out->getMaybeAllocationDomain());
+
+  if (in_alloc.size() != out_alloc.size()) {
+    // Non-trivial allocation domains cannot be in bijective correspondence if
+    // there are different numbers of them
+    return false;
+  }
+
+  std::vector<bool> in_contig;
+  std::vector<bool> out_contig;
+  for (const std::optional<bool>& c : in->getContiguity()) {
+    if (c.has_value()) {
+      in_contig.push_back(c.value());
+    }
+  }
+  for (const std::optional<bool>& c : out->getContiguity()) {
+    if (c.has_value()) {
+      out_contig.push_back(c.value());
+    }
+  }
+
+  const ValGraph& exact_graph =
+      GpuLower::current()->idModel().idGraph(IdMappingMode::EXACT);
+
+  for (size_t pos : c10::irange(in_alloc.size())) {
+    // At this point in_pos and out_pos are both in range and point to
+    // non-broadcast IDs
+    IterDomain* in_id = in_alloc.at(pos);
+    IterDomain* out_id = out_alloc.at(pos);
+
+    // If this allocation ID is parallelized such that its loop index is not
+    // used, then we can ignore it for this analysis.
+    const auto id_is_indexed = [](TensorView* tv, IterDomain* id) {
+      IterDomain* loop_id = lower_utils::getConcreteLoopID(id);
+      if (!loop_id->isParallelized()) {
+        return true;
+      }
+      if (loop_id->isDeviceDim()) {
+        return false;
+      }
+      return !(
+          loop_id->isThread() &&
path: csrc/device_lower/analysis/tensor_producer_aliases.cpp
user: login: jacobhinkle
body: Of course. Thank you! I started there and talked myself out of using that. I'll use it and put a comment there.
created_at: 2025-03-04T20:46:08Z
author_association: COLLABORATOR
diff_hunk: @@ -433,6 +433,9 @@ class ReadAfterWriteSyncs : public kir::ExprMutator {
     // TODO: unify the handle of cp.async
     std::unordered_map<AsyncOpType, std::unordered_set<Expr*>> input_async_ops;
     for (auto inp : expr->inputs()) {
+      if (auto* inp_tv = dynamic_cast<TensorView*>(inp)) {
+        inp = GpuLower::current()->getMaybeTensorProducerAlias(inp_tv);
path: csrc/device_lower/pass/insert_syncs.cpp
user: login: naoyam
body: > In practice I don't expect us to be doing this very often so it should not be a major concern when writing lowering passes.

I don't disagree, but let me think about it a little more. In this particular case, if we changed `Val::isFusionInput` to return true for `T1` (we could override it in KIR such that it follows the alias connection), I think no sync would be inserted. But more broadly, if we had something like:

```
T1 = ...; // intermediate global tensor
T2 = op1(T1);
T3 = set(T1); // will alias to T1
T4 = op2(T3);
```

Suppose, without aliasing, two grid sync were necessary between T1 and T2 and between T3 and T4. With aliasing, since T1 would be already coherent by the sync for T2, the second sync would not be necessary. In this case, `T3->isFusionInput` would not return true since `T1` is not a fusion input, so the above trick wouldn't help. Or, our sync insertion algorithm would be smart enough to consolidate two syncs into one?

In any case, I'd like to keep something like this in mind, but this shouldn't be a blocker.

created_at: 2025-03-04T22:54:47Z
author_association: COLLABORATOR

PR comments:
user: login: jacobhinkle
created_at: 2025-02-26T00:46:19Z
body: !test
author_association: COLLABORATOR
user: login: github-actions[bot]
created_at: 2025-02-26T00:47:01Z
body: <!-- llm action combined persistent comment, dont modify this line -->

*Review updated until commit https://github.com/NVIDIA/Fuser/commit/f4ff0dc5cfe90ce3ab604867cab211f566dff65e*


### **Description**
- Introduce analysis to identify and alias trivial global-to-global operations.

- Implement removal of trivial operations during lowering.

- Add tests for various trivial operations (set, broadcast, squeeze, permute).


___



### **Changes walkthrough** 📝
<table><thead><tr><th></th><th align="left">Relevant files</th></tr></thead><tbody><tr><td><strong>Enhancement</strong></td><td><details><summary>7 files</summary><table>
<tr>
  <td><strong>tensor_producer_aliases.cpp</strong><dd><code>Add analysis for trivial global-to-global operations</code>&nbsp; &nbsp; &nbsp; &nbsp; &nbsp; </dd></td>
  <td><a href="https://github.com/NVIDIA/Fuser/pull/3968/files#diff-6917cab257e5a109901a933f5429997edfb9b194bd16d54df4853725a3e6cfeb">+145/-0</a>&nbsp; </td>

</tr>

<tr>
  <td><strong>lower2device.cpp</strong><dd><code>Integrate tensor producer aliasing in lowering</code>&nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; </dd></td>
  <td><a href="https://github.com/NVIDIA/Fuser/pull/3968/files#diff-79930932c0557f58540f04c0c1930922e9f3b4a73965cf7e465ed7fa7fdc1184">+32/-1</a>&nbsp; &nbsp; </td>

</tr>

<tr>
  <td><strong>index.cpp</strong><dd><code>Update index lowering to respect tensor producer aliases</code>&nbsp; </dd></td>
  <td><a href="https://github.com/NVIDIA/Fuser/pull/3968/files#diff-60e927bdf6f0ab7d9206208a2515d01c1a0b2a355253224d8ac63a4a32a4aa69">+9/-1</a>&nbsp; &nbsp; &nbsp; </td>

</tr>

<tr>
  <td><strong>insert_syncs.cpp</strong><dd><code>Respect tensor producer aliases in sync insertion</code>&nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; </dd></td>
  <td><a href="https://github.com/NVIDIA/Fuser/pull/3968/files#diff-49bea61a8cde014ec0396c89d0654813136e0bd3ffe8b3a5974ee9ccf3a5fbb8">+5/-0</a>&nbsp; &nbsp; &nbsp; </td>

</tr>

<tr>
  <td><strong>indexing.cpp</strong><dd><code>Update allocation domain setup to respect tensor producer aliases</code></dd></td>
  <td><a href="https://github.com/NVIDIA/Fuser/pull/3968/files#diff-0c47d4d562dba268498573ced9dd7fc7d9638b5882aca4c4cc20e70a7ccec891">+6/-3</a>&nbsp; &nbsp; &nbsp; </td>

</tr>

<tr>
  <td><strong>tensor_producer_aliases.h</strong><dd><code>Add header for tensor producer alias analysis</code>&nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; </dd></td>
  <td><a href="https://github.com/NVIDIA/Fuser/pull/3968/files#diff-abf47eb7e040d0a84fedbd514d168bf0e0e05e5248e88cf49b7d41024b080178">+57/-0</a>&nbsp; &nbsp; </td>

</tr>

<tr>
  <td><strong>lower2device.h</strong><dd><code>Add methods for tensor producer aliasing</code>&nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; </dd></td>
  <td><a href="https://github.com/NVIDIA/Fuser/pull/3968/files#diff-31461d2f00df5b359cda3aac1bf66a9ac3a0e83960d14d99a5e305396c6f4fa3">+46/-0</a>&nbsp; &nbsp; </td>

</tr>
</table></details></td></tr><tr><td><strong>Tests</strong></td><td><details><summary>3 files</summary><table>
<tr>
  <td><strong>test_gpu2.cpp</strong><dd><code>Update test to reflect changes in trivial operations</code>&nbsp; &nbsp; &nbsp; &nbsp; &nbsp; </dd></td>
  <td><a href="https://github.com/NVIDIA/Fuser/pull/3968/files#diff-9f148c6abf4059e23d7b0bb5f76736b320a711f920e899836e15cfa2e5ece398">+5/-2</a>&nbsp; &nbsp; &nbsp; </td>

</tr>

<tr>
  <td><strong>test_matmul.cpp</strong><dd><code>Add test for broadcast operation removal</code>&nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; </dd></td>
  <td><a href="https://github.com/NVIDIA/Fuser/pull/3968/files#diff-64fc4e7bfbc5b9f95ac3dc5823bd99b683b048926805c13310ce6a8ef8032289">+202/-0</a>&nbsp; </td>

</tr>

<tr>
  <td><strong>test_remove_trivial_ops.cpp</strong><dd><code>Add comprehensive tests for trivial operation removal</code>&nbsp; &nbsp; &nbsp; &nbsp; </dd></td>
  <td><a href="https://github.com/NVIDIA/Fuser/pull/3968/files#diff-c6ceccea0eb68167e57534c309f6cb059b6e1539f0721512e74d5e9e7b6af5d8">+465/-0</a>&nbsp; </td>

</tr>
</table></details></td></tr><tr><td><strong>Configuration changes</strong></td><td><details><summary>1 files</summary><table>
<tr>
  <td><strong>CMakeLists.txt</strong><dd><code>Add new source and test files</code>&nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; </dd></td>
  <td><a href="https://github.com/NVIDIA/Fuser/pull/3968/files#diff-1e7de1ae2d059d21e1dd75d5812d5a34b0222cef273b7c3a2af62eb747f9d20a">+2/-0</a>&nbsp; &nbsp; &nbsp; </td>

</tr>
</table></details></td></tr></tr></tbody></table>

---

### PR Reviewer Guide 🔍

Here are some key observations to aid the review process:

<table>
<tr><td>🧪&nbsp;<strong>PR contains tests</strong></td></tr>
<tr><td>⚡&nbsp;<strong>Recommended focus areas for review</strong><br><br>

<details><summary><a href='https://github.com/NVIDIA/Fuser/pull/3968/files#diff-6917cab257e5a109901a933f5429997edfb9b194bd16d54df4853725a3e6cfebR41-R65'><strong>Contiguity Check</strong></a>

The current implementation assumes that both input and output tensors are fully contiguous. This might be too restrictive and could be relaxed in the future. Consider adding a comment or a TODO to revisit this assumption.</summary>

```c++
// tensor which we would always assume to be contiguous, but the input might
// not be. This would necessitate using a different linear index.
const std::vector<IterDomain*>& in_alloc = TensorDomain::noReductions(
    TensorDomain::noBroadcasts(in->getMaybeAllocationDomain()));
const std::vector<IterDomain*>& out_alloc =
    TensorDomain::noBroadcasts(out->getMaybeAllocationDomain());

if (in_alloc.size() != out_alloc.size()) {
  // Non-trivial allocation domains cannot be in bijective correspondence if
  // there are different numbers of them
  return false;
}

std::vector<bool> in_contig;
std::vector<bool> out_contig;
for (const std::optional<bool>& c : in->getContiguity()) {
  if (c.has_value()) {
    in_contig.push_back(c.value());
  }
}
for (const std::optional<bool>& c : out->getContiguity()) {
  if (c.has_value()) {
    out_contig.push_back(c.value());
  }
}
```

</details>

<details><summary><a href='https://github.com/NVIDIA/Fuser/pull/3968/files#diff-60e927bdf6f0ab7d9206208a2515d01c1a0b2a355253224d8ac63a4a32a4aa69R49-R63'><strong>Index Lowering</strong></a>

The modification in `IndexLowering::lowerSrcIndex` introduces a new `kir::TensorIndex` creation. Ensure that this does not introduce any performance regressions or correctness issues, especially when dealing with aliased tensors.</summary>

```c++
kir::TensorIndex* tind = Index::getProducerIndex(
    tv,
    dst->as<TensorView>(),
    for_loops_,
    getRotatedLoop(),
    override_index,
    generate_pointer,
    as_type);
if (TensorView* aliased_producer =
        GpuLower::current()->getTensorProducerAlias(tv)) {
  return IrBuilder::create<kir::TensorIndex>(
      aliased_producer, tind->index());
} else {
  return tind;
}
```

</details>

<details><summary><a href='https://github.com/NVIDIA/Fuser/pull/3968/files#diff-c6ceccea0eb68167e57534c309f6cb059b6e1539f0721512e74d5e9e7b6af5d8R1-R465'><strong>Test Coverage</strong></a>

The new tests cover various scenarios but might not cover all edge cases, such as nested trivial operations or more complex memory layouts. Consider adding additional test cases to ensure comprehensive coverage.</summary>

```c++
// clang-format off
/*
 * SPDX-FileCopyrightText: Copyright (c) 2023-present NVIDIA CORPORATION & AFFILIATES.
 * All rights reserved.
 * SPDX-License-Identifier: BSD-3-Clause
 */
// clang-format on
#include <csrc/exceptions.h>
#include <fusion.h>
#include <kernel_ir_dispatch.h>
#include <ops/all_ops.h>
#include <runtime/executor.h>
#include <tests/cpp/utils.h>
#include <tests/cpp/validator.h>

#include <gtest/gtest.h>

namespace nvfuser {

using RemoveTrivialOpsTest = NVFuserTest;

template <typename OpType>
class KernelContainsExpr : kir::IrVisitor {
 public:
  static bool check(const kir::Kernel* kernel) {
    KernelContainsExpr checker;
    checker.handle(kernel->topLevelExprs());
    return checker.has_op_;
  }

 private:
  using kir::IrVisitor::handle;

  void handle(OpType* op) final {
    has_op_ = true;
  }

 private:
  bool has_op_ = false;
};

// Test that we remove a trivial gmem->gmem set at lowering unless it is an
// alias.
TEST_F(RemoveTrivialOpsTest, Set) {
  Fusion fusion;
  FusionGuard fg(&fusion);

  auto tv0 = makeContigConcreteTensor({2, 3});
  fusion.addInput(tv0);
  auto tv1 = set(tv0);
  auto tv2 = neg(tv1);
  fusion.addOutput(tv2);

  auto options = at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0);
  auto t0 = at::randn({2, 3}, options);
  std::vector<c10::IValue> inputs{t0};

  {
    // In this case we do not remove the squeeze since it is G->L not G->G
    KernelExecutor ke;
    ke.compile(&fusion, inputs);
    const kir::Kernel* kernel = ke.compiledKernel()->kernel();
    EXPECT_TRUE(KernelContainsExpr<LoadStoreOp>::check(kernel));
    auto outputs = ke.run(inputs);
    testValidate(&fusion, outputs, inputs, __LINE__, __FILE__);
  }

  {
    // Setting tv1 to Global means the LoadStoreOp is now G->G. This means that
    // when we lower it, we will be able to safely remove it
    tv1->setMemoryType(MemoryType::Global);

    KernelExecutor ke;
    ke.compile(&fusion, inputs);
    const kir::Kernel* kernel = ke.compiledKernel()->kernel();
    EXPECT_FALSE(KernelContainsExpr<LoadStoreOp>::check(kernel));
    EXPECT_EQ(kernel->summary().global_allocations.size(), 0)
        << "Expected to have no intermediate global allocations";
    auto outputs = ke.run(inputs);
    testValidate(&fusion, outputs, inputs, __LINE__, __FILE__);
  }

  {
    NVF_ERROR(tv1->getMemoryType() == MemoryType::Global);
    // Transpose the input's allocation domain
    tv0->setAllocationDomain(
        {
            tv0->getLogicalDomain().at(1),
            tv0->getLogicalDomain().at(0),
        },
        true);

    KernelExecutor ke;

    std::vector<c10::IValue> transposed_inputs{at::randn({3, 2}, options).t()};
    ke.compile(&fusion, transposed_inputs);
    // We do not remove the set because the allocation domains do not match
    const kir::Kernel* kernel = ke.compiledKernel()->kernel();
    EXPECT_TRUE(KernelContainsExpr<LoadStoreOp>::check(kernel));
    EXPECT_EQ(kernel->summary().global_allocations.size(), 1)
        << "Expected to have one intermediate global allocation";
    auto outputs = ke.run(transposed_inputs);
    testValidate(&fusion, outputs, transposed_inputs, __LINE__, __FILE__);
  }

  {
    // Reset tv0's allocation domain but add a new output which reuses the
    // buffer of the input tv0
    tv0->setAllocationDomain(tv0->getLogicalDomain(), true);
    auto* tv3 = exp(tv2);
    fusion.addOutput(tv3);
    fusion.aliasOutputToInput(tv3, tv0, AllocationType::ReuseBuffer);

    KernelExecutor ke;
    ke.compile(&fusion, inputs);
    // We do not remove the squeeze because the input is target of an io alias
    const kir::Kernel* kernel = ke.compiledKernel()->kernel();
    EXPECT_TRUE(KernelContainsExpr<LoadStoreOp>::check(kernel));
    // Run with copy of inputs since we will overwrite the original inputs,
    // complicating validation
    std::vector<c10::IValue> inputs_copy{t0.clone()};
    auto outputs = ke.run(inputs_copy);
    testValidate(&fusion, outputs, inputs, __LINE__, __FILE__);
  }
}

// Test that we remove a trivial gmem->gmem broadcast at lowering
TEST_F(RemoveTrivialOpsTest, Broadcast) {
  Fusion fusion;
  FusionGuard fg(&fusion);

  auto tv0 = makeContigConcreteTensor({2, 3});
  fusion.addInput(tv0);
  auto tv1 = broadcast(tv0, {false, true, false});
  auto tv2 = neg(tv1);
  fusion.addOutput(tv2);

  auto options = at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0);
  auto t0 = at::randn({2, 3}, options);
  std::vector<c10::IValue> inputs{t0};

  {
    // In this case we do not remove the broadcast since it is G->L not G->G
    KernelExecutor ke;
    ke.compile(&fusion, inputs);
    const kir::Kernel* kernel = ke.compiledKernel()->kernel();
    EXPECT_TRUE(KernelContainsExpr<BroadcastOp>::check(kernel));
    auto outputs = ke.run(inputs);
    testValidate(&fusion, outputs, inputs, __LINE__, __FILE__);
  }

  {
    // Setting tv1 to Global means the BroadcastOp is now G->G. This means that
    // when we lower it, we will be able to safely remove it
    tv1->setMemoryType(MemoryType::Global);

    KernelExecutor ke;
    ke.compile(&fusion, inputs);
    const kir::Kernel* kernel = ke.compiledKernel()->kernel();
    EXPECT_FALSE(KernelContainsExpr<BroadcastOp>::check(kernel));
    EXPECT_EQ(kernel->summary().global_allocations.size(), 0)
        << "Expected to have no intermediate global allocations";
    auto outputs = ke.run(inputs);
    testValidate(&fusion, outputs, inputs, __LINE__, __FILE__);
  }

  {
    NVF_ERROR(tv1->getMemoryType() == MemoryType::Global);
    // Transpose the input's allocation domain
    tv0->setAllocationDomain(
        {
            tv0->getLogicalDomain().at(1),
            tv0->getLogicalDomain().at(0),
        },
        true);

    KernelExecutor ke;

    std::vector<c10::IValue> transposed_inputs{at::randn({3, 2}, options).t()};
    ke.compile(&fusion, transposed_inputs);
    // We do not remove the broadcast because the allocation domains do not
    // match
    const kir::Kernel* kernel = ke.compiledKernel()->kernel();
    EXPECT_TRUE(KernelContainsExpr<BroadcastOp>::check(kernel));
    EXPECT_EQ(kernel->summary().global_allocations.size(), 1)
        << "Expected to have one intermediate global allocation";
    auto outputs = ke.run(transposed_inputs);
    testValidate(&fusion, outputs, transposed_inputs, __LINE__, __FILE__);
  }

  {
    // Reset tv0's allocation domain but set tv1 as output
    tv0->setAllocationDomain(tv0->getLogicalDomain(), true);
    fusion.addOutput(tv1);

    KernelExecutor ke;
    ke.compile(&fusion, inputs);
    // We do not remove the broadcast because the output is a fusion output
    const kir::Kernel* kernel = ke.compiledKernel()->kernel();
    EXPECT_TRUE(KernelContainsExpr<BroadcastOp>::check(kernel));
    auto outputs = ke.run(inputs);
    testValidate(&fusion, outputs, inputs, __LINE__, __FILE__);
  }
}

// Test that we remove a trivial gmem->gmem squeeze at lowering
TEST_F(RemoveTrivialOpsTest, Squeeze) {
  Fusion fusion;
  FusionGuard fg(&fusion);

  auto tv0 = makeContigConcreteTensor({2, 1, 3});
  fusion.addInput(tv0);
  auto tv1 = squeeze(tv0, {1});
  auto tv2 = neg(tv1);
  fusion.addOutput(tv2);

  auto options = at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0);
  auto t0 = at::randn({2, 1, 3}, options);
  std::vector<c10::IValue> inputs{t0};

  {
    // In this case we do not remove the squeeze since it is G->L not G->G
    KernelExecutor ke;
    ke.compile(&fusion, inputs);
    const kir::Kernel* kernel = ke.compiledKernel()->kernel();
    // NOTE: SqueezeOp is converted to LoadStoreOp by the loadStoreOpInserter
    // lowering pass
    EXPECT_TRUE(KernelContainsExpr<LoadStoreOp>::check(kernel));
    auto outputs = ke.run(inputs);
    testValidate(&fusion, outputs, inputs, __LINE__, __FILE__);
  }

  {
    // Setting tv1 to Global means the SqueezeOp is now G->G. This means that
    // when we lower it, we will be able to safely remove it
    tv1->setMemoryType(MemoryType::Global);

    KernelExecutor ke;
    ke.compile(&fusion, inputs);
    const kir::Kernel* kernel = ke.compiledKernel()->kernel();
    EXPECT_FALSE(KernelContainsExpr<LoadStoreOp>::check(kernel));
    EXPECT_EQ(kernel->summary().global_allocations.size(), 0)
        << "Expected to have no intermediate global allocations";
    auto outputs = ke.run(inputs);
    testValidate(&fusion, outputs, inputs, __LINE__, __FILE__);
  }

  {
    NVF_ERROR(tv1->getMemoryType() == MemoryType::Global);
    // Transpose the input's allocation domain
    tv0->setAllocationDomain(
        {
            tv0->getLogicalDomain().at(2),
            tv0->getLogicalDomain().at(1),
            tv0->getLogicalDomain().at(0),
        },
        true);

    KernelExecutor ke;

    std::vector<c10::IValue> transposed_inputs{
        at::randn({3, 2}, options).t().unsqueeze(1)};
    ke.compile(&fusion, transposed_inputs);
    // We do not remove the broadcast because the allocation domains do not
    // match
    const kir::Kernel* kernel = ke.compiledKernel()->kernel();
    EXPECT_TRUE(KernelContainsExpr<LoadStoreOp>::check(kernel));
    EXPECT_EQ(kernel->summary().global_allocations.size(), 1)
        << "Expected to have one intermediate global allocation";
    auto outputs = ke.run(transposed_inputs);
    testValidate(&fusion, outputs, transposed_inputs, __LINE__, __FILE__);
  }

  {
    // Reset tv0's allocation domain but set tv1 as output
    tv0->setAllocationDomain(tv0->getLogicalDomain(), true);
    fusion.addOutput(tv1);

    KernelExecutor ke;
    ke.compile(&fusion, inputs);
    // We do not remove the squeeze because the output is a fusion output
    const kir::Kernel* kernel = ke.compiledKernel()->kernel();
    EXPECT_TRUE(KernelContainsExpr<LoadStoreOp>::check(kernel));
    auto outputs = ke.run(inputs);
    testValidate(&fusion, outputs, inputs, __LINE__, __FILE__);
  }
}

// Test that we remove a gmem->gmem permute whose allocation domain is
// unpermuted at lowering
TEST_F(RemoveTrivialOpsTest, Permute) {
  Fusion fusion;
  FusionGuard fg(&fusion);

  auto tv0 = makeContigConcreteTensor({2, 3});
  fusion.addInput(tv0);
  // This is a logical transpose. It implies data shuffling unless the
  // allocation domain is also transposed.
  auto tv1 = transpose(tv0);
  auto tv2 = neg(tv1);
  fusion.addOutput(tv2);

  auto options = at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0);
  auto t0 = at::randn({2, 3}, options);
  std::vector<c10::IValue> inputs{t0};

  tv1->setAllocationDomain(
      {
          tv1->getLogicalDomain().at(1),
          tv1->getLogicalDomain().at(0),
      },
      true);

  {
    // In this case we do not remove the permute since it is G->L not G->G
    KernelExecutor ke;
    ke.compile(&fusion, inputs);
    const kir::Kernel* kernel = ke.compiledKernel()->kernel();
    EXPECT_TRUE(KernelContainsExpr<LoadStoreOp>::check(kernel));
    auto outputs = ke.run(inputs);
    testValidate(&fusion, outputs, inputs, __LINE__, __FILE__);
  }

  {
    // Setting tv1 to Global means the LoadStoreOp is now G->G. This means that
    // when we lower it, we will be able to safely remove it
    tv1->setMemoryType(MemoryType::Global);

    KernelExecutor ke;
    ke.compile(&fusion, inputs);
    const kir::Kernel* kernel = ke.compiledKernel()->kernel();
    EXPECT_FALSE(KernelContainsExpr<LoadStoreOp>::check(kernel));
    EXPECT_EQ(kernel->summary().global_allocations.size(), 0)
        << "Expected to have no intermediate global allocations";
    auto outputs = ke.run(inputs);
    testValidate(&fusion, outputs, inputs, __LINE__, __FILE__);
  }

  {
    NVF_ERROR(tv1->getMemoryType() == MemoryType::Global);
    tv1->setAllocationDomain(tv1->getLogicalDomain(), true);

    KernelExecutor ke;

    ke.compile(&fusion, inputs);
    // We do not remove the transpose because the allocation domains do not
    // match
    const kir::Kernel* kernel = ke.compiledKernel()->kernel();
    EXPECT_TRUE(KernelContainsExpr<LoadStoreOp>::check(kernel));
    EXPECT_EQ(kernel->summary().global_allocations.size(), 1)
        << "Expected to have one intermediate global allocation";
    auto outputs = ke.run(inputs);
    testValidate(&fusion, outputs, inputs, __LINE__, __FILE__);
  }

  {
    // Reset tv1's allocation domain but also set it as an output
    tv1->setAllocationDomain(
        {
            tv1->getLogicalDomain().at(1),
            tv1->getLogicalDomain().at(0),
        },
        true);
    fusion.addOutput(tv1);

    KernelExecutor ke;
    ke.compile(&fusion, inputs);
    // We do not remove the transpose because the output is a fusion output
    const kir::Kernel* kernel = ke.compiledKernel()->kernel();
    EXPECT_TRUE(KernelContainsExpr<LoadStoreOp>::check(kernel));
    auto outputs = ke.run(inputs);
    testValidate(&fusion, outputs, inputs, __LINE__, __FILE__);
  }
}

// Test that we remove a gmem->gmem->gmem broadcast+squeeze at lowering
TEST_F(RemoveTrivialOpsTest, BroadcastSqueeze) {
  Fusion fusion;
  FusionGuard fg(&fusion);

  auto tv0 = makeContigConcreteTensor({2, 3});
  fusion.addInput(tv0);
  auto tv1 = broadcast(tv0, {false, true, false});
  auto tv2 = squeeze(tv1, {1});
  auto tv3 = neg(tv2);
  fusion.addOutput(tv3);

  auto options = at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0);
  auto t0 = at::randn({2, 3}, options);
  std::vector<c10::IValue> inputs{t0};

  {
    // In this case we do not remove the broadcast since it is G->L not G->G,
    // and the squeeze is L->L
    KernelExecutor ke;
    ke.compile(&fusion, inputs);
    const kir::Kernel* kernel = ke.compiledKernel()->kernel();
    // NOTE: SqueezeOp is converted to LoadStoreOp by the loadStoreOpInserter
    // lowering pass
    EXPECT_TRUE(KernelContainsExpr<BroadcastOp>::check(kernel));
    EXPECT_TRUE(KernelContainsExpr<LoadStoreOp>::check(kernel));
    auto outputs = ke.run(inputs);
    testValidate(&fusion, outputs, inputs, __LINE__, __FILE__);
  }

  {
    // Set G->G->G so that both ops get removed
    tv1->setMemoryType(MemoryType::Global);
    tv2->setMemoryType(MemoryType::Global);

    KernelExecutor ke;
    ke.compile(&fusion, inputs);
    const kir::Kernel* kernel = ke.compiledKernel()->kernel();
    EXPECT_FALSE(KernelContainsExpr<BroadcastOp>::check(kernel));
    EXPECT_FALSE(KernelContainsExpr<LoadStoreOp>::check(kernel));
    EXPECT_EQ(kernel->summary().global_allocations.size(), 0)
        << "Expected to have no intermediate global allocations";
    auto outputs = ke.run(inputs);
    testValidate(&fusion, outputs, inputs, __LINE__, __FILE__);
  }

  {
    NVF_ERROR(tv1->getMemoryType() == MemoryType::Global);
    NVF_ERROR(tv2->getMemoryType() == MemoryType::Global);
    // Transpose the input's allocation domain
    tv0->setAllocationDomain(
        {
            tv0->getLogicalDomain().at(1),
            tv0->getLogicalDomain().at(0),
        },
        true);

    KernelExecutor ke;

    std::vector<c10::IValue> transposed_inputs{at::randn({3, 2}, options).t()};
    ke.compile(&fusion, transposed_inputs);
    // We do not remove the broadcast because the allocation domains do not
    // match
    const kir::Kernel* kernel = ke.compiledKernel()->kernel();
    EXPECT_TRUE(KernelContainsExpr<BroadcastOp>::check(kernel));
    EXPECT_FALSE(KernelContainsExpr<LoadStoreOp>::check(kernel));
    EXPECT_EQ(kernel->summary().global_allocations.size(), 1)
        << "Expected to have one intermediate global allocation";
    auto outputs = ke.run(transposed_inputs);
    testValidate(&fusion, outputs, transposed_inputs, __LINE__, __FILE__);
  }

  {
    // Reset tv0's allocation domain but set tv1 as output
    tv0->setAllocationDomain(tv0->getLogicalDomain(), true);
    fusion.addOutput(tv1);

    KernelExecutor ke;
    ke.compile(&fusion, inputs);
    // We do not remove the broadcast because the output is a fusion output, but
    // we _do_ remove the squeeze
    const kir::Kernel* kernel = ke.compiledKernel()->kernel();
    EXPECT_TRUE(KernelContainsExpr<BroadcastOp>::check(kernel));
    EXPECT_FALSE(KernelContainsExpr<LoadStoreOp>::check(kernel));
    auto outputs = ke.run(inputs);
    testValidate(&fusion, outputs, inputs, __LINE__, __FILE__);
  }
}

} // namespace nvfuser
```

</details>

</td></tr>
</table>

author_association: NONE
user: login: jacobhinkle
created_at: 2025-02-26T01:33:47Z
body: !test
author_association: COLLABORATOR
user: login: jacobhinkle
created_at: 2025-02-26T12:22:49Z
body: !test
author_association: COLLABORATOR
user: login: jacobhinkle
created_at: 2025-02-26T16:00:19Z
body: !test --diff
author_association: COLLABORATOR
user: login: jacobhinkle
created_at: 2025-02-26T16:55:14Z
body: `NVFuserTest.FusionTestGridComm_CUDA` is currently failing. The scheduled fusion looks like this
```
Inputs:
  T0_g_float[iS0{3}, iS1{4}, iS2{2}]
  T1_g_float[iS3{3}, iS4{4}, iS5{2}]
Outputs:
  T5_g_float[iblockIdx.y15{3}, iblockIdx.x16{4}, iV17{2}]

%kernel_math {
T2_g_float[iblockIdx.y6{3}, iblockIdx.x7{4}, iV8{2}]
   = Set( T0_g_float[iS0{3}, iS1{4}, iS2{2}], cache_op=Streaming )
T3_g_float[iblockIdx.x9{3}, iblockIdx.y10{4}, iS11{2}]
   = T2_g_float[iblockIdx.y6{3}, iblockIdx.x7{4}, iV8{2}]
   + T1_g_float[iS3{3}, iS4{4}, iS5{2}];
T4_g_float[iblockIdx.y12{3}, iblockIdx.x13{4}, iS14{2}]
   = Set( T3_g_float[iblockIdx.x9{3}, iblockIdx.y10{4}, iS11{2}], cache_op=Streaming )
T5_g_float[iblockIdx.y15{3}, iblockIdx.x16{4}, iV17{2}]
   = Set( T4_g_float[iblockIdx.y12{3}, iblockIdx.x13{4}, iS14{2}], cache_op=Streaming )
} // %kernel_math
```
It produces this kernel on TOT:
```c++
__global__ void nvfuser_none_f0_c0_r0_g0(Tensor<float, 3, 3> T0, Tensor<float, 3, 3> T1, Tensor<float, 3, 3> T5, Tensor<float, 3, 3> T2, Tensor<float, 3, 3> T3, Tensor<int64_t, 1, 1> T6, Tensor<float, 3, 3> T4, Tensor<int64_t, 1, 1> T7) {
  NVFUSER_DEFINE_MAGIC_ZERO;
  nvfuser_index_t i0;
  i0 = (2 * ((nvfuser_index_t)blockIdx.x)) + (8 * ((nvfuser_index_t)blockIdx.y));
  nvfuser_index_t i1;
  i1 = (2 * ((nvfuser_index_t)blockIdx.y)) + (8 * ((nvfuser_index_t)blockIdx.x));
  bool b2;
  b2 = ((nvfuser_index_t)blockIdx.y) < 3;
  bool b3;
  b3 = ((nvfuser_index_t)blockIdx.x) < 3;
  // Allocate global tensor T2
  if (b2) {
    loadGlobalToGlobal<float, /*vec_size=*/2, /*is_volatile_to=*/true, /*is_volatile_from=*/false>( &*(volatile float*)&T2[i0],  &T0[i0]);
  }
  // Allocate global tensor T3
  // Allocate global tensor T6
  grid_sync::sync<true, true, false, true, true>(T6[index_utils::maskedOffset<false, false, true>(blockIdx, gridDim)], index_utils::maskedSize<true, true, false>(gridDim), DefaultBlockDim());
  #pragma unroll
  for(nvfuser_index_t i4 = 0; i4 < 2; ++i4) {
    nvfuser_index_t i5;
    i5 = i1 + (i4 + nvfuser_zero);
    if (b3) {
      *(volatile float*)&T3[i5]
        = *(volatile float*)&T2[i5]
        + T1[i5];
    }
  }
  NVFUSER_UPDATE_MAGIC_ZERO;
  // Allocate global tensor T4
  // Allocate global tensor T7
  grid_sync::sync<true, true, false, true, true>(T7[index_utils::maskedOffset<false, false, true>(blockIdx, gridDim)], index_utils::maskedSize<true, true, false>(gridDim), DefaultBlockDim());
  #pragma unroll
  for(nvfuser_index_t i6 = 0; i6 < 2; ++i6) {
    nvfuser_index_t i7;
    i7 = i0 + (i6 + nvfuser_zero);
    if (b2) {
      T4[i7]
         = *(volatile float*)&T3[i7];
    }
  }
  NVFUSER_UPDATE_MAGIC_ZERO;
  if (b2) {
    loadGlobalToGlobal<float, /*vec_size=*/2, /*is_volatile_to=*/false, /*is_volatile_from=*/false>( &T5[i0],  &T4[i0]);
  }
}
```
With this PR, we remove some of these G->G `set`s and the generated kernel becomes:
```c++
__global__ void nvfuser_none_f0_c0_r0_g0(Tensor<float, 3, 3> T0, Tensor<float, 3, 3> T1, Tensor<float, 3, 3> T5, Tensor<float, 3, 3> T3) {
  NVFUSER_DEFINE_MAGIC_ZERO;
  nvfuser_index_t i0;
  i0 = (2 * ((nvfuser_index_t)blockIdx.y)) + (8 * ((nvfuser_index_t)blockIdx.x));
  nvfuser_index_t i1;
  i1 = (2 * ((nvfuser_index_t)blockIdx.x)) + (8 * ((nvfuser_index_t)blockIdx.y));
  bool b2;
  b2 = ((nvfuser_index_t)blockIdx.x) < 3;
  // Allocate global tensor T3
  #pragma unroll
  for(nvfuser_index_t i3 = 0; i3 < 2; ++i3) {
    nvfuser_index_t i4;
    i4 = i0 + (i3 + nvfuser_zero);
    if (b2) {
      *(volatile float*)&T3[i4]
        = T0[i4]
        + T1[i4];
    }
  }
  NVFUSER_UPDATE_MAGIC_ZERO;
  if ((((nvfuser_index_t)blockIdx.y) < 3)) {
    loadGlobalToGlobal<float, /*vec_size=*/2, /*is_volatile_to=*/false, /*is_volatile_from=*/true>( &T5[i1],  &*(volatile float*)&T3[i1]);
  }
}
```
Notice we have lost the second grid sync. This is probably because the `loadGlobalToGlobal` expression is not seen as a "use" of `T3` when doing the sync analysis because `T3` is actually an alias. I'm investigating the sync analysis now to find a proper fix.
author_association: COLLABORATOR
user: login: jacobhinkle
created_at: 2025-02-26T18:07:56Z
body: > Notice we have lost the second grid sync. This is probably because the `loadGlobalToGlobal` expression is not seen as a "use" of `T3` when doing the sync analysis because `T3` is actually an alias. I'm investigating the sync analysis now to find a proper fix.

This is fixed in the latest small push by respecting aliases in `ReadAfterWriteSyncs`.
author_association: COLLABORATOR
user: login: jacobhinkle
created_at: 2025-02-26T18:08:03Z
body: !test --diff
author_association: COLLABORATOR
user: login: jacobhinkle
created_at: 2025-02-26T18:12:05Z
body: Now that RAW syncs are handled, my last remaining concern is about outputs that alias inputs for in-place operations (one of at least three uses of the term "alias"). I think we should probably just avoid doing this new "tensor producer aliasing" to any input tensor that is itself the target of an output alias, since it complicates the sync analysis. It also should not be necessary for the matmul use case.
author_association: COLLABORATOR
user: login: jacobhinkle
created_at: 2025-03-03T15:11:31Z
body: !test
author_association: COLLABORATOR
user: login: jacobhinkle
created_at: 2025-03-04T15:36:51Z
body: !test
author_association: COLLABORATOR
user: login: jacobhinkle
created_at: 2025-03-04T23:32:26Z
body: !test
author_association: COLLABORATOR