| #include <string> |
| |
| #include "third_party/absl/status/status.h" |
| #include "third_party/absl/status/statusor.h" |
| #include "third_party/llvm/llvm-project/mlir/include/mlir/IR/Attributes.h" |
| #include "third_party/llvm/llvm-project/mlir/include/mlir/IR/BuiltinAttributes.h" |
| #include "third_party/llvm/llvm-project/mlir/include/mlir/IR/Operation.h" |
| |
| namespace xla::mosaic { |
| |
| absl::StatusOr<std::string> CustomCallKernelName(mlir::Operation* op) { |
| mlir::Attribute kernel_name_attr = op->getAttr("kernel_name"); |
| auto kernel_name = kernel_name_attr.dyn_cast_or_null<mlir::StringAttr>(); |
| if (!kernel_name) { |
| return absl::InvalidArgumentError("Cannot find kernel_name"); |
| } |
| |
| return kernel_name.str(); |
| } |
| |
| } // namespace xla::mosaic |
