1
2
3
4
5
6

def main():
    parser = FlexibleArgumentParser(description="xFuser Arguments")
    args = xFuserArgs.add_cli_args(parser).parse_args()  # Add Command Line Interface (CLI) arguments
    engine_args = xFuserArgs.from_cli_args(args)  # Extract CLI args and pass them to xFuserArgs Constructor
    engine_config, input_config = engine_args.create_config()  # Init _WORLD. engine_config: model, run_time & parallel infos, input_config: input shape, prompt & sampler infos
    local_rank = get_world_group().local_rank

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18

Parallel Processing Options:
  --use_cfg_parallel    Use split batch in classifier_free_guidance. cfg_degree will be 2 if set
  --data_parallel_degree DATA_PARALLEL_DEGREE
                        Data parallel degree.
  --ulysses_degree ULYSSES_DEGREE
                        Ulysses sequence parallel degree. Used in attention layer.
  --ring_degree RING_DEGREE
                        Ring sequence parallel degree. Used in attention layer.
  --pipefusion_parallel_degree PIPEFUSION_PARALLEL_DEGREE
                        Pipefusion parallel degree. Indicates the number of pipeline stages.
  --num_pipeline_patch NUM_PIPELINE_PATCH
                        Number of patches the feature map should be segmented in pipefusion parallel.
  --attn_layer_num_for_pp [ATTN_LAYER_NUM_FOR_PP ...]
                        List representing the number of layers per stage of the pipeline in pipefusion parallel
  --tensor_parallel_degree TENSOR_PARALLEL_DEGREE
                        Tensor parallel degree.
  --split_scheme SPLIT_SCHEME
                        Split scheme for tensor parallel.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19

parallel_config = ParallelConfig(
    dp_config=DataParallelConfig(
        dp_degree=self.data_parallel_degree,
        use_cfg_parallel=self.use_cfg_parallel,
    ),
    sp_config=SequenceParallelConfig(
        ulysses_degree=self.ulysses_degree,
        ring_degree=self.ring_degree,
    ),
    tp_config=TensorParallelConfig(
        tp_degree=self.tensor_parallel_degree,
        split_scheme=self.split_scheme,
    ),
    pp_config=PipeFusionParallelConfig(
        pp_degree=self.pipefusion_parallel_degree,
        num_pipeline_patch=self.num_pipeline_patch,
        attn_layer_num_for_pp=self.attn_layer_num_for_pp,
    ),
)

1
2
3
4
5
6

    pipe = xFuserPixArtAlphaPipeline.from_pretrained(  # First construct a PixArtAlphaPipeline, then pass it and engine_config to xFuserPipelineBaseWrapper
        pretrained_model_name_or_path=engine_config.model_config.model,
        engine_config=engine_config,
        torch_dtype=torch.float16,
    ).to(f"cuda:{local_rank}")
    pipe.prepare_run(input_config)

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39

class xFuserPipelineBaseWrapper(xFuserBaseWrapper, metaclass=ABCMeta):

    def __init__(
        self,
        pipeline: DiffusionPipeline,
        engine_config: EngineConfig,
    ):
        self.module: DiffusionPipeline
        self._init_runtime_state(pipeline=pipeline, engine_config=engine_config)

        # backbone
        transformer = getattr(pipeline, "transformer", None)
        unet = getattr(pipeline, "unet", None)
        # vae
        vae = getattr(pipeline, "vae", None)
        # scheduler
        scheduler = getattr(pipeline, "scheduler", None)

        if transformer is not None:
            pipeline.transformer = self._convert_transformer_backbone(transformer)
        elif unet is not None:
            pipeline.unet = self._convert_unet_backbone(unet)

        if scheduler is not None:
            pipeline.scheduler = self._convert_scheduler(scheduler)

        super().__init__(module=pipeline)
   
   
    def _convert_transformer_backbone(
        self,
        transformer: nn.Module,
    ):
				#...

            logger.info("Transformer backbone found, paralleling transformer...")
            wrapper = **xFuserTransformerWrappersRegister.get_wrapper(transformer)**
            transformer = wrapper(transformer=transformer)
        return transformer

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56

