xref: /llvm-project/mlir/lib/Dialect/Func/Extensions/MeshShardingExtensions.cpp (revision adbf21f12b3069b2554efb39f2e92c6cf6f24940) 
1 //===- MeshShardingExtensions.cpp - ---------------------------------------===//
2 //
3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4 // See https://llvm.org/LICENSE.txt for license information.
5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6 //
7 //===----------------------------------------------------------------------===//
8 
9 #include "mlir/Dialect/Func/Extensions/MeshShardingExtensions.h"
10 #include "mlir/Dialect/Func/IR/FuncOps.h"
11 #include "mlir/Dialect/Mesh/Interfaces/ShardingInterfaceImpl.h"
12 #include "mlir/IR/MLIRContext.h"
13 
14 namespace mlir::func {
15 
registerShardingInterfaceExternalModels(DialectRegistry & registry)16 void registerShardingInterfaceExternalModels(DialectRegistry &registry) {
17   registry.addExtension(+[](MLIRContext *ctx, FuncDialect *dialect) {
18     ReturnOp::attachInterface<
19         mesh::IndependentParallelIteratorDomainShardingInterface<ReturnOp>>(
20         *ctx);
21   });
22 }
23 
24 } // namespace mlir::func
25