#sdy Copy over change to import_backend_func_calls

PiperOrigin-RevId: 738819740

Author: tomnatan30

shardy/round_trip_import/BUILD

@@ -42,9 +42,9 @@ cc_library(
 )
 
 cc_library(
-    name = "import_backend_func_calls",
-    srcs = ["import_backend_func_calls.cc"],
-    hdrs = ["import_backend_func_calls.h"],
+    name = "import_uninlineable_func_calls",
+    srcs = ["import_uninlineable_func_calls.cc"],
+    hdrs = ["import_uninlineable_func_calls.h"],
     deps = [
         ":constants",
         ":utils",
@@ -118,10 +118,10 @@ cc_library(
     srcs = ["pipelines.cc"],
     hdrs = ["pipelines.h"],
     deps = [
-        ":import_backend_func_calls",
         ":import_callback_custom_calls",
         ":import_sdy_custom_calls",
         ":import_shardy_attrs",
+        ":import_uninlineable_func_calls",
         ":shard_map_import",
         "@llvm-project//mlir:FuncDialect",
         "@llvm-project//mlir:Pass",

shardy/round_trip_import/constants.h

@@ -46,6 +46,9 @@ inline constexpr llvm::StringRef kFFIPythonGpuCallbackCustomCallTargetName =
 // The attribute name for backend config.
 inline constexpr llvm::StringRef kXlaBackendConfigAttr = "backend_config";
 
+// The attribute name for inlineable.
+inline constexpr llvm::StringRef kXlaInlineableAttr = "inlineable";
+
 // Attribute name for temporarily storing the Shardy sharding during HLO
 // sdy-round-trip. It cannot match the name `kShardingAttr` ("sdy.sharding"), as
 // during sdy-round-trip, going from HLO to StableHLO, the code removes

shardy/round_trip_import/import_backend_func_calls.cc

@@ -0,0 +1,158 @@
+/* Copyright 2025 The Shardy Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "shardy/round_trip_import/import_uninlineable_func_calls.h"
+
+#include <cassert>
+#include <iterator>
+#include <memory>
+#include <string>
+
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/Support/FormatVariadic.h"
+#include "llvm/Support/Threading.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/IR/Attributes.h"
+#include "mlir/IR/BuiltinAttributes.h"
+#include "mlir/IR/BuiltinOps.h"
+#include "mlir/IR/Diagnostics.h"
+#include "mlir/IR/OperationSupport.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/IR/SymbolTable.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Pass/PassRegistry.h"
+#include "mlir/Support/LLVM.h"
+#include "mlir/Support/TypeID.h"
+#include "mlir/Transforms/DialectConversion.h"
+#include "shardy/dialect/sdy/ir/constants.h"
+#include "shardy/dialect/sdy/ir/dialect.h"
+#include "shardy/dialect/sdy/ir/utils.h"
+#include "shardy/round_trip_import/constants.h"
+#include "shardy/round_trip_import/utils.h"
+
+namespace mlir {
+namespace sdy {
+
+namespace {
+
+using func::CallOp;
+using func::FuncOp;
+
+bool isInlineableCallOp(CallOp callOp) {
+  if (hasFrontendAttr(callOp, kXlaBackendConfigAttr)) {
+    return false;
+  }
+  auto inlineableAttr =
+      tryGetFrontendAttr<BoolAttr>(callOp, kXlaInlineableAttr);
+  return !inlineableAttr || inlineableAttr->getValue();
+}
+
+void importCallOp(
+    CallOp callOp,
+    llvm::SmallDenseMap<StringRef, Region*>& calleeNameToMovedRegion,
+    IRRewriter& rewriter, SymbolTable& symbolTable) {
+  SmallVector<NamedAttribute> namedCompAttrs;
+  llvm::copy_if(callOp->getDiscardableAttrs(),
+                std::back_inserter(namedCompAttrs),
+                [](const NamedAttribute& attr) {
+                  return attr.getName() != kShardingAttr;
+                });
+
+  StringRef calleeName = callOp.getCallee();
+  rewriter.setInsertionPoint(callOp);
+  auto namedCompOp = rewriter.create<NamedComputationOp>(
+      callOp->getLoc(), callOp->getResultTypes(), calleeName,
+      callOp.getOperands(),
+      /*inShardings=*/nullptr,
+      /*outShardings=*/getShardingPerValue(callOp));
+  namedCompOp->setAttrs(namedCompAttrs);
+
+  Region& namedCompRegion = namedCompOp.getRegion();
+  if (auto movedRegionIt = calleeNameToMovedRegion.find(calleeName);
+      movedRegionIt != calleeNameToMovedRegion.end()) {
+    static llvm::once_flag onceFlag;
+    emitOpWarningOnce(
+        onceFlag, callOp,
+        llvm::formatv("uninlineable function @{0} has multiple call ops, we "
+                      "need to clone the function body for each call",
+                      calleeName)
+            .str());
+    rewriter.cloneRegionBefore(*movedRegionIt->second, namedCompRegion,
+                               namedCompRegion.begin());
+  } else {
+    FuncOp funcOp = symbolTable.lookup<FuncOp>(calleeName);
+    assert(funcOp &&
+           ("Failed to lookup function: " + std::string(calleeName)).c_str());
+    inlineRegionAndConvertTerminatorOp<ReturnOp>(funcOp.getBody(),
+                                                 namedCompRegion);
+    calleeNameToMovedRegion[calleeName] = &namedCompRegion;
+  }
+
+  rewriter.replaceOp(callOp, namedCompOp);
+}
+
+class ImportUninlineableFuncCallsPass
+    : public PassWrapper<ImportUninlineableFuncCallsPass,
+                         OperationPass<ModuleOp>> {
+ public:
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(ImportUninlineableFuncCallsPass)
+
+  void runOnOperation() final {
+    ModuleOp moduleOp = getOperation();
+    IRRewriter rewriter(moduleOp.getContext());
+    SymbolTable symbolTable(moduleOp);
+    // For every callee name, the first CallOp encountered with that symbol will
+    // move the body of the callee into the created NamedComputationOp, and map
+    // the symbol name to the moved region. Subsequent CallOps with that symbol
+    // will clone the mapped region.
+    llvm::SmallDenseMap<StringRef, Region*> calleeNameToMovedRegion;
+
+    moduleOp->walk([&](CallOp op) {
+      if (isInlineableCallOp(op)) {
+        return;
+      }
+      importCallOp(op, calleeNameToMovedRegion, rewriter, symbolTable);
+    });
+
+    // Erase all func ops that now have no call ops.
+    for (auto [calleeName, _] : calleeNameToMovedRegion) {
+      symbolTable.erase(symbolTable.lookup(calleeName));
+    }
+  }
+
+  StringRef getArgument() const override {
+    return "xla-sdy-import-uninlineable-func-calls";
+  }
+
+  StringRef getDescription() const override {
+    return "Creates a pass that converts a `CallOp` with a `backend_config` "
+           "or `inlineable=false` frontend attr to a `NamedComputationOp` with "
+           "the function body inlined and the name of the callee.";
+  }
+};
+
+}  // namespace
+
+std::unique_ptr<Pass> createImportUninlineableFuncCallsPass() {
+  return std::make_unique<ImportUninlineableFuncCallsPass>();
+}
+
+void registerImportUninlineableFuncCallsPass() {
+  registerPass(createImportUninlineableFuncCallsPass);
+}
+
+}  // namespace sdy
+}  // namespace mlir