def generate_masked_orthogonal_rank_groups(
    world_size: int, parallel_size: List[int], mask: List[bool]
) -> List[List[int]]:

    def prefix_product(a: List[int], init=1) -> List[int]:  # Exclusive
        r = [init]
        for v in a:
            init = init * v
            r.append(init)
        return r

    def inner_product(a: List[int], b: List[int]) -> int:
        return sum([x * y for x, y in zip(a, b)])

    def decompose(index, shape, stride=None):  # index: 第几个并行组  # shape: 并行组大小的 list
        """
        This function solve the math problem below:
            There is an equation: index = sum(idx[i] * stride[i])
            And given the value of index, stride.
            Return the idx.
        This function will used to get the pp/dp/pp_rank from group_index and rank_in_group.
        """
        if stride is None:
            stride = prefix_product(shape)
        idx = [(index // d) % s for s, d in zip(shape, stride)]  #  计算在每个并行维度上的索引
        # stride is a prefix_product result. And the value of stride[-1]
        # is not used.
        assert (
            sum([x * y for x, y in zip(idx, stride[:-1])]) == index
        ), "idx {} with shape {} mismatch the return idx {}".format(index, shape, idx)
        return idx

    masked_shape = [s for s, m in zip(parallel_size, mask) if m]  # 需要采取并行的维度
    unmasked_shape = [s for s, m in zip(parallel_size, mask) if not m]  # 不需要的

    global_stride = prefix_product(parallel_size)  # exclusive 前缀积 表示大的并行维度包括几个设备
    masked_stride = [d for d, m in zip(global_stride, mask) if m]
    unmasked_stride = [d for d, m in zip(global_stride, mask) if not m]

    group_size = prefix_product(masked_shape)[-1]  # 最大的一个并行维度包括几个设备
    num_of_group = world_size // group_size  # 分成几个大组

    ranks = []  
    for group_index in range(num_of_group):  # 遍历每个设备组
        # get indices from unmaksed for group_index.
        decomposed_group_idx = decompose(group_index, unmasked_shape)  # 得到在不需要采取并行的维度上的索引
        rank = []
        for rank_in_group in range(group_size):  # 遍历该组中的每个设备 local rank
            # get indices from masked for rank_in_group.
            decomposed_rank_idx = decompose(rank_in_group, masked_shape)  # 得到最大并行组的每个设备在采取并行的维度上的索引
            rank.append(  // 相加得到全局rank
                inner_product(decomposed_rank_idx, masked_stride)  
                + inner_product(decomposed_group_idx, unmasked_stride)
            )
        ranks.append(rank)
    return ranks

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27

def _split_transformer_blocks(self, transformer: nn.Module,):
    # omit

    # transformer layer split
    attn_layer_num_for_pp = (  # 获取每个 pipeFusion 的设备被分配的 transformer block 数量
        get_runtime_state().parallel_config.pp_config.attn_layer_num_for_pp
    )
    pp_rank = get_pipeline_parallel_rank()
    pp_world_size = get_pipeline_parallel_world_size()
    if attn_layer_num_for_pp is not None:
        if is_pipeline_first_stage():
            transformer.transformer_blocks = transformer.transformer_blocks[ : attn_layer_num_for_pp[0]]
        else:
            transformer.transformer_blocks = transformer.transformer_blocks[sum(attn_layer_num_for_pp[: pp_rank - 1]) : 
                                                                            sum(attn_layer_num_for_pp[:pp_rank])]
    else:  # 没有指定则平均分
        num_blocks_per_stage = (len(transformer.transformer_blocks) + pp_world_size - 1) // pp_world_size
        start_idx = pp_rank * num_blocks_per_stage
        end_idx = min((pp_rank + 1) * num_blocks_per_stage, len(transformer.transformer_blocks),)
        transformer.transformer_blocks = transformer.transformer_blocks[start_idx:end_idx]
    # position embedding
    if not is_pipeline_first_stage():
        transformer.pos_embed = None
    if not is_pipeline_last_stage():
        transformer.norm_out = None
        transformer.proj_out = None
    return transformer

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46

@xFuserTransformerWrappersRegister.register(PixArtTransformer2DModel)
class xFuserPixArtTransformer2DWrapper(xFuserTransformerBaseWrapper):
    def __init__(
        self,
        transformer: PixArtTransformer2DModel,
    ):
        super().__init__(
            transformer=transformer,
            submodule_classes_to_wrap=[nn.Conv2d, PatchEmbed],
            submodule_name_to_wrap=["attn1"],
        )

    @xFuserBaseWrapper.forward_check_condition
    def forward(
        self,
        hidden_states: torch.Tensor,
        encoder_hidden_states: Optional[torch.Tensor] = None,
        timestep: Optional[torch.LongTensor] = None,
        added_cond_kwargs: Dict[str, torch.Tensor] = None,
        cross_attention_kwargs: Dict[str, Any] = None,
        attention_mask: Optional[torch.Tensor] = None,
        encoder_attention_mask: Optional[torch.Tensor] = None,
        return_dict: bool = True,
    ):  
	'''
	......
	'''
	height, width = self._get_patch_height_width()
        # * only pp rank 0 needs pos_embed (patchify)
        if is_pipeline_first_stage():
            hidden_states = self.pos_embed(hidden_states)
	'''
	......
	'''
	if is_pipeline_last_stage():
	'''
	......
	'''
	else:
	    output = hidden_states

        if not return_dict:
            return (output,)

        return Transformer2DModelOutput(sample=output)

1
2
3
4
5
6

get_runtime_state().set_input_parameters(
    height=height,
    width=width,
    batch_size=batch_size,
    num_inference_steps=num_inference_steps,
)

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25

if do_classifier_free_guidance:
    (prompt_embeds, 
    prompt_attention_mask,) = self._process_cfg_split_batch(negative_prompt_embeds,
                                                                            prompt_embeds,
                                                                            negative_prompt_attention_mask,
                                                                            prompt_attention_mask,)

def _process_cfg_split_batch(self,
        concat_group_0_negative: torch.Tensor,
        concat_group_0: torch.Tensor,
        concat_group_1_negative: torch.Tensor,
        concat_group_1: torch.Tensor,):
  
  if get_classifier_free_guidance_world_size() == 1:
      concat_group_0 = torch.cat([concat_group_0_negative, concat_group_0], dim=0)
      concat_group_1 = torch.cat([concat_group_1_negative, concat_group_1], dim=0)
  elif get_classifier_free_guidance_rank() == 0:
      concat_group_0 = concat_group_0_negative
      concat_group_1 = concat_group_1_negative
  elif get_classifier_free_guidance_rank() == 1:
      concat_group_0 = concat_group_0
      concat_group_1 = concat_group_1
  else:
      raise ValueError("Invalid classifier free guidance rank")
  return concat_group_0, concat_group_1

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42

if len(timesteps) == 0:
    return latents
num_pipeline_patch = get_runtime_state().num_pipeline_patch
num_pipeline_warmup_steps = get_runtime_state().runtime_config.warmup_steps
patch_latents = self._init_async_pipeline(
    num_timesteps=len(timesteps),
    latents=latents,
    num_pipeline_warmup_steps=num_pipeline_warmup_steps,
)
last_patch_latents = (  # 每个 pipeline group 最后的设备接收所有的 patch
    [None for _ in range(num_pipeline_patch)]
    if (is_pipeline_last_stage())
    else None
)

def _init_async_pipeline(
    self,
    num_timesteps: int,
    latents: torch.Tensor,
    num_pipeline_warmup_steps: int,
):
    get_runtime_state().set_patched_mode(patch_mode=True)

    if is_pipeline_first_stage():
        # get latents computed in warmup stage
        # ignore latents after the last timestep
        latents = (get_pp_group().pipeline_recv()
                  if num_pipeline_warmup_steps > 0
                  else latents)
        patch_latents = list(latents.split(get_runtime_state().pp_patches_height, dim=2))
    elif is_pipeline_last_stage():
        patch_latents = list(latents.split(get_runtime_state().pp_patches_height, dim=2))
    else:
        patch_latents = [None for _ in range(get_runtime_state().num_pipeline_patch)]

    recv_timesteps = (num_timesteps - 1 if is_pipeline_first_stage() else num_timesteps)
    # construct receive tasks for each patch
    for _ in range(recv_timesteps):
        for patch_idx in range(get_runtime_state().num_pipeline_patch):
            get_pp_group().add_pipeline_recv_task(patch_idx)

    return patch_latents

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63

first_async_recv = True
for i, t in enumerate(timesteps):
    for patch_idx in range(num_pipeline_patch):
        if is_pipeline_last_stage():
            last_patch_latents[patch_idx] = patch_latents[patch_idx]

        if is_pipeline_first_stage() and i == 0:
            pass
        else:
            if first_async_recv:
                get_pp_group().recv_next()  
                first_async_recv = False
            patch_latents[patch_idx] = get_pp_group().get_pipeline_recv_data(
                idx=patch_idx
            )
        patch_latents[patch_idx] = self._backbone_forward(
            latents=patch_latents[patch_idx],
            prompt_embeds=prompt_embeds,
            prompt_attention_mask=prompt_attention_mask,
            added_cond_kwargs=added_cond_kwargs,
            t=t,
            guidance_scale=guidance_scale,
        )
        if is_pipeline_last_stage():
            patch_latents[patch_idx] = self._scheduler_step(
                patch_latents[patch_idx],  # pred noise
                last_patch_latents[patch_idx],  # last timestep noise
                t,
                extra_step_kwargs,
            )
            if i != len(timesteps) - 1:
                get_pp_group().pipeline_isend(
                    patch_latents[patch_idx], segment_idx=patch_idx
                )
        else:
            get_pp_group().pipeline_isend(
                patch_latents[patch_idx], segment_idx=patch_idx
            )

        if is_pipeline_first_stage() and i == 0:
            pass
        else:
            if i == len(timesteps) - 1 and patch_idx == num_pipeline_patch - 1:
                pass
            else:
                get_pp_group().recv_next()

        get_runtime_state().next_patch()  # switch to next: (self.pipeline_patch_idx + 1) % self.num_pipeline_patch

    if i == len(timesteps) - 1 or (
        (i + num_pipeline_warmup_steps + 1) > num_warmup_steps
        and (i + num_pipeline_warmup_steps + 1) % self.scheduler.order == 0
    ):
        progress_bar.update()
        assert callback is None, "callback not supported in async " "pipeline"
        if (
            callback is not None
            and i + num_pipeline_warmup_steps % callback_steps == 0
        ):
            step_idx = (i + num_pipeline_warmup_steps) // getattr(
                self.scheduler, "order", 1
            )
            callback(step_idx, t, patch_latents[patch_idx])

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20

latents = None
if is_pipeline_last_stage():
    latents = torch.cat(patch_latents, dim=2)
    if get_sequence_parallel_world_size() > 1:
        sp_degree = get_sequence_parallel_world_size()
        sp_latents_list = get_sp_group().all_gather(
            latents, separate_tensors=True
        )
        latents_list = []
        for pp_patch_idx in range(get_runtime_state().num_pipeline_patch):
            latents_list += [
                sp_latents_list[sp_patch_idx][
                    ...,
                    get_runtime_state().pp_patches_start_idx_local[pp_patch_idx] : get_runtime_state().pp_patches_start_idx_local[pp_patch_idx + 1],
                    :,
                ]
                for sp_patch_idx in range(sp_degree)
            ]
        latents = torch.cat(latents_list, dim=-2)
return latents

1
2
3
4
5
6
7

class PatchGroupNorm(nn.Module):
  
    ''' def __init__(self, ...)'''

    def forward(self, x: Tensor) -> Tensor:
        height = torch.tensor(x.shape[-2], dtype=torch.int64, device=x.device)
        dist.all_reduce(height)  # 收集所有进程的高度并汇总。最终每个进程的 height 都将表示全局的高度和。

1
2
3

channels_per_group = x.shape[1] // self.num_groups  # 每个组的通道数量
nelements_rank = channels_per_group * x.shape[-2] * x.shape[-1]  # 当前进程负责的每个组中的元素总
nelements = channels_per_group * height * x.shape[-1]  # 所有进程的每个组中的元素总数

1
2
3
4
5

x = x.view(x.shape[0], self.num_groups, -1, x.shape[-2], x.shape[-1])  #  [batch_size, num_groups, channels_per_group, height, width]
group_sum = x.mean(dim=(2,3,4), dtype=torch.float32)  # 对每个组的所有元素 (channels_per_group, height, width) 求平均
group_sum = group_sum * nelements_rank  # 加权后的局部和 = 局部均值 * 当前进程的元素数量
dist.all_reduce(group_sum)  # 收集并汇总所有进程的局部和，得到全局和
E = (group_sum / nelements)[:, :, None, None, None].to(x.dtype)  # 计算全局的均值 E

1
2
3
4
5
6

# 和计算均值同样的操作
group_var_sum = torch.empty((x.shape[0], self.num_groups), dtype=torch.float32, device=x.device)
torch.var(x, dim=(2,3,4), out=group_var_sum)  
group_var_sum = group_var_sum * nelements_rank
dist.all_reduce(group_var_sum)
var = (group_var_sum / nelements)[:, :, None, None, None].to(x.dtype)

1
2
3

x = (x - E) / torch.sqrt(var + self.eps)
x = x * self.weight[:, :, None, None, None] + self.bias[:, :, None, None, None]
return x

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26

class PatchConv2d(nn.Conv2d):
    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        kernel_size: _size_2_t,
        stride: _size_2_t = 1,
        padding: Union[str, _size_2_t] = 0,
        dilation: _size_2_t = 1,
        groups: int = 1,
        bias: bool = True,
        padding_mode: str = 'zeros',  # TODO: refine this type
        device=None,
        dtype=None,
        block_size: Union[int, Tuple[int, int]] = 0
    ) -> None:

        if isinstance(dilation, int):
            assert dilation == 1, "dilation is not supported in PatchConv2d"
        else:
            for i in dilation:
                assert i == 1, "dilation is not supported in PatchConv2d"
        self.block_size = block_size
        super().__init__(
            in_channels, out_channels, kernel_size, stride, padding, dilation, 
            groups, bias, padding_mode, device, dtype)

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12

