MCPcopy Create free account
hub / github.com/ml-explore/mlx / append_binary_kernels

Function append_binary_kernels

mlx/backend/metal/jit_kernels.cpp:54–102  ·  view source on GitHub ↗

Source from the content-addressed store, hash-verified

52}
53
54void append_binary_kernels(
55 const std::string& lib_name,
56 Dtype in_type,
57 Dtype out_type,
58 const char* op,
59 std::string& kernel_source) {
60 const std::array<std::pair<std::string, std::string>, 7> kernel_types = {{
61 {"ss", "binary_ss"},
62 {"vs2", "binary_vs2"},
63 {"sv2", "binary_sv2"},
64 {"vv2", "binary_vv2"},
65 {"g1large", "binary_g_nd1"},
66 {"g2large", "binary_g_nd2"},
67 {"g3large", "binary_g_nd3"},
68 }};
69 auto in_t = get_type_string(in_type);
70 auto out_t = get_type_string(out_type);
71
72 for (auto& [name, func] : kernel_types) {
73 kernel_source +=
74 get_template_definition(name + "_" + lib_name, func, in_t, out_t, op);
75 }
76 kernel_source += get_template_definition(
77 "vs_" + lib_name, "binary_vs", in_t, out_t, op, 1);
78 kernel_source += get_template_definition(
79 "sv_" + lib_name, "binary_sv", in_t, out_t, op, 1);
80 kernel_source += get_template_definition(
81 "vv_" + lib_name, "binary_vv", in_t, out_t, op, 1);
82
83 if (get_work_per_thread(in_type) > 1) {
84 kernel_source += get_template_definition(
85 "vsn_" + lib_name, "binary_vs", in_t, out_t, op);
86 kernel_source += get_template_definition(
87 "svn_" + lib_name, "binary_sv", in_t, out_t, op);
88 kernel_source += get_template_definition(
89 "vvn_" + lib_name, "binary_vv", in_t, out_t, op);
90 }
91
92 kernel_source += get_template_definition(
93 "g1_" + lib_name, "binary_g_nd1", in_t, out_t, op, "int");
94 kernel_source += get_template_definition(
95 "g2_" + lib_name, "binary_g_nd2", in_t, out_t, op, "int");
96 kernel_source += get_template_definition(
97 "g3_" + lib_name, "binary_g_nd3", in_t, out_t, op, "int");
98 kernel_source += get_template_definition(
99 "gn2_" + lib_name, "binary_g", in_t, out_t, op, 2, "int");
100 kernel_source += get_template_definition(
101 "gn4large_" + lib_name, "binary_g", in_t, out_t, op, 4);
102}
103
104MTL::ComputePipelineState* get_binary_kernel(
105 metal::Device& d,

Callers 2

get_binary_kernelFunction · 0.85
get_binary_two_kernelFunction · 0.85

Calls 3

get_type_stringFunction · 0.85
get_template_definitionFunction · 0.85
get_work_per_threadFunction · 0.85

Tested by

no test coverage detected