xref: /aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/service/gpu/xlir_ops.h (revision b6fb3261f9314811a0f4371741dbb8839866f948)
1 /*
2  * Copyright 2021 The TensorFlow Authors
3  *
4  * Licensed under the Apache License, Version 2.0 (the "License");
5  * you may not use this file except in compliance with the License.
6  * You may obtain a copy of the License at
7  *
8  *      http://www.apache.org/licenses/LICENSE-2.0
9  *
10  * Unless required by applicable law or agreed to in writing, software
11  * distributed under the License is distributed on an "AS IS" BASIS,
12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13  * See the License for the specific language governing permissions and
14  * limitations under the License.
15  */
16 
17 // MLIR op definitions for xlir_ops library
18 //
19 // This file declares the 'xlir' dialect as well as the operators that make up
20 // the xlir_ops library.
21 
22 #ifndef TENSORFLOW_COMPILER_XLA_SERVICE_GPU_XLIR_OPS_H_
23 #define TENSORFLOW_COMPILER_XLA_SERVICE_GPU_XLIR_OPS_H_
24 
25 #include "mlir/IR/Dialect.h"
26 #include "mlir/IR/OpDefinition.h"
27 #include "mlir/Interfaces/InferTypeOpInterface.h"
28 #include "mlir/Interfaces/SideEffectInterfaces.h"
29 #include "tfrt/gpu/kernels/gpu_ops.h"  // from @tf_runtime
30 #include "tfrt/basic_kernels/opdefs/basic_kernels.h"  // from @tf_runtime
31 #include "tfrt/tensor/opdefs/host_tensor.h"  // from @tf_runtime
32 #include "tfrt/tensor/opdefs/tensor.h"  // from @tf_runtime
33 #include "tfrt/tensor/opdefs/tensor_shape.h"  // from @tf_runtime
34 
35 namespace xla {
36 namespace gpu {
37 
38 // Dialect for XLIR operations.
39 class XlirDialect : public mlir::Dialect {
40  public:
getDialectNamespace()41   static llvm::StringRef getDialectNamespace() { return "xlir"; }
42   explicit XlirDialect(mlir::MLIRContext* context);
43 };
44 
45 }  // namespace gpu
46 }  // namespace xla
47 
48 // TableGen'd declarations
49 #define GET_OP_CLASSES
50 #include "tensorflow/compiler/xla/service/gpu/xlir_opdefs.h.inc"
51 
52 #endif  // TENSORFLOW_COMPILER_XLA_SERVICE_GPU_XLIR_OPS_H_
53