Add support for custom sharding rules via ShardingRuleOpInterface

tenstorrent · tenstorrent · commit 4eb2bcd900fd · 2025-11-14T18:59:23.000Z
Allows custom ops to supply their own sharding rules when not covered
by the default registry.
diff --git a/shardy/dialect/sdy/transforms/propagation/op_sharding_rule_registry.cc b/shardy/dialect/sdy/transforms/propagation/op_sharding_rule_registry.cc
@@ -687,6 +687,11 @@ OpShardingRuleAttr createOpShardingRule(Operation* op,
           }
           return OpShardingRuleBuilder::buildPointwise(customCall);
         }
+        if (auto shardingRuleOp =
+              llvm::dyn_cast<ShardingRuleOpInterface>(
+                customCall.getOperation())) {
+          return shardingRuleOp.getShardingRule();
+        }
         // TODO(b/327191011): output unregistered op stats instead.
         static llvm::once_flag onceFlag;
         emitOpWarningOnce(

Original file line number	Diff line number	Diff line change
`@@ -687,6 +687,11 @@ OpShardingRuleAttr createOpShardingRule(Operation* op,`
`687`	`687`	`}`
`688`	`688`	`return OpShardingRuleBuilder::buildPointwise(customCall);`
`689`	`689`	`}`
	`690`	`+ if (auto shardingRuleOp =`
	`691`	`+ llvm::dyn_cast<ShardingRuleOpInterface>(`
	`692`	`+ customCall.getOperation())) {`
	`693`	`+ return shardingRuleOp.getShardingRule();`
	`694`	`+ }`
`690`	`695`	`// TODO(b/327191011): output unregistered op stats instead.`
`691`	`696`	`static llvm::once_flag onceFlag;`
`692`	`697`	`emitOpWarningOnce(`