xref: /aosp_15_r20/external/pytorch/aten/src/ATen/LinalgBackend.h (revision da0073e96a02ea20f0ac840b70461e3646d07c45)
1*da0073e9SAndroid Build Coastguard Worker #pragma once
2*da0073e9SAndroid Build Coastguard Worker 
3*da0073e9SAndroid Build Coastguard Worker #include <c10/util/Exception.h>
4*da0073e9SAndroid Build Coastguard Worker 
5*da0073e9SAndroid Build Coastguard Worker #include <ostream>
6*da0073e9SAndroid Build Coastguard Worker #include <string>
7*da0073e9SAndroid Build Coastguard Worker 
8*da0073e9SAndroid Build Coastguard Worker namespace at {
9*da0073e9SAndroid Build Coastguard Worker 
10*da0073e9SAndroid Build Coastguard Worker enum class LinalgBackend : int8_t { Default, Cusolver, Magma };
11*da0073e9SAndroid Build Coastguard Worker 
LinalgBackendToString(at::LinalgBackend backend)12*da0073e9SAndroid Build Coastguard Worker inline std::string LinalgBackendToString(at::LinalgBackend backend) {
13*da0073e9SAndroid Build Coastguard Worker   switch (backend) {
14*da0073e9SAndroid Build Coastguard Worker     case LinalgBackend::Default:
15*da0073e9SAndroid Build Coastguard Worker       return "at::LinalgBackend::Default";
16*da0073e9SAndroid Build Coastguard Worker     case LinalgBackend::Cusolver:
17*da0073e9SAndroid Build Coastguard Worker       return "at::LinalgBackend::Cusolver";
18*da0073e9SAndroid Build Coastguard Worker     case LinalgBackend::Magma:
19*da0073e9SAndroid Build Coastguard Worker       return "at::LinalgBackend::Magma";
20*da0073e9SAndroid Build Coastguard Worker     default:
21*da0073e9SAndroid Build Coastguard Worker       TORCH_CHECK(false, "Unknown linalg backend");
22*da0073e9SAndroid Build Coastguard Worker   }
23*da0073e9SAndroid Build Coastguard Worker }
24*da0073e9SAndroid Build Coastguard Worker 
25*da0073e9SAndroid Build Coastguard Worker inline std::ostream& operator<<(
26*da0073e9SAndroid Build Coastguard Worker     std::ostream& stream,
27*da0073e9SAndroid Build Coastguard Worker     at::LinalgBackend backend) {
28*da0073e9SAndroid Build Coastguard Worker   return stream << LinalgBackendToString(backend);
29*da0073e9SAndroid Build Coastguard Worker }
30*da0073e9SAndroid Build Coastguard Worker 
31*da0073e9SAndroid Build Coastguard Worker } // namespace at
32