#sdy Add custom-call sharding rule registry.

PiperOrigin-RevId: 658549415

shardy/dialect/sdy/transforms/propagation/BUILD

@@ -263,3 +263,33 @@ cc_test(
         "@llvm-project//mlir:Support",
     ],
 )
+
+cc_library(
+    name = "custom_call_sharding_rule_registry",
+    srcs = ["custom_call_sharding_rule_registry.cc"],
+    hdrs = ["custom_call_sharding_rule_registry.h"],
+    deps = [
+        "//shardy/dialect/sdy/ir:dialect",
+        "//shardy/dialect/sdy/transforms/common:macros",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:Support",
+    ],
+)
+
+cc_test(
+    name = "custom_call_sharding_rule_registry_test",
+    srcs = ["custom_call_sharding_rule_registry_test.cc"],
+    deps = [
+        ":custom_call_sharding_rule_registry",
+        ":op_sharding_rule_registry",
+        "//shardy/dialect/sdy/ir:dialect",
+        "@com_google_googletest//:gtest_main",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:FuncDialect",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:Parser",
+        "@llvm-project//mlir:Support",
+        "@stablehlo//:stablehlo_ops",
+    ],
+)

shardy/dialect/sdy/transforms/propagation/custom_call_sharding_rule_registry.cc

@@ -0,0 +1,58 @@
+/* Copyright 2024 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/dialect/sdy/transforms/propagation/custom_call_sharding_rule_registry.h"
+
+#include <memory>
+#include <optional>
+#include <string>
+#include <utility>
+
+#include "llvm/Support/ManagedStatic.h"
+#include "llvm/Support/Mutex.h"
+#include "mlir/Support/LogicalResult.h"
+
+namespace mlir {
+namespace sdy {
+
+namespace {
+static llvm::ManagedStatic<CustomCallShardingRuleRegistry> cc_registry;
+}  // namespace
+
+std::optional<CustomCallShardingRuleRegistry::ShardingRuleCallBack>
+CustomCallShardingRuleRegistry::GetShardingRuleCallBack(
+    const std::string& op_name) {
+  llvm::sys::ScopedLock scopedLock(cc_registry->mutex_);
+  if (auto iter = cc_registry->name_to_call_back_.find(op_name);
+      iter != cc_registry->name_to_call_back_.end()) {
+    return *iter->second;
+  }
+  return std::nullopt;
+}
+
+LogicalResult CustomCallShardingRuleRegistry::RegisterShardingRuleCallBack(
+    const std::string& op_name, ShardingRuleCallBack call_back_func) {
+  llvm::sys::ScopedLock scopedLock(cc_registry->mutex_);
+  auto [it, emplaced] = cc_registry->name_to_call_back_.try_emplace(
+      op_name,
+      std::make_unique<ShardingRuleCallBack>(std::move(call_back_func)));
+  if (!emplaced) {
+    return failure();
+  }
+  return success();
+}
+
+}  // namespace sdy
+}  // namespace mlir

shardy/dialect/sdy/transforms/propagation/custom_call_sharding_rule_registry.h

@@ -0,0 +1,60 @@
+/* Copyright 2024 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.
+==============================================================================*/
+
+#ifndef SHARDY_DIALECT_SDY_TRANSFORMS_PROPAGATION_CUSTOM_CALL_SHARDING_RULE_REGISTRY_H_
+#define SHARDY_DIALECT_SDY_TRANSFORMS_PROPAGATION_CUSTOM_CALL_SHARDING_RULE_REGISTRY_H_
+
+#include <functional>
+#include <optional>
+
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/Support/Mutex.h"
+#include "mlir/IR/Operation.h"
+#include "mlir/Support/LogicalResult.h"
+#include "shardy/dialect/sdy/ir/dialect.h"
+namespace mlir {
+namespace sdy {
+
+// A registry for custom-call sharding rule information.
+//
+// This registry records the name of a custom-call and its corresponding
+// callback function that inspects the custom-call op to produce an
+// `OpShardingRuleAttr`.
+//
+// A process can have at most one globally shared instance of this object. This
+// instance is created lazily the first time the registry is accessed.
+//
+// This class is thread-safe.
+class CustomCallShardingRuleRegistry {
+ public:
+  using ShardingRuleCallBack =
+      std::function<OpShardingRuleAttr(mlir::Operation*)>;
+
+  static LogicalResult RegisterShardingRuleCallBack(
+      const std::string& op_name, ShardingRuleCallBack call_back_func);
+
+  static std::optional<ShardingRuleCallBack> GetShardingRuleCallBack(
+      const std::string& op_name);
+
+ private:
+  llvm::sys::Mutex mutex_;
+  std::unordered_map<std::string, std::unique_ptr<ShardingRuleCallBack>>
+      name_to_call_back_;
+};
+
+}  // namespace sdy
+}  // namespace mlir
+
+#endif  // SHARDY_DIALECT_SDY_TRANSFORMS_PROPAGATION_CUSTOM_CALL_SHARDING_RULE_REGISTRY_H_

shardy/dialect/sdy/transforms/propagation/custom_call_sharding_rule_registry_test.cc

@@ -0,0 +1,58 @@
+/* Copyright 2024 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/dialect/sdy/transforms/propagation/custom_call_sharding_rule_registry.h"
+
+#include <optional>
+
+#include "mlir/IR/Operation.h"
+#include "mlir/Support/LogicalResult.h"
+#include "shardy/dialect/sdy/ir/dialect.h"
+#include <gtest/gtest.h>
+
+namespace mlir {
+namespace sdy {
+
+namespace {
+
+using CustomCallShardingRuleRegistryTest = ::testing::Test;
+
+TEST_F(CustomCallShardingRuleRegistryTest, SimplyRegister) {
+  auto callBackFunc = [](mlir::Operation* op) { return OpShardingRuleAttr(); };
+  const char kCustomCallOpName[] = "sdy_testonly1";
+
+  std::optional<CustomCallShardingRuleRegistry::ShardingRuleCallBack>
+      queryResult = CustomCallShardingRuleRegistry::GetShardingRuleCallBack(
+          kCustomCallOpName);
+  EXPECT_FALSE(queryResult.has_value());
+
+  LogicalResult registerResult =
+      CustomCallShardingRuleRegistry::RegisterShardingRuleCallBack(
+          kCustomCallOpName, callBackFunc);
+  EXPECT_TRUE(registerResult.succeeded());
+
+  queryResult = CustomCallShardingRuleRegistry::GetShardingRuleCallBack(
+      kCustomCallOpName);
+  EXPECT_TRUE(queryResult.has_value());
+
+  registerResult = CustomCallShardingRuleRegistry::RegisterShardingRuleCallBack(
+      kCustomCallOpName, callBackFunc);
+  EXPECT_TRUE(registerResult.failed());
+}
+
+}  // namespace
+
+}  // namespace sdy
+}  // namespace mlir