def _conv_forward(self, input: Tensor, weight: Tensor, bias: Optional[Tensor]):
        bs, channels, h, w = input.shape

        world_size, rank = self._get_world_size_and_rank()

        if (world_size == 1):  # 处理非分布式情况
            if self.padding_mode != 'zeros':
                return F.conv2d(F.pad(input, self._reversed_padding_repeated_twice, mode=self.padding_mode),
                                weight, bias, self.stride,
                                _pair(0), self.dilation, self.groups)
            return F.conv2d(input, weight, bias, self.stride,
                            self.padding, self.dilation, self.groups)

1
2
3
4

patch_height_list = [torch.zeros(1, dtype=torch.int64, device=f"cuda:{rank}") for _ in range(dist.get_world_size())]
dist.all_gather(patch_height_list, torch.tensor([h], dtype=torch.int64, device=f"cuda:{rank}"))  # 收集所有进程的输入高度
patch_height_index = self._calc_patch_height_index(patch_height_list)  # 计算每个进程块的起始高度和结束高度的索引
halo_width = self._calc_halo_width_in_h_dim(rank, patch_height_index, self.kernel_size[0], self.padding[0], self.stride[0])  # 计算当前进程块的上下 halo 区域的宽度

1
2
3
4
5
6
7

prev_bottom_halo_width: int = 0
next_top_halo_width: int = 0
if rank != 0:
    prev_bottom_halo_width = self._calc_bottom_halo_width(rank - 1, patch_height_index, self.kernel_size[0], self.padding[0], self.stride[0])
