| 33 | namespace { |
| 34 | |
| 35 | void PrintSupportedOps(const string& device, const string& regen_run) { |
| 36 | XlaOpRegistry::RegisterCompilationKernels(); |
| 37 | |
| 38 | std::vector<const KernelDef*> kdefs = |
| 39 | XlaOpRegistry::DeviceKernels(device, |
| 40 | /*include_compilation_only_kernels=*/true); |
| 41 | std::sort( |
| 42 | kdefs.begin(), kdefs.end(), |
| 43 | [](const KernelDef* a, const KernelDef* b) { return a->op() < b->op(); }); |
| 44 | |
| 45 | std::cout << "**Supported operators for device: " << device << "**\n\n" |
| 46 | << "Operator | Type Constraint\n" |
| 47 | << "-------- | ---------------" << std::endl; |
| 48 | for (const KernelDef* kdef : kdefs) { |
| 49 | std::vector<string> constraints; |
| 50 | for (const KernelDef::AttrConstraint& constraint : kdef->constraint()) { |
| 51 | std::vector<string> types; |
| 52 | for (int type : constraint.allowed_values().list().type()) { |
| 53 | types.push_back(DataTypeString(static_cast<DataType>(type))); |
| 54 | } |
| 55 | std::sort(types.begin(), types.end()); |
| 56 | constraints.push_back("`" + constraint.name() + "={" + |
| 57 | absl::StrJoin(types, ",") + "}`"); |
| 58 | } |
| 59 | std::cout << "`" << kdef->op() << "` | " |
| 60 | << absl::StrJoin(constraints, "<br>") << std::endl; |
| 61 | } |
| 62 | |
| 63 | std::cout << "\nTo regenerate this table, run:\n\n```shell\n" |
| 64 | << regen_run << " --device=" << device << "\n```" << std::endl; |
| 65 | } |
| 66 | |
| 67 | } // namespace |
| 68 |
no test coverage detected