郑州建设网站企业使用session和cookie实现网站自动登录 .net

郑州建设网站企业,使用session和cookie实现网站自动登录 .net,新闻资讯app制作公司,建立有效的()0x00 概要此处的”转译系统“包含两部分#xff1a;把计算图转换为任务图。将 Mirage 生成的#xff08;优化过的#xff09;计算图转换为高效的 CUDA 代码0x01 Task和Event在 Mirage 持久化内核#xff08;Persistent Kernel#xff09;的设计与实现中#xff0c;需突破…0x00 概要此处的”转译系统“包含两部分把计算图转换为任务图。将 Mirage 生成的优化过的计算图转换为高效的 CUDA 代码0x01 Task和Event在 Mirage 持久化内核Persistent Kernel的设计与实现中需突破三个关键技术瓶颈如何将抽象算子转化为可执行任务。如何处理任务间的数据依赖。如何高效分配任务至 GPU 计算单元。这三个问题的解决直接决定了内核能否充分发挥 GPU 并行性能适配复杂张量计算场景如大语言模型推理。Mirage 通过引入Task和Event与三层图一起来解决上述问题Kernel Graph 定义张量数据流Block Graph 定义内存访问模式Task 执行具体计算Event 管理任务依赖关系Thread Graph 执行底层并行计算1.1 可执行任务GPU 执行 CUDA 或 Triton 代码时需将算子的整体计算逻辑切分为多个 “计算块”Block—— 每个计算块对应 GPU 流式多处理器SM可承载的基本计算单元最终由调度系统分配至不同 SM 并行执行。基于这一硬件特性Mirage 持久化内核将 “单个计算块的计算” 定义为最小任务单元Task实现算子到任务的结构化转化。1.1.1 任务定义任务的由TaskDesc 来实现。struct TaskDesc {TaskDesc(TaskType t, int _variant_id): task_type(t), variant_id(_variant_id), num_inputs(0), num_outputs(0),trigger_event(EVENT_INVALID_ID), dependent_event(EVENT_INVALID_ID) {}TaskDesc() {}TaskType task_type; // 任务类型unsigned variant_id; // 变体IDint num_inputs, num_outputs;EventId trigger_event; // 触发事件EventId dependent_event; // 依赖事件TensorDesc inputs[MAX_INPUTS_PER_TASK]; // 张量描述TensorDesc outputs[MAX_OUTPUTS_PER_TASK];};1.1.2 任务类型任务类型如下enum TaskType {TASK_TERMINATE 0, // 终止任务TASK_BEGIN_TASK_GRAPH 10, // 人物图开始标记// compute task starts from 100TASK_EMBEDDING 101, // 嵌入层TASK_RMS_NORM_LINEAR 102, // RMS归一化和线性层组合TASK_ATTENTION_1 103, // 注意力机制第一部分TASK_ATTENTION_2 104, // 注意力机制第二部分TASK_SILU_MUL_LINEAR_WITH_RESIDUAL 105,TASK_ALLREDUCE 106,TASK_REDUCE 107,TASK_LINEAR_WITH_RESIDUAL 108,TASK_ARGMAX 109,TASK_ARGMAX_PARTIAL 110,TASK_ARGMAX_REDUCE 111,TASK_FIND_NGRAM_PARTIAL 112, //部分n-gram查找TASK_FIND_NGRAM_GLOBAL 113, // 全局n-gram查找TASK_TARGET_VERIFY_GREEDY 114, // 贪心目标验证TASK_SINGLE_BATCH_EXTEND_ATTENTION 115,TASK_NVSHMEM_COPY 199, // 使用NVSHMEM进行跨GPU的数据复制TASK_SCHD_TASKS 200, // 调度任务TASK_SCHD_EVENTS 201, // 调度事件TASK_GET_EVENT 202, // 获取事件TASK_GET_NEXT_TASK 203, // 获取任务};1.2 事件传统内核设计中数据依赖关系以算子为单位定义 —— 只有前一个算子的所有计算完全结束后一个算子才能启动这种粗粒度依赖会导致大量计算资源闲置例如前一算子仅剩余少量计算未完成时后一算子需持续等待。Mirage 持久化内核将依赖关系下沉至任务级别实现更精细的并行调度。具体而言算子级依赖会被拆解为任务间的依赖链即事件。1.2.1 事件定义事件的由 EventDesc 来实现。struct EventDesc {EventDesc(void): event_type(EVENT_INVALID), num_triggers(0),first_task_id(TASK_INVALID_ID), last_task_id(TASK_INVALID_ID) {}EventDesc(EventType type, int nt, TaskId f, TaskId l): event_type(type), num_triggers(nt), first_task_id(f), last_task_id(l) {}EventType event_type;int num_triggers; // 触发器数量TaskId first_task_id, last_task_id; // 首尾任务ID范围};1.2.2 事件类型事件类型如下enum EventType {EVENT_EMPTY 900, // 空事件EVENT_LAUNCH_TASKS 901, // 启动任务EVENT_LAUNCH_MASSIVE_TASKS 902, // 启动大规模任务EVENT_LAUNCH_DEPENDENT_TASKS 903, // 启动依赖任务EVENT_END_OF_TASK_GRAPH 910, // 任务图结束EVENT_TERMINATION 911, // 终止事件EVENT_INVALID 999, //无效事件};下图展示了如何确定事件类型。mirage-4-10x02 生成CUDA代码TaskDesc 结构体本身并不直接包含可执行代码。它更像是一个任务的描述符或配置信息包含了任务执行所需的一些元数据。2.1 生成代码实际的可执行代码是通过以下方式来生成的。register_muggraph在 runtime.cc 的 register_mugraph 函数中会遍历 Graph 中的 KN_CUSTOMIZED_OP 操作符。对于每个操作符它会从 task_configs即 Graph::task_config中查找对应的配置输入数、输出数、TaskType, variant_id。创建 TaskDesc 结构体会将获取到的 TaskType 和 variant_id 填入 TaskDesc。在生成计算图时候会调用 register_task实际上是生成CUDA代码比如def embed_layer(self,input: DTensor, # [batch_size, num_spec_tokens]weight: DTensor, # [vocab_size, hidden_size]output: DTensor, # [batch_size, hidden_size]grid_dim: tuple,block_dim: tuple,input_source: int 0, # 0: all_tokens, 1: input_token):tb_graph TBGraph(CyTBGraph(grid_dim, block_dim, 1, 64))tb_graph.new_input(input, (-1, 1, -1), -1, True)tb_graph.new_input(weight, (1, -1, -1), -1, True)tb_graph.new_input(output, (1, 0, -1), -1, True)self.kn_graph.customized([input, weight, output], tb_graph)# 会生成CUDA代码self.kn_graph.register_task(tb_graph, embedding, [input_source])当用户调用 Graph::register_task 时它会获取当前图中最后一个操作符必须是 KN_CUSTOMIZED_OP根据传入的 task_type 字符串和参数调用 TaskRegister 对应的 register_*_task 函数。注册成功后它会将任务的输入/输出数量、TaskType 和 variant_id 存储在 Graph 的 task_config 映射中以 KNOperator* 为键。register_task的实现位于graph.cc具体代码如下void Graph::register_task(char const *task_type, std::vectorint params) {std::string name std::string(task_type);KNOperator const *op operators.back();assert(op-op_type type::KN_CUSTOMIZED_OP);KNCustomizedOp const *customized static_castKNCustomizedOp const *(op);TaskRegister *task_register TaskRegister::get_instance();if (name embedding) {int variant_id task_register-register_embedding_task(customized-bgraph, params);task_config[op] std::make_tuple(2, 1, TASK_EMBEDDING, variant_id);} else if (name rmsnorm_linear) {int variant_id task_register-register_rmsnorm_linear_task(customized-bgraph, params);task_config[op] std::make_tuple(3, 1, TASK_RMS_NORM_LINEAR, variant_id);} else if (name attention) {int variant_id task_register-register_attention_task(customized-bgraph, params);task_config[op] std::make_tuple(7, 1, TASK_ATTENTION_1, variant_id);} else if (name single_batch_extend_attention) {int variant_id task_register-register_single_batch_extend_attention_task(customized-bgraph, params);task_config[op] std::make_tuple(7, 1, TASK_SINGLE_BATCH_EXTEND_ATTENTION, variant_id);} else if (name linear_with_residual) {int variant_id task_register-register_linear_with_residual_task(customized-bgraph, params);task_config[op] std::make_tuple(3, 1, TASK_LINEAR_WITH_RESIDUAL, variant_id);} else if (name silu_mul_linear_with_residual) {int variant_id task_register-register_silu_mul_linear_with_residual_task(customized-bgraph, params);task_config[op] std::make_tuple(3, 1, TASK_SILU_MUL_LINEAR_WITH_RESIDUAL, variant_id);} else if (name argmax) {task_config[op] std::make_tuple(1, 1, TASK_ARGMAX, 0);} else if (name argmax_partial) {int variant_id task_register-register_arrrgmax_partial_task(customized-bgraph, params);task_config[op] std::make_tuple(1, 2, TASK_ARGMAX_PARTIAL, variant_id);} else if (name argmax_reduce) {int variant_id task_register-register_argmax_reduce_task(customized-bgraph, params);task_config[op] std::make_tuple(2, 1, TASK_ARGMAX_REDUCE, variant_id);} else if (name allreduce) {task_config[op] std::make_tuple(2, 1, TASK_ALLREDUCE, 0);} else if (name find_ngram_partial) {int variant_id task_register-register_find_ngram_partial_task(customized-bgraph, params);task_config[op] std::make_tuple(1, 1, TASK_FIND_NGRAM_PARTIAL, variant_id);} else if (name find_ngram_global) {int variant_id task_register-register_find_ngram_global_task(customized-bgraph, params);task_config[op] std::make_tuple(2, 1, TASK_FIND_NGRAM_GLOBAL, variant_id);} else if (name target_verify_greedy) {int variant_id task_register-register_target_verify_greedy_task(customized-bgraph, params);task_config[op] std::make_tuple(2, 1, TASK_TARGET_VERIFY_GREEDY, variant_id);}}以register_embedding_task为例其代码如下int TaskRegister::register_embedding_task(threadblock::Graph const bgraph,std::vectorint const params) {assert(params.size() 1);// params[0]: input source (0: tokens, 1: input_token)int batch_size 0, output_size 0, output_stride 0;std::vectortb::TBInputOp * input_ops;std::vectortb::TBInputOp * output_ops;int num_inputs 2;int num_outputs 1;assert(bgraph.operators.size() (size_t)num_inputs num_outputs);for (auto const op : bgraph.operators) {assert(op-op_type mirage::type::TB_INPUT_OP);if (input_ops.size() (size_t)num_inputs) {input_ops.push_back(static_casttb::TBInputOp *(op));} else {output_ops.push_back(static_casttb::TBInputOp *(op));}}assert(output_ops[0]-output_tensors[0].num_dims 2);batch_size output_ops[0]-output_tensors[0].dim[0];output_size output_ops[0]-output_tensors[0].dim[1];kn::KNInputOp *kn_input_op static_castkn::KNInputOp *(output_ops[0]-dtensor.owner_op);output_stride static_castint(kn_input_op-input_strides[0]);mirage::transpiler::CodeKeeper code;code.inc_indent();code.e(kernel::embedding_kernelbfloat16, $, $, $(,batch_size,output_size,output_stride);if (params[0] 0) {code.e( runtime_config.tokens runtime_config.step[0], );} else if (params[0] 1) {code.e( task_desc.inputs[0].base_ptr,);}code.e( task_desc.inputs[1].base_ptr,);code.e( task_desc.outputs[0].base_ptr););return register_task_variant(TASK_EMBEDDING, code.to_string());}最终算子embedding_kernel定义如下namespace kernel {template typename T, int BATCH_SIZE, int CHUNK_SIZE, int OUTPUT_DIM_SIZE__device__ __forceinline__ voidembedding_kernel(void const *__restrict__ input_ptr,void const *__restrict__ embedding_ptr,void *__restrict__ output_ptr) {int64_t const *__restrict__ input_ids static_castint64_t const *(input_ptr);T const *__restrict__ embedding static_castT const *(embedding_ptr);T *__restrict__ output static_castT *(output_ptr);#pragma unrollfor (int batch_idx 0; batch_idx BATCH_SIZE; batch_idx) {int64_t wordIdx input_ids[batch_idx];if (wordIdx 0) {#pragma unrollfor (int i threadIdx.x; i CHUNK_SIZE; i NUM_THREADS) {output[batch_idx * OUTPUT_DIM_SIZE i] embedding[wordIdx * OUTPUT_DIM_SIZE i];}} else {// TODO: This might not be necessaryfor (int i threadIdx.x; i CHUNK_SIZE;i NUM_THREADS) { // writing 0 to outputoutput[batch_idx * OUTPUT_DIM_SIZE i] T(0.0f);}}}}} // namespace kernel2.2 注册代码上述代码TaskRegister::register_embedding_task 调用了 register_task_variant 函数来对all_task_variants 进行设置。TaskRegister:register_*_task 函数如 register_embedding_task, register_custom_task 等会根据 TaskBlock::Graph 和参数生成特定的 CUDA 调用代码字符串并将其注册到 all_task_variants 中返回该变体在向量中的索引即 variant_id。TaskRegister 单例mirage::runtime::TaskRegister 是一个单例类负责管理和注册所有可能的任务变体代码。它内部维护一个映射std::mapruntime::TaskType, std::vectorstd::string all_task_variants。all_task_variants 的作用是存储和管理不同类型任务的代码变体。键是任务类型TaskTypetask_type 指定了任务的大类例如 TASK_EMBEDDING, TASK_ATTENTION_1, TASK_LINEAR_WITH_RESIDUAL 等。值是该类型任务的代表变体列表。all_task_variants为每种任务类型维护一个代码变体集合。在register_task_variant中会检查是否存在相同的代码变体避免重复存储。这样可以允许同一种任务类型有不同的实现方式。variant_id 指定了同一任务类型下的具体变体因为同一逻辑任务可能有多种不同的实现或参数配置。即all_task_variants这个映射将每个 TaskType 关联到一个字符串向量向量中的每个字符串代表该任务类型的一个具体实现代码通常是以字符串形式生成的 CUDA kernel 调用代码。register_task_variant函数register_task_variant函数代码如下int TaskRegister::register_task_variant(runtime::TaskType type,std::string const code) {std::vectorstd::string variants all_task_variants[type];for (size_t i 0; i variants.size(); i) {if (variants[i] code) {return (int)(i);}}// Add a new variantvariants.push_back(code);return (int)(variants.size() - 1);}2.3 获取代码回忆下在生成任务图时会做如下操作。在 runtime.cc 的 register_mugraph 函数中会遍历 Graph 中的 KN_CUSTOMIZED_OP 操作符。对于每个操作符它会从 task_configs即 Graph::task_config中查找对应的配置输入数、输出数、TaskType, variant_id。创建 TaskDesc 结构体会将获取到的 TaskType 和 variant_id 填入 TaskDesc。运行时获取代码的过程如下当持久化内核persistent kernel运行时执行到某个 TaskDesc它会根据其 task_type 和 variant_id进行操作。task_type 指定了任务的大类例如 TASK_EMBEDDING, TASK_ATTENTION_1, TASK_LINEAR_WITH_RESIDUAL 等。variant_id 指定了同一任务类型下的具体变体因为同一逻辑任务可能有多种不同的实现或参数配置。在 TaskRegister::all_task_variants 中找到对应的任务类型向量。使用 variant_id 作为索引从该向量中取出预先生成好的 CUDA kernel 调用代码字符串。这个字符串通常会被编译成实际的 kernel 函数可能通过 JIT 编译或预先编译的库然后通过 CUDA API如 cudaLaunchKernel 或类似的封装来执行。0x03 生成任务图3.1 入口persistent_kernel.py 的 compile 函数会调用kn_graph.generate_task_graph来生成任务图即从计算图生成cu文件。def compile(self,**kwargs,):output_dir kwargs.get(output_dir, None)MIRAGE_ROOT, INCLUDE_PATH, DEPS_PATH get_key_paths()tempdir_obj tempfile.TemporaryDirectory()tempdir tempdir_obj.nameresults self.kn_graph.generate_task_graph(num_gpusself.world_size, my_gpu_idself.mpi_rank)generate_task_graph的代码如下def generate_task_graph(self, num_gpus: int, my_gpu_id: int):return self.cygraph.generate_task_graph(num_gpus, my_gpu_id)3.2 runtime.cc主体generate_task_graph 调用register_mugraph来进行转换建立event和task调用print_task_graph把代码转换出来。TaskGraphResult Graph::generate_task_graph(int _num_gpus, int _my_gpu_id) {std::vectorTaskDesc all_tasks;std::vectorEventDesc all_events;std::vectorTaskId first_tasks;int num_gpus, my_gpu_id;std::mapkernel::KNOperator *, std::mapdim3, TaskId, Dim3Comparatorall_task_maps;num_gpus _num_gpus;my_gpu_id _my_gpu_id;// add the termination event to the event listsEventDesc e(EVENT_TERMINATION, 1, 0, 0);all_events.push_back(e);TaskDesc t(TASK_TERMINATE, 0 /*variant_id*/);all_tasks.push_back(t);register_mugraph(*this,num_gpus,my_gpu_id,all_tasks,all_events,first_tasks,all_task_maps,task_config);assert(sanity_check(*this, all_tasks, all_events, first_tasks));return print_task_graph(*this,num_gpus,my_gpu_id,all_tasks,all_events,first_tasks,all_task_maps,task_config,io_config,true /*use_json_format*/);}这些代码都位于runtime.cc。3.2.1 runtime.cc的功能runtime.cc本质是转译器将高级内核图转换为可以在持久化内核运行时系统中执行的低级任务图表示。runtime.cc和persistent_kernel.py共同构成了Mirage系统中持久化内核执行系统的核心部分。runtime.ccC实现负责底层的任务图生成、事件管理和代码生成。persistent_kernel.pyPython实现提供高层接口和抽象用于定义和配置持久化内核的数据流关系。persistent_kernel.py中定义的内核配置和图结构会被传递给runtime.ccruntime.cc会使用这些信息生成实际的CUDA代码和任务图。两者的协同工作流程如下mirage-4-2.5具体交互点如下任务配置传递。persistent_kernel.py的配置通过task_config传递给runtime.ccruntime.cc的register_mugraph函数使用这些配置来创建任务I/O配置传递persistent_kernel.py定义的I/O配置通过io_config传递给runtime.ccruntime.cc的print_task_graph函数使用这些配置来生成正确的内存分配代码。代码生成runtime.cc的print_task_graph函数生成实际的CUDA代码生成的代码例如_init_persistent_kernel和_execute_task 函数这些生成的函数会被persistent_kernel.py使用来执行实际的内核事件和任务管理runtime.cc负责创建和管理事件及任务之间的依赖关系这些事件如EVENT_LAUNCH_TASKS在两个文件中都 被使用。3.2.2 runtime.cc总体流程runtime.cc总体流程如下mirage-4-23.2.3 runtime.cc的具体函数具体函数如下generate_task_graph主入口点协调整个任务图的生成过程。register_mugraph核心函数负责1 将内核图转换为任务和事件即TaskDesc和EventDesc序列2 处理特殊操作如ALLREDUCE。3 使用事件设置任务间的正确依赖关系。4 根据任务数量确定适当的事件类型。5 建立操作符到任务ID的映射关系dfs_create_events_add_tasks 递归函数负责1 使用深度优先搜索方法创建事件和任务。2 处理多维任务分区。3 在生成者和消费者任务之间分配正确的依赖关系。sanity_check()验证函数负责1 确保所有任务都能被执行。2 验证所有事件都能被触发。print_task_graph输出生成函数负责1 创建用于初始化持久化内核的CUDA代码2 生成任务图的JSON表示3 生成执行任务的设备函数3.3 建立依赖关系register_mugraph函数完成了从内核图由KNOperator组成到可执行的任务图的关键转换过程图结构转换将 KNOperator 图转换为 TaskDesc 和 EventDesc 序列依赖关系建立通过事件机制建立任务间的依赖关系分布式支持特殊处理 ALLREDUCE 等分布式操作任务映射建立操作符到任务ID的映射关系资源配置为运行时执行准备必要的任务和事件描述register_mugraph函数是连接计算图定义和实际 GPU 执行的重要桥梁。3.3.1 流程具体流程如下初始化任务图结构添加开始任务和事件来启动依赖任务。遍历图中所有操作符。特殊处理ALLREDUCE操作等分布式操作。创建NVSHMEM复制任务用于跨GPU数据传输创建REDUCE任务用于规约操作。为每个操作创建任务描述创建操作间依赖事件。更新触发事件。其中 num_shared_tensors 变量的作用时统计当前操作符与前一个操作符之间共享的张量数量。当找到共享变量时会记录下相关的映射信息这些信息会在后续创建事件和任务时会使用。mirage-4-33.3.2 结果register_mugraph生成的主要结果为任务描述列表all_tasks包含所有需要执行的任务描述TaskDesc每个任务包含任务类型、变体ID、输入输出张量等描述信息。任务按照执行顺序排列。事件描述列表all_events包含所有事件的描述EventDesc。每个事件描述包含事件类型、触发任务数量、任务ID范围等。控制任务间的依赖关系和执行顺序。首任务列表 first_tasks包含任务图中第一批可以执行的任务ID任务映射表 all_tasks_maps映射每个操作符到其对应的任务ID映射表用于定位特定操作符生成的任务。后续print_task_graph会利用这些生成结果。3.3.3 代码register_mugraph具体代码如下void register_mugraph( // 接受一个kernel图GPU数量当前GPU ID以及任务和事件相关容器mirage::kernel::Graph const graph,int num_gpus,int my_gpu_id,std::vectorTaskDesc all_tasks,std::vectorEventDesc all_events,std::vectorTaskId first_tasks,std::mapkernel::KNOperator *, std::mapdim3, TaskId, Dim3Comparatorall_task_maps,std::unordered_mapkn::KNOperator const *,std::tupleint, int, TaskType, int consttask_configs) {// push a begin-graph task and a event to launch dependent asks// 添加一个开始任务图的事件和任务即初始化任务图结构{EventDesc e(EVENT_LAUNCH_DEPENDENT_TASKS, 1, 0, 0);TaskDesc t(TASK_BEGIN_TASK_GRAPH, 0 /*variant_id*/);// 设置任务触发事件IDt.trigger_event get_event_id(my_gpu_id, all_events.size(), false);all_tasks.push_back(t);all_events.push_back(e);}// 保存前一个操作的输出操作符和映射关系std::vectortb::TBInputOp * pre_output_ops;kn::KNCustomizedOp const *pre_op nullptr;std::mapdim3, TaskId, Dim3Comparator pre_task_map;// 遍历图中所有的操作符for (auto const op : graph.operators) {// 跳过输入操作符if (op-op_type type::KNOperatorType::KN_INPUT_OP) {continue;}// 获取当前操作的任务配置std::tupleint, int, TaskType, int task_config task_configs.find(op)-second;// 获取当前操作的任务映射std::mapdim3, TaskId, Dim3Comparator cur_task_map;assert(op-op_type type::KNOperatorType::KN_CUSTOMIZED_OP);// Customized op// 将操作转换为自定义操作类型kn::KNCustomizedOp const *cur_op dynamic_castkn::KNCustomizedOp const *(op);// 获取线程块图tb::Graph const bgraph cur_op-bgraph;dim3 bid;// 存储任务描述的向量std::vectorTaskDesc tasks;// 存储输入输出操作符std::vectortb::TBInputOp * input_ops;std::vectortb::TBInputOp * output_ops;// 从配置中获取输入输出数量和任务类型int num_inputs std::get0(task_config);int num_outputs std::get1(task_config);TaskType task_type std::get2(task_config);int variant_id std::get3(task_config);// 确保操作符数量为输出输出之和assert(bgraph.operators.size() (size_t)num_inputs num_outputs);// 分离输入输出操作符for (auto const op : bgraph.operators) {assert(op-op_type mirage::type::TB_INPUT_OP);if (input_ops.size() (size_t)num_inputs) {input_ops.push_back(static_casttb::TBInputOp *(op));} else {output_ops.push_back(static_casttb::TBInputOp *(op));}}// Specical handling for ALLREDUCEif (task_type TASK_ALLREDUCE) {// Shouldnt have AllReduce when num_gpus 1assert(num_gpus 1); // 需要多个GPUassert(input_ops.size() 2); // 确保输入输出数量正确assert(output_ops.size() 1);// To simplify the implementation, asserting that// produce/consumer must have the same partitionint num_shared_tensors 0;int3 input_map, output_map;// 查找共享张量并获取映射关系for (auto const input : input_ops) {for (auto const output : pre_output_ops) {if (input-dtensor.guid output-dtensor.guid) {input_map input-input_map;output_map output-input_map;num_shared_tensors;}}}assert(num_shared_tensors 1); // 确保有一个共享张量assert(input_map output_map); // 确保映射关系相同且网格维度一致assert(bgraph.grid_dim pre_op-bgraph.grid_dim);dim3 bid;// 存储ALLGather前任务映射std::mapdim3, std::mapint, TaskId, Dim3Comparator ag_pre_task_map;// 遍历所有线程块维度for (bid.x 0; bid.x bgraph.grid_dim.x; bid.x) {for (bid.y 0; bid.y bgraph.grid_dim.y; bid.y) {for (bid.z 0; bid.z bgraph.grid_dim.z; bid.z) {// event_desc_0 is the trigger_event of previous_task// event_desc_1 is the trigger_event of allgather// 创建事件描述用于触发前一个任务EventDesc event_desc_0;event_desc_0.event_type EVENT_LAUNCH_TASKS;event_desc_0.num_triggers 1;event_desc_0.first_task_id all_tasks.size();event_desc_0.last_task_id all_tasks.size() num_gpus - 1;// 确保前一个任务映射中存在当前块assert(pre_task_map.find(bid) ! pre_task_map.end());int task_id pre_task_map.find(bid)-second;// 设置前一个任务的触发事件all_tasks[task_id].trigger_event get_event_id(my_gpu_id, all_events.size(), false);all_events.push_back(event_desc_0);// Step 1: create (num_gpus - 1) tasks for allgatherstd::mapint, TaskId pre_tasks;for (int tgt_gpu_id 0; tgt_gpu_id num_gpus; tgt_gpu_id) {if (tgt_gpu_id my_gpu_id) {continue; // 跳过当前GPU}// 创建 TASK_NVSHMEM_COPY 复制任务TaskDesc task(TASK_NVSHMEM_COPY, 0 /*variant_id*/);// task.trigger_event get_event_id(// tgt_gpu_id, all_events.size(), true /*nvshmem_event*/);// Initialize input tensors to the task{TensorDesc desc;assert(input_ops[0]-output_tensors.size() 1);tb::STensor stensor input_ops[0]-output_tensors[0];desc.num_dims stensor.num_dims;desc.data_type stensor.data_type;for (int d stensor.num_dims - 1; d 0; d--) {desc.dim[d] stensor.dim[d];desc.stride[d] (d stensor.num_dims - 1)? 1: desc.stride[d 1] *input_ops[0]-dtensor.dim[d 1];}task.inputs[task.num_inputs] desc;}// Initialize output tensors to the task{TensorDesc desc;assert(input_ops[1]-output_tensors.size() 1);tb::STensor stensor input_ops[1]-output_tensors[0];desc.num_dims stensor.num_dims;desc.data_type stensor.data_type;for (int d stensor.num_dims - 1; d 0; d--) {desc.dim[d] stensor.dim[d];desc.stride[d] (d stensor.num_dims - 1)? 1: desc.stride[d 1] *input_ops[1]-dtensor.dim[d 1];}task.outputs[task.num_outputs] desc;}all_tasks.push_back(task);pre_tasks[tgt_gpu_id] all_tasks.size() - 1;} // for tgt_gpu_idag_pre_task_map[bid] pre_tasks;} // for bid.z} // for bid.y} // for bid.x// 遍历所有线程块维度处理reduce 任务for (bid.x 0; bid.x bgraph.grid_dim.x; bid.x) {for (bid.y 0; bid.y bgraph.grid_dim.y; bid.y) {for (bid.z 0; bid.z bgraph.grid_dim.z; bid.z) {// event_desc_1 is the trigger_event of allgather// 创建allgather 的触发事件EventDesc event_desc_1;event_desc_1.event_type EVENT_LAUNCH_TASKS;event_desc_1.first_task_id all_tasks.size();event_desc_1.last_task_id all_tasks.size() 1;event_desc_1.num_triggers num_gpus - 1;// 确保存在当前任务映射assert(ag_pre_task_map.find(bid) ! ag_pre_task_map.end());std::mapint, TaskId pre_tasks ag_pre_task_map.find(bid)-second;// 设置所有前任务的触发事件for (auto const t : pre_tasks) {all_tasks[t.second].trigger_event get_event_id(t.first, all_events.size(), true);}all_events.push_back(event_desc_1);// Step 2: create a task for reduceTaskDesc task(TASK_REDUCE, 0 /*variant_id*/);// 初始化输入张量for (int i 0; i 2; i) {TensorDesc desc;tb::STensor stensor input_ops[i]-output_tensors[0];desc.num_dims stensor.num_dims;desc.data_type stensor.data_type;for (int d stensor.num_dims - 1; d 0; d--) {desc.dim[d] stensor.dim[d];desc.stride[d] (d stensor.num_dims - 1)? 1: desc.stride[d 1] * input_ops[1]-dtensor.dim[d 1];}task.inputs[task.num_inputs] desc;}// Create output tensor{TensorDesc desc;tb::STensor stensor output_ops[0]-output_tensors[0];desc.num_dims stensor.num_dims;desc.data_type stensor.data_type;for (int d stensor.num_dims - 1; d 0; d--) {desc.dim[d] stensor.dim[d];desc.stride[d] (d stensor.num_dims - 1)? 1: desc.stride[d 1] *output_ops[0]-dtensor.dim[d 1];}task.inputs[task.num_outputs] desc;all_tasks.push_back(task);// Update current task map// 当前任务映射cur_task_map[bid] all_tasks.size() - 1;}}}}// 更新前操作相关变量pre_output_ops output_ops;pre_op cur_op;pre_task_map cur_task_map;all_task_maps.emplace(op, cur_task_map);continue;}// Step 1: add all tasks based on their blockIdx// (bid.x, bid.y, bid.z) ordering// 根据 blockIdx 添加所有任务 (bid.x, bid.y, bid.z)的顺序for (bid.x 0; bid.x bgraph.grid_dim.x; bid.x) {for (bid.y 0; bid.y bgraph.grid_dim.y; bid.y) {for (bid.z 0; bid.z bgraph.grid_dim.z; bid.z) {TaskDesc task(task_type, variant_id); // 创建任务描述// Initialize input tensors to the taskfor (auto const input : input_ops) { // 初始化任务的输入张量TensorDesc desc;assert(input-output_tensors.size() 1);tb::STensor stensor input-output_tensors[0];desc.num_dims stensor.num_dims;desc.data_type stensor.data_type;for (int d stensor.num_dims - 1; d 0; d--) {desc.dim[d] stensor.dim[d];desc.stride[d] (d stensor.num_dims - 1)? 1: desc.stride[d 1] * input-dtensor.dim[d 1];}task.inputs[task.num_inputs] desc;}// Initialize output tensors to the taskfor (auto const output : output_ops) { // 初始化任务的输出张量TensorDesc desc;assert(output-output_tensors.size() 1);tb::STensor stensor output-output_tensors[0];desc.num_dims stensor.num_dims;desc.data_type stensor.data_type;for (int d stensor.num_dims - 1; d 0; d--) {desc.dim[d] stensor.dim[d];desc.stride[d] (d stensor.num_dims - 1)? 1: desc.stride[d 1] * output-dtensor.dim[d 1];}task.outputs[task.num_outputs] desc;}tasks.push_back(task);}}}// Step 2: create events between operators// 在操作符之间创建事件if (pre_op nullptr) {// 如果是第一个操作符添加到first_tasksdim3 bid;for (bid.x 0; bid.x bgraph.grid_dim.x; bid.x) {for (bid.y 0; bid.y bgraph.grid_dim.y; bid.y) {for (bid.z 0; bid.z bgraph.grid_dim.z; bid.z) {cur_task_map[bid] all_tasks.size();int offset bid.x * bgraph.grid_dim.y * bgraph.grid_dim.z bid.y * bgraph.grid_dim.z bid.z;first_tasks.push_back(all_tasks.size());all_tasks.push_back(tasks[offset]);}}}} else {// Step 2.1: analyze dependencies between thread blocks of the two ops// 分析两个操作之间线程块的依赖关系std::vectorint producer_partition(mirage::config::MAX_TENSOR_DIMS, 1);std::vectorint consumer_partition(mirage::config::MAX_TENSOR_DIMS, 1);int num_shared_tensors 0;int3 input_map, output_map;// 查找共享张量并获取映射关系for (auto const input : input_ops) {for (auto const output : pre_output_ops) {if (input-dtensor.guid output-dtensor.guid) {input_map input-input_map;output_map output-input_map;num_shared_tensors;}}}// assert that their is at least a single tensor shared between opsassert(num_shared_tensors 1); // 确保至少有一个共享张量// 设置生产者和消费者的分区for (int d 0; d mirage::config::MAX_TENSOR_DIMS; d) {if (d input_map.x) {consumer_partition[d] bgraph.grid_dim.x;}if (d input_map.y) {consumer_partition[d] bgraph.grid_dim.y;}if (d input_map.z) {consumer_partition[d] bgraph.grid_dim.z;}if (d output_map.x) {producer_partition[d] pre_op-bgraph.grid_dim.x;}if (d output_map.y) {producer_partition[d] pre_op-bgraph.grid_dim.y;}if (d output_map.z) {producer_partition[d] pre_op-bgraph.grid_dim.z;}}// Step 2.2: create events and add tasks 创建事件并添加任务// number of events is the product of gcd of producer/consumerstd::vectorint event_dims(mirage::config::MAX_TENSOR_DIMS, 1);for (int d 0; d mirage::config::MAX_TENSOR_DIMS; d) {event_dims[d] std::gcd(producer_partition[d], consumer_partition[d]);}// 利用深度优先搜索创建事件和添加任务dfs_create_events_add_tasks(0, /*depth*/my_gpu_id, /*my_gpu_id*/event_dims, /*event_dims*/input_map, /*input_map*/output_map, /*output_map*/bgraph.grid_dim, /*consumer_grid_dim*/pre_op-bgraph.grid_dim, /*producer_grid_dim*/dim3(0, 0, 0), /*consumer_lo_bid*/bgraph.grid_dim, /*consumer_hi_bid*/dim3(0, 0, 0), /*producer_lo_bid*/pre_op-bgraph.grid_dim, /*producer_hi_bid*/all_events,all_tasks,tasks, /*cur_op_tasks*/pre_task_map, /*pre_task_map*/cur_task_map /*cur_task_map)*/);}pre_output_ops output_ops;pre_op cur_op;pre_task_map cur_task_map;all_task_maps.emplace(op, cur_task_map);}// Update the trigger event for all tasks in pre_task_mapfor (auto const it : pre_task_map) {all_tasks[it.second].trigger_event get_event_id(my_gpu_id, all_events.size(), false /*nvshmem_event*/);}// 添加任务图结束事件all_events.push_back(EventDesc(EVENT_END_OF_TASK_GRAPH, pre_task_map.size(), 0, 0));// Prelaunch all tasks at the begining of an iteration// 迭代开始时预启动所有任务all_events[1].first_task_id 2;all_events[1].last_task_id all_tasks.size();for (size_t e 2; e all_events.size(); e) {// 对于任务启动事件将其转换为空事件if (all_events[e].event_type EVENT_LAUNCH_TASKS ||all_events[e].event_type EVENT_LAUNCH_MASSIVE_TASKS) {all_events[e].event_type EVENT_EMPTY;// 为相关任务设置依赖事件for (size_t t all_events[e].first_task_id;t all_events[e].last_task_id;t) {all_tasks[t].dependent_event get_event_id(my_gpu_id, e, false /*nvshmem_event*/);}}}}3.4 输出代码print_task_graph包括两部分。代码生成在print_task_graph中生成完整的CUDA源文件。文件输出将生成的CUDA代码写入.cu文件供后续编译使用。上述方式允许系统根据计算图结构动态生成优化的CUDA kernel代码。mirage-4-43.4.1 逻辑print_task_graph接受register_mugraph生成的所有关键数据结构all_tasks包含所有任务描述的向量。all_events包含所有事件描述的向量。first_tasks包含第一批任务ID的向量。all_task_maps操作符到任务的映射表。print_task_graph生成的CUDA代码包括任务图构造函数 construct_task_graph任务和事件的初始化代码 _init_persistent_kernel。内存分配代码CUDANVSHMEM张量_execute_taskprint_task_graph生成的JSON包括从task_graph.json文件读取任务信息解析任务输入输出张量描述重建完整的任务结构。print_task_graph 利用如下信息生成任务依赖关系。all_tasks中的trigger_event和dependent_event字段all_events中的事件触发关系first_tasks确定任务图的入口点。3.4.2 代码print_task_graph具体代码如下TaskGraphResult print_task_graph(// 函数参数内核图、GPU数量、当前GPU ID、所有任务描述、所有事件描述、首任务列表mirage::kernel::Graph const graph,int num_gpus,int my_gpu_id,std::vectorTaskDesc const all_tasks,std::vectorEventDesc const all_events,std::vectorTaskId const first_tasks,// 所有操作符到任务映射的映射std::mapkernel::KNOperator *, std::mapdim3, TaskId, Dim3Comparator constall_task_maps,// 操作符到任务设置的映射std::unordered_mapkn::KNOperator const *,std::tupleint, int, TaskType, int const task_configs,// 输入输出配置映射std::mapmirage::type::GuidType, IODesc const io_configs,bool use_json_format) {using mirage::runtime::IODesc;// 创建代码生成器实例mirage::transpiler::CodeKeeper code;mirage::transpiler::CodeKeeper tgbody;tgbody.inc_indent();// 添加必要的头文件包含code.e(#include \persistent_kernel.cuh\);if (use_json_format) {code.e(#include nlohmann/json.hpp);code.e(#include fstream);code.e(#include filesystem);code.e(using json nlohmann::json;);}// 添加运行时命名空间声明code.e(using namespace mirage::runtime;);// 生成获取事件ID的函数code.e(size_t get_event_id(int my_gpu_id, size_t event_pos, bool nvshmem_event) {);code.e(size_t event_id ((static_castsize_t(my_gpu_id) 32) | event_pos););code.e(if (nvshmem_event) {);code.e(event_id event_id | EVENT_NVSHMEM_TAG;);code.e(});code.e(return event_id;);code.e(});code.e();// function that loads json file and generates task graph// 如果使用JSON格式生成从JSON文件构造人物图的函数if (use_json_format) {code.e(void construct_task_graph(int num_gpus,);code.e( int my_gpu_id,);code.e( std::vectorTaskDesc all_tasks,);code.e( std::vectorEventDesc all_events,);code.e( std::vectorTaskId first_tasks,);code.e( std::mapstd::string, void* const all_tensors) {);code.e(std::filesystem::path file_path(__FILE__););code.e(std::ifstream json_file(file_path.parent_path().string()\/task_graph.json\););code.e(nlohmann::json json_task_graph;);code.e(json_file json_task_graph;);// load tasks// 加载任务code.e(for (json const task : json_task_graph[\all_tasks\]) {);code.e(TaskDesc task_desc(static_castTaskType(task.at(\task_type\)),);code.e( task.at(\variant_id\)););code.e(if (task.at(\trigger_event\).is_number_integer()) {);code.e(task_desc.trigger_event task.at(\trigger_event\).getunsigned long long int(););code.e(});code.e(else {);code.e(assert(false););code.e(});code.e(if (task.at(\dependent_event\).is_number_integer()) {);code.e(task_desc.dependent_event task.at(\dependent_event\).getunsigned long long int(););code.e(});code.e(else {);code.e(assert(false););code.e(});// load inputs 加载输入张量code.e(task_desc.num_inputs 0;);code.e(for (json const tensor : task[\inputs\]) {);code.e(TensorDesc input;);code.e(std::string name tensor.at(\base_ptr\).getstd::string(););code.e(assert(all_tensors.find(name) ! all_tensors.end()););code.e(off_t offset tensor.at(\offset\).getoff_t(););code.e(input.base_ptr static_castchar*(all_tensors.at(name))offset;);code.e(assert(tensor.at(\dims\).size() tensor.at(\strides\).size()););code.e(input.num_dims tensor.at(\dims\).size(););code.e(input.data_type tensor.at(\data_type\).getint(););code.e(for (int i 0; i input.num_dims; i) {);code.e(input.dim[i] tensor[\dims\][i].getint(););code.e(input.stride[i] tensor[\strides\][i].getint(););code.e(});code.e(task_desc.inputs[task_desc.num_inputs] input;);code.e(});// load outputs 加载输出张量code.e(task_desc.num_outputs 0;);code.e(for (json const tensor : task[\outputs\]) {);code.e(TensorDesc output;);code.e(std::string name tensor.at(\base_ptr\).getstd::string(););code.e(assert(all_tensors.find(name) ! all_tensors.end()););code.e(off_t offset tensor.at(\offset\).getoff_t(););code.e(output.base_ptr static_castchar*(all_tensors.at(name))offset;);code.e(assert(tensor.at(\dims\).size() tensor.at(\strides\).size()););code.e(output.num_dims tensor.at(\dims\).size(););code.e(output.data_type tensor.at(\data_type\).getint(););code.e(for (int i 0; i output.num_dims; i) {);code.e(output.dim[i] tensor[\dims\][i];);code.e(output.stride[i] tensor[\strides\][i];);code.e(});code.e(task_desc.outputs[task_desc.num_outputs] output;);code.e(});code.e(all_tasks.push_back(task_desc););code.e(});// load events 加载事件code.e(for (json const e : json_task_graph[\all_events\]) {);code.e(EventType event_type static_castEventType(e.at(\event_type\).getint()););code.e(int num_triggers e.at(\num_triggers\).getint(););code.e(int first_task_id e.at(\first_task_id\).getint(););code.e(int last_task_id e.at(\last_task_id\).getint(););code.e(all_events.push_back(EventDesc(event_type, num_triggers, first_task_id, last_task_id)););code.e(});// load first tasks 加载首任务code.e(for (json const t : json_task_graph[\first_tasks\]) {);code.e(first_tasks.push_back(t.getint()););code.e(});code.e(});code.e();}// 生成初始化持久内核的函数code.e(static void _init_persistent_kernel(std::vectorTaskDesc all_tasks,);code.e( std::vectorEventDesc all_events,);code.e( std::vectorTaskId first_tasks,);code.e( int num_gpus,);code.e( int my_gpu_id) {);code.e(assert(num_gpus $);, num_gpus);if (use_json_format) {// 创建张量映射code.e(std::mapstd::string, void* all_tensors;);}for (auto const iter : io_configs) { // 输出输入输出配置IODesc desc iter.second;switch (desc.type) {case IODesc::TorchTensor: { // 处理Torch张量code.e(char *$ (char*)($);, desc.name, desc.torch_data_ptr);if (use_json_format) {code.e(all_tensors[\$\] $;, desc.name, desc.name);}break;}case IODesc::FusedTorchTensor: { // 处理融合张量for (auto const sdesc : desc.sub_descs) {code.e(char *$ (char*)($);, sdesc.name, sdesc.torch_data_ptr);if (use_json_format) {code.e(all_tensors[\$\] $;, sdesc.name, sdesc.name);}}break;}case IODesc::CUDAMallocTensor: { // 处理CUDA分配张量code.e(void *$;, desc.name);size_t size mirage::type::get_datatype_size(static_casttype::DataType(desc.tensor.data_type));for (int i 0; i desc.tensor.num_dims; i) {size * desc.tensor.dim[i];}code.e(cudaMalloc($, $);, desc.name, size);if (use_json_format) {code.e(all_tensors[\$\] $;, desc.name, desc.name);}break;}case IODesc::NVSHMEMMallocTensor: { // 处理NVSHMEM分配张量size_t size mirage::type::get_datatype_size(static_casttype::DataType(desc.tensor.data_type));for (int i 0; i desc.tensor.num_dims; i) {size * desc.tensor.dim[i];}code.e(void *$ nvshmem_malloc($);, desc.name, size);if (use_json_format) {code.e(all_tensors[\$\] $;, desc.name, desc.name);}break;}default:assert(false);}}json json_task_graph { // 创建jSON任务图对象{all_tasks, {}}, {all_events, {}}, {first_tasks, {}}};// generate task[0] 终止任务{tgbody.e(all_tasks.push_back(TaskDesc(TASK_TERMINATE)););json_task_graph[all_tasks].push_back(json{{task_type, TASK_TERMINATE},{variant_id, 0},{inputs, {}},{outputs, {}},{trigger_event, EVENT_INVALID_ID},{dependent_event, EVENT_INVALID_ID}});}// generate task[1] 任务图任务{tgbody.e(all_tasks.push_back(TaskDesc(TASK_BEGIN_TASK_GRAPH)););json_task_graph[all_tasks].push_back(json{{task_type, TASK_BEGIN_TASK_GRAPH},{variant_id, 0},{inputs, {}},{outputs, {}},{trigger_event,get_event_id(my_gpu_id, 1 /*event_pos*/, false /*is_nvshmem*/)},{dependent_event, EVENT_INVALID_ID}});}// generate all other tasks 生成所有其它任务size_t task_pos 2;for (auto const op : graph.operators) {if (op-op_type type::KNOperatorType::KN_INPUT_OP) {continue;}assert(op-op_type type::KNOperatorType::KN_CUSTOMIZED_OP);std::tupleint, int, TaskType, int task_config task_configs.find(op)-second;assert(all_task_maps.find(op) ! all_task_maps.end());std::mapdim3, TaskId, Dim3Comparator const task_map all_task_maps.find(op)-second;// Customized opkn::KNCustomizedOp const *cur_op dynamic_castkn::KNCustomizedOp const *(op);tb::Graph const bgraph cur_op-bgraph;dim3 bid;std::vectortb::TBInputOp * input_ops;std::vectortb::TBInputOp * output_ops;int num_inputs std::get0(task_config);// int num_outputs std::get1(task_config);TaskType task_type std::get2(task_config);// 收集输入和输出操作for (auto const op : bgraph.operators) {assert(op-op_type mirage::type::TB_INPUT_OP);if (input_ops.size() (size_t)num_inputs) {input_ops.push_back(static_casttb::TBInputOp *(op));} else {output_ops.push_back(static_casttb::TBInputOp *(op));}}if (task_type TASK_ALLREDUCE) { // 处理特殊任务for (bid.x 0; bid.x bgraph.grid_dim.x; bid.x) {for (bid.y 0; bid.y bgraph.grid_dim.y; bid.y) {for (bid.z 0; bid.z bgraph.grid_dim.z; bid.z) {// To perform allreduce, we first launch (num_gpus-1) tasks for// allgatherfor (int tgt_gpu_id 0; tgt_gpu_id num_gpus; tgt_gpu_id) {if (tgt_gpu_id my_gpu_id) {continue;}TaskDesc task_desc all_tasks[task_pos];assert(task_desc.task_type TASK_NVSHMEM_COPY);tgbody.e(// task[$], task_pos);tgbody.e({);tgbody.e(TaskDesc task_desc(static_castTaskType($));,task_desc.task_type);bool is_nvshmem_event ((task_desc.trigger_event EVENT_NVSHMEM_TAG) 0);assert(is_nvshmem_event);assert(task_desc.dependent_event ! EVENT_INVALID_ID);assert(task_desc.num_inputs 1);assert(task_desc.num_outputs 1);json json_task {{task_type, task_desc.task_type},{variant_id, task_desc.variant_id},{inputs, {}},{outputs, {}},{trigger_event, task_desc.trigger_event},{dependent_event, task_desc.dependent_event}};off_t offset 0;// Add inputint3 input_map input_ops[0]-input_map;IODesc io_desc io_configs.find(input_ops[0]-dtensor.guid)-second;if (input_map.x 0) {size_t block_size io_desc.tensor.dim[input_map.x] / bgraph.grid_dim.x;offset block_size * bid.x * io_desc.tensor.stride[input_map.x];}if (input_map.y 0) {size_t block_size io_desc.tensor.dim[input_map.y] / bgraph.grid_dim.y;offset block_size * bid.y * io_desc.tensor.stride[input_map.y];}if (input_map.z 0) {size_t block_size io_desc.tensor.dim[input_map.z] / bgraph.grid_dim.z;offset block_size * bid.z * io_desc.tensor.stride[input_map.z];}tgbody.e(TensorDesc input$;, 0);tgbody.e(input$.base_ptr static_castchar*($) $;,0,io_desc.name,offset *type::get_datatype_size(static_casttype::DataType(io_desc.tensor.data_type)));tgbody.e(input$.num_dims $;, 0, task_desc.inputs[0].num_dims);tgbody.e(input$.data_type $;, 0, task_desc.inputs[0].data_type);json json_dims json::array(), json_strides json::array();for (int d 0; d task_desc.inputs[0].num_dims; d) {tgbody.e(input$.dim[$] $;, 0, d, task_desc.inputs[0].dim[d]);tgbody.e(input$.stride[$] $;,0,d,task_desc.inputs[0].stride[d]);json_dims.push_back(task_desc.inputs[0].dim[d]);json_strides.push_back(task_desc.inputs[0].stride[d]);}tgbody.e(task_desc.inputs[$] input$;, 0, 0);json_task[inputs].push_back(json{{base_ptr, io_desc.name},{offset,offset * type::get_datatype_size(static_casttype::DataType(io_desc.tensor.data_type))},{data_type, task_desc.inputs[0].data_type},{dims, json_dims},{strides, json_strides}});// Add nvshmem_copy output// Note that nvshmem_copys output is stored in input_ops[1]offset my_gpu_id * input_ops[0]-dtensor.num_elements();int3 output_map input_ops[1]-input_map;io_desc io_configs.find(input_ops[1]-dtensor.guid)-second;if (output_map.x 0) {size_t block_size io_desc.tensor.dim[output_map.x] / bgraph.grid_dim.x;offset block_size * bid.x * io_desc.tensor.stride[output_map.x];}if (output_map.y 0) {size_t block_size io_desc.tensor.dim[output_map.y] / bgraph.grid_dim.y;offset block_size * bid.y * io_desc.tensor.stride[output_map.y];}if (output_map.z 0) {size_t block_size io_desc.tensor.dim[output_map.z] / bgraph.grid_dim.z;offset block_size * bid.z * io_desc.tensor.stride[output_map.z];}tgbody.e(TensorDesc output$;, 0);tgbody.e(output$.base_ptr static_castchar*($) $;,0,io_desc.name,offset *type::get_datatype_size(static_casttype::DataType(io_desc.tensor.data_type)));tgbody.e(output$.num_dims $;, 0, task_desc.outputs[0].num_dims);tgbody.e(output$.data_type $;, 0, task_desc.outputs[0].data_type);json_dims json::array();json_strides json::array();for (int d 0; d task_desc.outputs[0].num_dims; d) {tgbody.e(output$.dim[$] $;, 0, d, task_desc.outputs[0].dim[d]);tgbody.e(output$.stride[$] $;,0,d,task_desc.outputs[0].stride[d]);json_dims.push_back(task_desc.outputs[0].dim[d]);json_strides.push_back(task_desc.outputs[0].stride[d]);}tgbody.e(task_desc.outputs[$] output$;, 0, 0);json_task[outputs].push_back(json{{base_ptr, io_desc.name},{offset,offset * type::get_datatype_size(static_casttype::DataType(io_desc.tensor.data_type))},{data_type, task_desc.outputs[0].data_type},{dims, json_dims},{strides, json_strides}});tgbody.e(all_tasks.push_back(task_desc););json_task_graph[all_tasks].push_back(json_task);tgbody.e(});task_pos;} // for tgt_gpu_id} // for bid.z} // for bid.y} // for bid.x} // if task_type TASK_ALLREDUCE// 为每个线程块生成任务for (bid.x 0; bid.x bgraph.grid_dim.x; bid.x) {for (bid.y 0; bid.y bgraph.grid_dim.y; bid.y) {for (bid.z 0; bid.z bgraph.grid_dim.z; bid.z) {TaskId task_id task_map.at(bid);TaskDesc task_desc all_tasks[task_pos];assert(task_desc.task_type task_type ||task_type TASK_ALLREDUCE);assert(task_pos (task_id 0xffffffff));tgbody.e(// task[$], task_pos);