FullReduceTransposeReshape rewrite

src/enzyme_ad/jax/Passes/EnzymeHLOOpt.cpp

@@ -4319,6 +4319,32 @@ struct ReduceConcat final
 struct FullReduceReshapeOrTranspose final
     : CheckedOpRewritePattern<stablehlo::ReduceOp,
                               FullReduceReshapeOrTranspose> {
+private:
+  // Affine map sending an index of `fromShape` to the corresponding index of
+  // `toShape` under the row-major element ordering used by reshape.
+  static AffineMap reshapeMap(MLIRContext *ctx, ArrayRef<int64_t> fromShape,
+                              ArrayRef<int64_t> toShape) {
+    AffineExpr linear = getAffineConstantExpr(0, ctx);
+    int64_t stride = 1;
+    for (int i = fromShape.size() - 1; i >= 0; --i) {
+      linear = linear + getAffineDimExpr(i, ctx) * stride;
+      stride *= fromShape[i];
+    }
+
+    SmallVector<AffineExpr> exprs(toShape.size());
+    stride = 1;
+    for (int i = toShape.size() - 1; i >= 0; --i) {
+      AffineExpr expr = linear.floorDiv(stride);
+      if (i > 0)
+        expr = expr % toShape[i];
+      exprs[i] = expr;
+      stride *= toShape[i];
+    }
+
+    return AffineMap::get(fromShape.size(), 0, exprs, ctx);
+  }
+
+public:
   using CheckedOpRewritePattern::CheckedOpRewritePattern;
 
   LogicalResult matchAndRewriteImpl(stablehlo::ReduceOp op,
@@ -4330,115 +4356,224 @@ struct FullReduceReshapeOrTranspose final
     if (op.getDimensions().size() != inpTy.getShape().size())
       return failure();
 
-    RankedTensorType changeType = nullptr;
+    MLIRContext *ctx = op.getContext();
+
+    // The reduce spans every dimension, so it is invariant under any bijective
+    // reindexing (reshape/transpose) of its input as long as, for every
+    // elementwise op in the cone, all operands agree on the same reindexing.
+    // For every value in the cone we track that reindexing as an affine map
+    // from the value's index space to the index space of the underlying
+    // pre-reshape/transpose data.
+    llvm::MapVector<Value, AffineMap> transforms;
+    // Maps reshape/transpose results in the cone to their pre-transform value.
+    IRMapping mapping;
+    // Elementwise ops whose operand transforms all agree, in post-order (defs
+    // before users).
+    SetVector<Value> tape;
+    // All rewritable cone ops in post-order, for the final erasure sweep.
+    SmallVector<Operation *> coneOps;
     SmallVector<Operation *> reshapeOrTransposes;
-    llvm::MapVector<Operation *, int> toclone;
+
+    auto getTransform = [&](Value val) -> AffineMap {
+      auto found = transforms.find(val);
+      if (found != transforms.end())
+        return found->second;
+      auto ty = dyn_cast<RankedTensorType>(val.getType());
+      return AffineMap::getMultiDimIdentityMap(ty ? ty.getRank() : 0, ctx);
+    };
+
     {
       SmallVector<Value> todo = {op.getInputs()[0]};
-      while (todo.size()) {
-        auto cur = todo.pop_back_val();
-        auto curOp = cur.getDefiningOp();
-        if (!curOp)
-          return failure();
-        if (auto rs = dyn_cast<stablehlo::ReshapeOp>(curOp)) {
-          if (changeType != nullptr) {
-            if (rs.getOperand().getType() != changeType)
-              return failure();
-          } else {
-            changeType = rs.getOperand().getType();
-          }
-          reshapeOrTransposes.push_back(rs);
+      DenseSet<Value> processed;
+      DenseMap<Operation *, unsigned> childIdx;
+      while (!todo.empty()) {
+        Value v = todo.back();
+        if (processed.contains(v)) {
+          todo.pop_back();
           continue;
         }
-        if (auto rs = dyn_cast<stablehlo::TransposeOp>(curOp)) {
-          if (changeType != nullptr) {
-            if (rs.getOperand().getType() != changeType)
-              return failure();
-          } else {
-            changeType = rs.getOperand().getType();
-          }
-          reshapeOrTransposes.push_back(rs);
+
+        Operation *curOp = v.getDefiningOp();
+        auto vTy = dyn_cast<RankedTensorType>(v.getType());
+        bool leaf =
+            !curOp || !vTy || !vTy.hasStaticShape() ||
+            curOp->getNumResults() != 1 || curOp->getNumOperands() == 0 ||
+            curOp->getNumRegions() != 0 ||
+            (!isa<stablehlo::ReshapeOp, stablehlo::TransposeOp>(curOp) &&
+             !hasTraitElementwise(curOp)) ||
+            !isMemoryEffectFree(curOp) ||
+            !llvm::all_of(curOp->getOperands(), [](Value operand) {
+              auto ty = dyn_cast<RankedTensorType>(operand.getType());
+              return ty && ty.hasStaticShape();
+            });
+        if (leaf) {
+          processed.insert(v);
+          todo.pop_back();
           continue;
         }
-        if (!hasTraitElementwise(curOp))
-          return failure();
-        if (!isMemoryEffectFree(curOp))
-          return failure();
-        for (auto op : curOp->getOperands())
-          todo.push_back(op);
-        toclone[curOp] = curOp->getNumOperands();
-      }
-    }
 
-    IRMapping map;
-    SmallVector<Operation *> todo;
-    for (auto reshape : reshapeOrTransposes) {
-      map.map(reshape->getResult(0), reshape->getOperand(0));
-      for (auto u : reshape->getResult(0).getUsers()) {
-        if (toclone.contains(u)) {
-          toclone[u]--;
-          if (toclone[u] == 0) {
-            todo.push_back(u);
-            toclone.erase(u);
-          }
-        }
-      }
-    }
-    for (auto pair : toclone) {
-      for (auto u : pair.first->getResult(0).getUsers()) {
-        if (u == op)
+        unsigned &idx = childIdx[curOp];
+        if (idx < curOp->getNumOperands()) {
+          todo.push_back(curOp->getOperand(idx++));
           continue;
-        if (llvm::is_contained(reshapeOrTransposes, u))
+        }
+
+        // All operands of curOp have been processed already.
+        processed.insert(v);
+        todo.pop_back();
+
+        if (auto ts = dyn_cast<stablehlo::TransposeOp>(curOp)) {
+          // Maps a result index to the corresponding operand index.
+          auto map = AffineMap::getPermutationMap(ts.getPermutation(), ctx);
+          transforms[v] = getTransform(ts.getOperand()).compose(map);
+          mapping.map(v, mapping.lookupOrDefault(ts.getOperand()));
+          coneOps.push_back(curOp);
+          reshapeOrTransposes.push_back(curOp);
           continue;
-        if (toclone.contains(u))
+        }
+
+        if (auto rs = dyn_cast<stablehlo::ReshapeOp>(curOp)) {
+          auto inTy = cast<RankedTensorType>(rs.getOperand().getType());
+          auto outTy = cast<RankedTensorType>(rs.getResult().getType());
+          // Maps a result index to the corresponding operand index.
+          auto map = reshapeMap(ctx, outTy.getShape(), inTy.getShape());
+          transforms[v] = getTransform(rs.getOperand()).compose(map);
+          mapping.map(v, mapping.lookupOrDefault(rs.getOperand()));
+          coneOps.push_back(curOp);
+          reshapeOrTransposes.push_back(curOp);
           continue;
-        return failure();
-      }
-    }
+        }
 
-    for (auto cur : todo) {
-      for (auto curOp : cur->getOperands()) {
-        if (!map.contains(curOp))
-          return failure();
+        // Elementwise op: rewritable iff all operands agree on the transform.
+        // Operands outside the cone (or unrewritable ones) count as identity.
+        AffineMap ref = getTransform(curOp->getOperand(0));
+        if (llvm::all_of(curOp->getOperands(), [&](Value operand) {
+              return getTransform(operand) == ref;
+            })) {
+          transforms[v] = ref;
+          tape.insert(v);
+          coneOps.push_back(curOp);
+        }
       }
     }
 
-    while (todo.size()) {
-      auto cur = todo.pop_back_val();
-
-      SmallVector<Value> vals;
-      for (auto op : cur->getOperands())
-        vals.push_back(map.lookup(op));
-
-      auto changeType2 = RankedTensorType::get(
-          changeType.getShape(),
-          cast<RankedTensorType>(cur->getResult(0).getType()).getElementType());
-      auto res =
-          rewriter.create(cur->getLoc(), cur->getName().getIdentifier(), vals,
-                          TypeRange(changeType2), cur->getAttrs(), {}, {});
-
-      map.map(cur->getResult(0), res->getResult(0));
+    // A reshape/transpose only becomes dead if all its users are rewritten.
+    // If it has a user outside the cone, rewriting on top of its pre-transform
+    // value would duplicate the computation instead of removing it, so block
+    // rewrites building on that value.
+    DenseSet<Value> blockedValues;
+    auto inCone = [&](Operation *user) {
+      if (user == op)
+        return true;
+      if (user->getNumResults() != 1)
+        return false;
+      Value result = user->getResult(0);
+      return tape.contains(result) || mapping.contains(result);
+    };
+    for (Operation *rt : reshapeOrTransposes)
+      if (!llvm::all_of(rt->getResult(0).getUsers(), inCone))
+        blockedValues.insert(mapping.lookup(rt->getResult(0)));
+
+    // Decide, without mutating the IR yet, what the rewritten form of each
+    // tape value is: unchanged (all operands unchanged), a clone on the
+    // pre-transform operands, or blocked (depends on a blocked value; keeping
+    // the original avoids mixing pre- and post-transform shapes).
+    enum class Status { Unchanged, Clone, Blocked };
+    DenseMap<Value, Status> status;
+    DenseMap<Value, RankedTensorType> newTypes;
+
+    struct OperandInfo {
+      bool blocked;
+      bool changed;
+      RankedTensorType type;
+    };
+    auto operandInfo = [&](Value operand) -> OperandInfo {
+      if (mapping.contains(operand)) {
+        Value pre = mapping.lookup(operand);
+        if (blockedValues.contains(pre))
+          return {true, false, nullptr};
+        return {false, true, cast<RankedTensorType>(pre.getType())};
+      }
+      auto found = status.find(operand);
+      if (found != status.end() && found->second == Status::Blocked)
+        return {true, false, nullptr};
+      if (found != status.end() && found->second == Status::Clone)
+        return {false, true, newTypes.lookup(operand)};
+      return {false, false, dyn_cast<RankedTensorType>(operand.getType())};
+    };
 
-      for (auto u : cur->getResult(0).getUsers()) {
-        if (toclone.contains(u)) {
-          toclone[u]--;
-          if (toclone[u] == 0) {
-            todo.push_back(u);
-            toclone.erase(u);
-          }
+    for (Value v : tape) {
+      Operation *curOp = v.getDefiningOp();
+      bool anyBlocked = false, anyChanged = false;
+      RankedTensorType operandTy = nullptr;
+      for (Value operand : curOp->getOperands()) {
+        OperandInfo info = operandInfo(operand);
+        if (info.blocked) {
+          anyBlocked = true;
+          break;
+        }
+        anyChanged |= info.changed;
+        if (!operandTy)
+          operandTy = info.type;
+        else if (operandTy.getShape() != info.type.getShape()) {
+          // Equal transforms should imply equal pre-transform shapes; be
+          // conservative if they do not.
+          anyBlocked = true;
+          break;
         }
       }
+      if (anyBlocked)
+        status[v] = Status::Blocked;
+      else if (!anyChanged)
+        status[v] = Status::Unchanged;
+      else {
+        status[v] = Status::Clone;
+        newTypes[v] = RankedTensorType::get(
+            operandTy.getShape(),
+            cast<RankedTensorType>(v.getType()).getElementType());
+      }
+    }
+
+    // No-op detection: only rewrite if the reduce input actually changes,
+    // i.e. at least one reshape/transpose is bypassed on the way to it.
+    OperandInfo rootInfo = operandInfo(op.getInputs()[0]);
+    if (rootInfo.blocked || !rootInfo.changed)
+      return failure();
+
+    rewriter.setInsertionPoint(op);
+    for (Value v : tape) {
+      if (status.lookup(v) != Status::Clone)
+        continue;
+      Operation *curOp = v.getDefiningOp();
+      SmallVector<Value> newOperands;
+      for (Value operand : curOp->getOperands())
+        newOperands.push_back(mapping.lookupOrDefault(operand));
+      Type newType = newTypes.lookup(v);
+      Operation *newOp = rewriter.create(
+          curOp->getLoc(), curOp->getName().getIdentifier(), newOperands,
+          TypeRange(newType), curOp->getAttrs(), {}, {});
+      mapping.map(v, newOp->getResult(0));
     }
 
+    Value newInit = mapping.lookupOrDefault(op.getInputs()[0]);
+
     SmallVector<int64_t> newReduceDimensions;
-    for (size_t i = 0, end = changeType.getShape().size(); i < end; i++)
+    for (int64_t i = 0, e = cast<RankedTensorType>(newInit.getType()).getRank();
+         i < e; ++i)
       newReduceDimensions.push_back(i);
 
     auto newReduction = stablehlo::ReduceOp::create(
-        rewriter, op.getLoc(), op->getResultTypes(),
-        map.lookup(op.getInputs()[0]), op.getInitValues(), newReduceDimensions);
+        rewriter, op.getLoc(), op->getResultTypes(), ValueRange{newInit},
+        op.getInitValues(), newReduceDimensions);
     newReduction.getRegion().takeBody(op.getRegion());
     rewriter.replaceOp(op, newReduction);
+
+    // Erase the bypassed ops in reverse post-order (users before defs) so a
+    // single sweep removes the whole dead cone.
+    for (Operation *cur : llvm::reverse(coneOps))
+      if (isOpTriviallyDead(cur))
+        rewriter.eraseOp(cur);
+
     return success();
   }
 };

test/lit_tests/full_reduce_transpose.mlir

@@ -0,0 +1,21 @@
+// RUN: enzymexlamlir-opt %s "--enzyme-hlo-generate-td=patterns=full_reduce_reshape_or_transpose" --transform-interpreter --enzyme-hlo-remove-transform | FileCheck %s
+
+func.func @main(%g: tensor<2x2x2xf64>, %w: tensor<2x2x2xf64>) -> tensor<f64> {
+  %cst = stablehlo.constant dense<0.000000e+00> : tensor<f64>
+  %a = stablehlo.transpose %g, dims = [0, 2, 1] : (tensor<2x2x2xf64>) -> tensor<2x2x2xf64>
+  %b = stablehlo.transpose %w, dims = [2, 1, 0] : (tensor<2x2x2xf64>) -> tensor<2x2x2xf64>
+  %m = stablehlo.multiply %a, %b : tensor<2x2x2xf64>
+  %n = stablehlo.transpose %m, dims = [0, 2, 1] : (tensor<2x2x2xf64>) -> tensor<2x2x2xf64>
+  %r = stablehlo.reduce(%n init: %cst) applies stablehlo.add across dimensions = [0, 1, 2]
+      : (tensor<2x2x2xf64>, tensor<f64>) -> tensor<f64>
+  return %r : tensor<f64>
+}
+
+// CHECK:  func.func @main(%arg0: tensor<2x2x2xf64>, %arg1: tensor<2x2x2xf64>) -> tensor<f64> {
+// CHECK-NEXT:    %cst = stablehlo.constant dense<0.000000e+00> : tensor<f64>
+// CHECK-NEXT:    %0 = stablehlo.transpose %arg0, dims = [0, 2, 1] : (tensor<2x2x2xf64>) -> tensor<2x2x2xf64>
+// CHECK-NEXT:    %1 = stablehlo.transpose %arg1, dims = [2, 1, 0] : (tensor<2x2x2xf64>) -> tensor<2x2x2xf64>
+// CHECK-NEXT:    %2 = stablehlo.multiply %0, %1 : tensor<2x2x2xf64>
+// CHECK-NEXT:    %3 = stablehlo.reduce(%2 init: %cst) applies stablehlo.add across dimensions = [0, 1, 2] : (tensor<2x2x2xf64>, tensor<f64>) -> tensor<f64>
+// CHECK-NEXT:    return %3 : tensor<f64>
+// CHECK-NEXT:  }

test/lit_tests/reducetransposereshape.mlir

@@ -0,0 +1,26 @@
+// RUN: enzymexlamlir-opt %s "--enzyme-hlo-generate-td=patterns=full_reduce_reshape_or_transpose" --transform-interpreter --enzyme-hlo-remove-transform | FileCheck %s
+
+module {
+  func.func @main(%arg0: tensor<4x5xf32>, %arg1: tensor<20xf32>) -> tensor<f32> {
+    %cst = stablehlo.constant dense<0.0> : tensor<f32>
+
+    %ts = stablehlo.transpose %arg0, dims = [1, 0] : (tensor<4x5xf32>) -> tensor<5x4xf32>
+    %rs = stablehlo.reshape %ts : (tensor<5x4xf32>) -> tensor<20xf32>
+    %add = stablehlo.add %rs, %arg1 : tensor<20xf32>
+
+    %v =  stablehlo.reshape %add : (tensor<20xf32>) -> tensor<4x5xf32>
+    %red = stablehlo.reduce(%v init: %cst) applies stablehlo.add across dimensions = [0, 1]
+        : (tensor<4x5xf32>, tensor<f32>) -> tensor<f32>
+    return %red : tensor<f32>
+  }
+
+}
+
+// CHECK:  func.func @main(%arg0: tensor<4x5xf32>, %arg1: tensor<20xf32>) -> tensor<f32> {
+// CHECK-NEXT:    %cst = stablehlo.constant dense<0.000000e+00> : tensor<f32>
+// CHECK-NEXT:    %0 = stablehlo.transpose %arg0, dims = [1, 0] : (tensor<4x5xf32>) -> tensor<5x4xf32>
+// CHECK-NEXT:    %1 = stablehlo.reshape %0 : (tensor<5x4xf32>) -> tensor<20xf32>
+// CHECK-NEXT:    %2 = stablehlo.add %1, %arg1 : tensor<20xf32>
+// CHECK-NEXT:    %3 = stablehlo.reduce(%2 init: %cst) applies stablehlo.add across dimensions = [0] : (tensor<20xf32>, tensor<f32>) -> tensor<f32>
+// CHECK-NEXT:    return %3 : tensor<f32>
+// CHECK-NEXT:  }

test/lit_tests/transpose_if2.mlir

@@ -43,7 +43,7 @@ module @reactant_conditi... attributes {mhlo.num_partitions = 1 : i64, mhlo.num_
 // CHECK-NEXT:     %cst_1 = stablehlo.constant dense<1.000000e+00> : tensor<10x2xf64>
 // CHECK-NEXT:     %0 = stablehlo.add %arg0, %cst_1 : tensor<10x2xf64>
 // CHECK-NEXT:     %1 = stablehlo.negate %0 : tensor<10x2xf64>
-// CHECK-NEXT:     %2 = stablehlo.reduce(%0 init: %cst_0) applies stablehlo.add across dimensions = [0, 1] {enzymexla.non_negative = [#enzymexla<guaranteed NOTGUARANTEED>]} : (tensor<10x2xf64>, tensor<f64>) -> tensor<f64>
+// CHECK-NEXT:     %2 = stablehlo.reduce(%0 init: %cst_0) applies stablehlo.add across dimensions = [1, 0] {enzymexla.non_negative = [#enzymexla<guaranteed NOTGUARANTEED>]} : (tensor<10x2xf64>, tensor<f64>) -> tensor<f64>
 // CHECK-NEXT:     %3 = stablehlo.compare  GT, %2, %cst_0 : (tensor<f64>, tensor<f64>) -> tensor<i1>
 // CHECK-NEXT:     %4 = stablehlo.select %3, %arg0, %0 : tensor<i1>, tensor<10x2xf64>
 // CHECK-NEXT:     %5 = stablehlo.select %3, %0, %1 : tensor<i1>, tensor<10x2xf64>