shardy/round_trip_import/import_uninlineable_func_calls.cc

@@ -13,8 +13,8 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
-#ifndef SHARDY_ROUND_TRIP_IMPORT_IMPORT_BACKEND_FUNC_CALLS_H_
-#define SHARDY_ROUND_TRIP_IMPORT_IMPORT_BACKEND_FUNC_CALLS_H_
+#ifndef SHARDY_ROUND_TRIP_IMPORT_IMPORT_UNINLINEABLE_FUNC_CALLS_H_
+#define SHARDY_ROUND_TRIP_IMPORT_IMPORT_UNINLINEABLE_FUNC_CALLS_H_
 
 #include <memory>
 
@@ -23,19 +23,21 @@ limitations under the License.
 namespace mlir {
 namespace sdy {
 
-// Creates a pass that converts a `CallOp` with a `backend_config` attr to a
-// `NamedComputationOp` with the function body inlined and name of the callee.
+// Creates a pass that converts a `CallOp` with a `backend_config` or
+// `inlineable=false` frontend attr to a `NamedComputationOp` with the function
+// body inlined and name of the callee.
 //
-// This pass is used to handle host offloading calls which are non inlined
-// functions that require the callee to be propagated through.
+// This pass is used to handle host offloading and GPU stream calls which are
+// non inlined functions that require the callee to be propagated through.
 //
-// NOTE: it assumes that there is a unique callee for each caller.
-std::unique_ptr<Pass> createImportBackendFuncCallsPass();
+// NOTE: In case there are multiple call ops for the same callee, we will clone
+// the function body for each call op and emit a warning.
+std::unique_ptr<mlir::Pass> createImportUninlineableFuncCallsPass();
 
-// Register the xla-sdy-import-backend-func-calls pass.
-void registerImportBackendFuncCallsPass();
+// Register the xla-sdy-import-uninlineable-calls pass.
+void registerImportUninlineableFuncCallsPass();
 
 }  // namespace sdy
 }  // namespace mlir
 
-#endif  // SHARDY_ROUND_TRIP_IMPORT_IMPORT_BACKEND_FUNC_CALLS_H_
+#endif  // SHARDY_ROUND_TRIP_IMPORT_IMPORT_UNINLINEABLE_FUNC_CALLS_H_