Reverts 847b2a2

PiperOrigin-RevId: 899495748

Author: ekayaaslan

shardy/dialect/sdy/transforms/import/constant_or_scalar_splitter.cc

@@ -26,7 +26,6 @@ limitations under the License.
 #include "mlir/IR/OpDefinition.h"
 #include "mlir/IR/Operation.h"
 #include "mlir/IR/PatternMatch.h"
-#include "mlir/IR/SymbolTable.h"
 #include "mlir/IR/Value.h"
 #include "mlir/IR/Visitors.h"
 #include "mlir/Interfaces/SideEffectInterfaces.h"
@@ -49,7 +48,6 @@ namespace sdy {
 
 namespace {
 
-using func::CallOp;
 using func::FuncOp;
 
 void cloneShardingGroupUsers(OpResult opResult, IRMapping& mapping,
@@ -67,12 +65,14 @@ void cloneShardingGroupUsers(OpResult opResult, IRMapping& mapping,
 // given op is either:
 // - A broadcast, reshape or slice op.
 // - An elementwise op.
-// - A call to a func that all operations are constant preserving.
+// - A named computation all operations are constant preserving.
 // Assumes the op is not constant or iota.
 bool isConstantPreserving(
-    Operation* op, const llvm::SmallDenseSet<StringRef>& nonConstFuncOps) {
-  if (CallOp callOp = dyn_cast<CallOp>(op)) {
-    return !nonConstFuncOps.contains(callOp.getCallee());
+    Operation* op,
+    const llvm::SmallDenseSet<StringRef>& nonConstantNamedComputationOps) {
+  if (auto namedComputationOp = dyn_cast<NamedComputationOp>(op)) {
+    return !nonConstantNamedComputationOps.contains(
+        namedComputationOp.getName());
   }
   if (!isPure(op)) {
     return false;
@@ -93,11 +93,11 @@ bool isConstantPreserving(
 // constants, that is, exist in `constantOps`.
 bool isConstantExpression(
     Operation* op, const llvm::SetVector<Operation*>& constantOps,
-    const llvm::SmallDenseSet<StringRef>& nonConstFuncOps) {
+    const llvm::SmallDenseSet<StringRef>& nonConstantNamedComputationOps) {
   if (isa<ConstantOp, stablehlo::IotaOp>(op)) {
     return true;
   }
-  return isConstantPreserving(op, nonConstFuncOps) &&
+  return isConstantPreserving(op, nonConstantNamedComputationOps) &&
          llvm::all_of(op->getOperands(), [&](Value operand) {
            return operand.getDefiningOp() &&
                   constantOps.contains(operand.getDefiningOp());
@@ -117,26 +117,20 @@ bool isScalarExpansion(Operation* op) {
 // Recursively clones all operands of the given op, that are not already mapped
 // in `mapping`, and finally clones the op itself. We do not clone scalars as
 // they do not get sharded.
-void cloneSubComputation(OpResult opResult, IRMapping& mapping,
-                         SymbolTable& symbolTable) {
+void cloneSubComputation(OpResult opResult, IRMapping& mapping) {
   if (isScalar(opResult) || mapping.lookupOrNull(opResult)) {
     return;
   }
   Operation* op = opResult.getOwner();
   for (Value operand : op->getOperands()) {
     if (auto defOpResult = dyn_cast<OpResult>(operand)) {
-      cloneSubComputation(defOpResult, mapping, symbolTable);
+      cloneSubComputation(defOpResult, mapping);
     }
   }
 
   // This will insert the cloned op right before the original op.
   OpBuilder builder(op);
-  Operation* clonedOp = builder.clone(*op, mapping);
-  if (CallOp callOp = dyn_cast<CallOp>(clonedOp)) {
-    FuncOp funcOp = symbolTable.lookup<FuncOp>(callOp.getCallee());
-    callOp.setCallee(
-        symbolTable.insert(cloneFuncRecursively(funcOp, symbolTable)));
-  }
+  builder.clone(*op, mapping);
   cloneShardingGroupUsers(opResult, mapping, builder);
 }
 
@@ -145,19 +139,18 @@ void cloneSubComputation(OpResult opResult, IRMapping& mapping,
 // sharded.
 //
 // Returns the cloned op result.
-Value cloneSubComputation(OpResult opResult, SymbolTable& symbolTable) {
+Value cloneSubComputation(OpResult opResult) {
   if (isScalar(opResult)) {
     return opResult;
   }
   IRMapping mapping;
-  cloneSubComputation(opResult, mapping, symbolTable);
+  cloneSubComputation(opResult, mapping);
   return mapping.lookup(opResult);
 }
 
 void cloneSubComputationOnOperands(
     Operation* op, const llvm::SetVector<Operation*>& constantOps,
-    const llvm::SetVector<Operation*>& scalarExpansionOps,
-    SymbolTable& symbolTable) {
+    const llvm::SetVector<Operation*>& scalarExpansionOps) {
   for (OpOperand& operand : op->getOpOperands()) {
     if (auto defOpResult = dyn_cast<OpResult>(operand.get());
         defOpResult && (constantOps.contains(defOpResult.getOwner()) ||
@@ -167,38 +160,38 @@ void cloneSubComputationOnOperands(
       // `defOpResult`, and replace the `operand` with the cloned defining
       // op. The cloned constant sub-computation has only one user `op`,
       // so that it is isolated from the rest of the computation.
-      operand.set(cloneSubComputation(defOpResult, symbolTable));
+      operand.set(cloneSubComputation(defOpResult));
     }
   }
 }
 
-void processOp(Operation* op, FuncOp funcOp,
-               llvm::SetVector<Operation*>& constantOps,
+void processOp(Operation* op, llvm::SetVector<Operation*>& constantOps,
                llvm::SetVector<Operation*>& scalarExpansionOps,
-               llvm::SmallDenseSet<StringRef>& nonConstFuncOps,
-               SymbolTable& symbolTable) {
-  if (isa<FuncOp, ShardingGroupOp>(op)) {
+               llvm::SmallDenseSet<StringRef>& nonConstantNamedComputationOps) {
+  if (isa<ShardingGroupOp>(op)) {
     return;
   }
-  if (isConstantExpression(op, constantOps, nonConstFuncOps)) {
+  if (isConstantExpression(op, constantOps, nonConstantNamedComputationOps)) {
     constantOps.insert(op);
     return;
   }
   // NOTE: There are cases that op is an constant expression but may not pass
   // the following check such as constant and iota ops. That is fine because if
   // the op is a constant expression it is a stronger condition than being just
-  // constant preserving and it does not make the `funcOp` non-const, and at
-  // this point, it is guaranteed that the op is not constant expression.
-  if (!isConstantPreserving(op, nonConstFuncOps) &&
+  // constant preserving and it does not make the parent named computation
+  // non-const, and at this point, it is guaranteed that the op is not constant
+  // expression.
+  if (!isConstantPreserving(op, nonConstantNamedComputationOps) &&
       !op->hasTrait<OpTrait::IsTerminator>()) {
-    nonConstFuncOps.insert(funcOp.getName());
+    if (auto namedCompuationOp = op->getParentOfType<NamedComputationOp>()) {
+      nonConstantNamedComputationOps.insert(namedCompuationOp.getName());
+    }
   }
   if (isScalarExpansion(op)) {
     scalarExpansionOps.insert(op);
     return;
   }
-  cloneSubComputationOnOperands(op, constantOps, scalarExpansionOps,
-                                symbolTable);
+  cloneSubComputationOnOperands(op, constantOps, scalarExpansionOps);
 }
 
 // Converts stablehlo::ConstantOp to sdy::ConstantOp.
@@ -247,14 +240,21 @@ struct ConstantOrScalarSplitterPass
     }
   }
 
-  void walkOnRegion(FuncOp funcOp,
-                    llvm::SmallDenseSet<StringRef>& nonConstFuncOps,
-                    SymbolTable& symbolTable) {
+  // Assumes that the `NamedComputationOp` of the region are already walked, and
+  // skips walking on them.
+  void walkOnRegion(
+      mlir::Region& region,
+      llvm::SmallDenseSet<StringRef>& nonConstantNamedComputationOps) {
     llvm::SetVector<Operation*> constantOps;
     llvm::SetVector<Operation*> scalarExpansionOps;
-    funcOp.walk<WalkOrder::PreOrder>([&](Operation* op) {
-      processOp(op, funcOp, constantOps, scalarExpansionOps, nonConstFuncOps,
-                symbolTable);
+    region.walk<WalkOrder::PreOrder>([&](Operation* op) {
+      processOp(op, constantOps, scalarExpansionOps,
+                nonConstantNamedComputationOps);
+      // Skip walking on the `NamedComputationOp`.
+      if (isa<NamedComputationOp>(op)) {
+        return WalkResult::skip();
+      }
+      return WalkResult::advance();
     });
     // Since for every op in `constantOps` that has a use that isn't in
     // `constantOps`, we replaced the use with a clone of the entire
@@ -267,7 +267,6 @@ struct ConstantOrScalarSplitterPass
 
   void runOnOperation() final {
     ModuleOp moduleOp = getOperation();
-    SymbolTable symbolTable(moduleOp);
 
     // We first convert any `stablehlo::ConstantOp` to an `sdy::ConstantOp`, so
     // that constants won't be deduped via folding.
@@ -276,13 +275,17 @@ struct ConstantOrScalarSplitterPass
     }
 
     // Then we split constant sub-computations for each non-constant user.
-    llvm::SmallDenseSet<StringRef> nonConstFuncOps;
-    FuncOp mainFuncOp = walkCallsOrDie(moduleOp, [&](CallOp callOp) {
-      FuncOp funcOp = symbolTable.lookup<FuncOp>(callOp.getCallee());
-      walkOnRegion(funcOp, nonConstFuncOps, symbolTable);
-      return WalkResult::advance();
+    llvm::SmallDenseSet<StringRef> nonConstantNamedComputationOps;
+    // Iterate on a post-order of NamedComputationOp blocks.
+    moduleOp.walk([&](NamedComputationOp namedComputationOp) {
+      walkOnRegion(namedComputationOp.getBody(),
+                   nonConstantNamedComputationOps);
+    });
+    // Iterate order does not matter. Funcs do not call each other. The calls
+    // are inlined to NamedComputationOps.
+    moduleOp.walk([&](FuncOp funcOp) {
+      walkOnRegion(funcOp.getBody(), nonConstantNamedComputationOps);
     });
-    walkOnRegion(mainFuncOp, nonConstFuncOps, symbolTable);
   }
 
  private:

shardy/dialect/sdy/transforms/import/import_pipeline.cc

@@ -30,6 +30,9 @@ void addImportPipeline(OpPassManager& pm, int& dumpIndex,
   pm.addPass(createSymbolDCEPass());
   pm.addPass(createLiftInlinedMeshesPass());
   pm.addPass(createRemoveSizeOneAxesPass());
+  pm.addPass(createImportFuncCallsPass());
+  // Keep SymbolDCEPass after ImportFuncCallsPass.
+  pm.addPass(createSymbolDCEPass());
   pm.addPass(createConstantOrScalarSplitterPass());
   pm.addPass(createSymbolDCEPass());
   pm.addPass(createManualAxesCleanupPass());
@@ -40,7 +43,6 @@ void addImportPipeline(OpPassManager& pm, int& dumpIndex,
   // of the propagation itself.
   pm.addPass(mlir::sdy::createSaveModuleOpPass(
       options.dumpDirectory, "before_propagation", dumpIndex++));
-  pm.addPass(createImportFuncCallsPass());
 
   pm.addNestedPass<func::FuncOp>(createAddDataFlowEdgesPass(
       AddDataFlowEdgesPassOptions{options.enableNativeNonFlatSupport}));