if rank != world_size - 1:
    next_top_halo_width = self._calc_top_halo_width(rank + 1, patch_height_index, self.kernel_size[0], self.padding[0], self.stride[0])
    next_top_halo_width = max(0, next_top_halo_width)

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19

to_next = None
to_prev = None
top_halo_recv = None
bottom_halo_recv = None
if next_top_halo_width > 0:
    bottom_halo_send = input[:, :, -next_top_halo_width:, :].contiguous()
    to_next = dist.isend(bottom_halo_send, rank + 1)

if halo_width[0] > 0:  # not rank 0
    top_halo_recv = torch.empty([bs, channels, halo_width[0], w], dtype=input.dtype, device=f"cuda:{rank}")
    dist.recv(top_halo_recv, rank - 1)

if prev_bottom_halo_width > 0:  # not rank N-1
    top_halo_send = input[:, :, :prev_bottom_halo_width, :].contiguous()
    to_prev = dist.isend(top_halo_send, rank - 1)

if halo_width[1] > 0:
    bottom_halo_recv = torch.empty([bs, channels, halo_width[1], w], dtype=input.dtype, device=f"cuda:{rank}")
    dist.recv(bottom_halo_recv, rank + 1)

1
2
3
4
5
6
7

if halo_width[0] < 0:  # Remove redundancy at the top of the input
input = input[:, :, -halo_width[0]:, :]
 
