-
Notifications
You must be signed in to change notification settings - Fork 7
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
#sdy Add custom-call sharding rule registry.
PiperOrigin-RevId: 658549415
- Loading branch information
1 parent
541bb58
commit 5b4c52d
Showing
4 changed files
with
205 additions
and
0 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
57 changes: 57 additions & 0 deletions
57
shardy/dialect/sdy/transforms/propagation/custom_call_sharding_rule_registry.cc
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,57 @@ | ||
/* 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 "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 |
60 changes: 60 additions & 0 deletions
60
shardy/dialect/sdy/transforms/propagation/custom_call_sharding_rule_registry.h
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -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_ |
58 changes: 58 additions & 0 deletions
58
shardy/dialect/sdy/transforms/propagation/custom_call_sharding_rule_registry_test.cc
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -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 |