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

Function get_flt_align

dnn/scripts/cutlass_generator/conv2d_operation.py:731–753  ·  view source on GitHub ↗
(tile: TileDescription)

Source from the content-addressed store, hash-verified

729
730 # rule for filter alignment
731 def get_flt_align(tile: TileDescription) -> int:
732 nonlocal flt_align
733 if (
734 tile.math_instruction.opcode_class == OpcodeClass.Simt
735 and tile.math_instruction.element_accumulator == DataType.s32
736 ):
737 thread_num = (
738 tile.warp_count[0] * tile.warp_count[1] * tile.warp_count[2] * 32
739 )
740 flt_block = (
741 tile.threadblock_shape[0]
742 * tile.threadblock_shape[2]
743 * DataTypeSize[tile.math_instruction.element_a]
744 )
745 load_per_thread = flt_block // thread_num
746 if load_per_thread >= 128:
747 flt_align = 128
748 elif load_per_thread >= 64:
749 flt_align = 64
750 else:
751 assert load_per_thread >= 32
752 flt_align = 32
753 return flt_align
754
755 def get_dst_align(tile: TileDescription, out_layout: LayoutType) -> int:
756 nonlocal dst_align

Callers 1

GenerateConv2dFunction · 0.85

Calls

no outgoing calls

Tested by

no test coverage detected