if top_halo_recv is not None:  # concat the halo region to the input tensor 
    input = torch.cat([top_halo_recv, input], dim=-2)
if bottom_halo_recv is not None:
    input = torch.cat([input, bottom_halo_recv], dim=-2)

1
2
3
4

if to_next is not None:
    to_next.wait()
if to_prev is not None:
    to_prev.wait()

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47

padding = self._adjust_padding_for_patch(self._reversed_padding_repeated_twice, rank=rank, world_size=world_size)
if self.block_size == 0 or (h <= self.block_size and w <= self.block_size):
    if self.padding_mode != 'zeros':
        conv_res = F.conv2d(F.pad(input, padding, mode=self.padding_mode),
                            weight, bias, self.stride, _pair(0), self.dilation, self.groups)
    else:
        conv_res = F.conv2d(input, weight, bias, self.stride,
                            self.padding, self.dilation, self.groups)
    return conv_res
else:
    if self.padding_mode != "zeros":
        input = F.pad(input, padding, mode=self.padding_mode)
    elif self.padding != 0:
        input = F.pad(input, padding, mode="constant")

    _, _, h, w = input.shape
    num_chunks_in_h = (h + self.block_size - 1) // self.block_size  # h 维度的 block 数量
    num_chunks_in_w = (w + self.block_size - 1) // self.block_size  # w ...
    unit_chunk_size_h = h // num_chunks_in_h
    unit_chunk_size_w = w // num_chunks_in_w
  
outputs = []
for idx_h in range(num_chunks_in_h):
    inner_output = []
    for idx_w in range(num_chunks_in_w):
        start_w = idx_w * unit_chunk_size_w
        start_h = idx_h * unit_chunk_size_h
        end_w = (idx_w + 1) * unit_chunk_size_w
        end_h = (idx_h + 1) * unit_chunk_size_h

        # 计算每个块的开始和结束索引，调整块的边界
        # ...

        # 对当前块执行卷积操作
        inner_output.append(
            F.conv2d(
                input[:, :, start_h:end_h, start_w:end_w],
                weight,
                bias,
                self.stride,
                0,
                self.dilation,
                self.groups,
            )
        )
    outputs.append(torch.cat(inner_output, dim=-1))
return torch.cat(outputs, dim=-2)