MCPcopy Create free account
hub / github.com/MegEngine/MegEngine / gen_fastdiv_offset

Function gen_fastdiv_offset

src/jit/impl/nvrtc/codegen_cuda.cpp:30–68  ·  view source on GitHub ↗

Source from the content-addressed store, hash-verified

28}
29
30std::string gen_fastdiv_offset(size_t nr_inps) {
31 std::string res = "";
32
33 char tmp[100];
34#define APPEND(fmt...) \
35 do { \
36 snprintf(tmp, sizeof(tmp), fmt); \
37 res += tmp; \
38 } while (0)
39
40 for (size_t i = 0; i < nr_inps; ++i) {
41 APPEND("offset_%zu = 0;\n", i);
42 }
43 for (size_t i = 0; i < nr_inps; ++i) {
44 APPEND("tmp_idx = global_idx;\n");
45 APPEND("#pragma unroll\n");
46 APPEND("for (int j = {{NDIM}} - 1; j >= 1; --j) {\n");
47 APPEND("Uint32Fastdiv& shp = "
48 "visitors.m[%zu].m_shape_highdim[j-1];\n",
49 i);
50 res += R"(
51 unsigned int
52 ans_for_one = tmp_idx & ~shp.m_divisor_is_not_1,
53 dfix = tmp_idx + shp.m_inc_dividend,
54 hi32 = __umulhi(dfix, shp.m_mul),
55 ans = hi32 >> shp.m_shift,
56 idx_div = (ans & shp.m_divisor_is_not_1) | ans_for_one;
57 )";
58 APPEND("offset_%zu += (tmp_idx - idx_div * shp.m_divisor) * "
59 "visitors.m[%zu].m_stride[j];\n",
60 i, i);
61 APPEND("tmp_idx = idx_div;\n");
62 APPEND("}\n");
63 APPEND("offset_%zu += tmp_idx * visitors.m[%zu].m_stride[0];\n", i, i);
64 }
65
66#undef APPEND
67 return res;
68}
69
70ASTPtr gen_data_ast(size_t input_id, const JITExecutor::Args::Data& n) {
71 auto res = ssprintf(

Callers 1

codegen_cudaMethod · 0.85

Calls

no outgoing calls

Tested by

no test coverage detected