1 /* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
2 Licensed under the Apache License, Version 2.0 (the "License");
3 you may not use this file except in compliance with the License.
4 You may obtain a copy of the License at
5 http://www.apache.org/licenses/LICENSE-2.0
6 Unless required by applicable law or agreed to in writing, software
7 distributed under the License is distributed on an "AS IS" BASIS,
8 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
9 See the License for the specific language governing permissions and
10 limitations under the License.
11 ==============================================================================*/
12
13 #include "mlir-c/Bindings/Python/Interop.h"
14 #include "mlir-c/Registration.h"
15 #include "mlir-hlo-c/Dialects.h"
16 #include "mlir/Bindings/Python/PybindAdaptors.h"
17
18 namespace py = pybind11;
19
PYBIND11_MODULE(_mlirHlo,m)20 PYBIND11_MODULE(_mlirHlo, m) {
21 m.doc() = "mlir-hlo main python extension";
22
23 m.def(
24 "register_mhlo_dialect",
25 [](MlirContext context, bool load) {
26 MlirDialectHandle mhloDialect = mlirGetDialectHandle__mhlo__();
27 mlirDialectHandleRegisterDialect(mhloDialect, context);
28 if (load) {
29 mlirDialectHandleLoadDialect(mhloDialect, context);
30 }
31 },
32 py::arg("context"), py::arg("load") = true);
33
34 m.def(
35 "register_chlo_dialect",
36 [](MlirContext context, bool load) {
37 MlirDialectHandle chloDialect = mlirGetDialectHandle__chlo__();
38 mlirDialectHandleRegisterDialect(chloDialect, context);
39 if (load) {
40 mlirDialectHandleLoadDialect(chloDialect, context);
41 }
42 },
43 py::arg("context"), py::arg("load") = true);
44 }
45