Home
last modified time | relevance | path

Searched refs:torch_layout (Results 1 – 2 of 2) sorted by relevance

/external/pytorch/torch/_inductor/codegen/rocm/
Dck_universal_gemm_template.py28 def torch_layout_to_ck_layout(torch_layout): argument
29 if torch_layout.stride[-1] == 1:
31 elif torch_layout.stride[-2] == 1:
/external/pytorch/torch/_inductor/codegen/cuda/
Dgemm_template.py558 …def cutlass_layout(torch_layout: ir.Layout) -> "Optional[cutlass_lib.LayoutType]": # type: ignore…
573 if torch_layout.stride[-1] == 1:
575 elif torch_layout.stride[-2] == 1:
596 torch_layout: ir.Layout,
600 return CUTLASSGemmTemplate.cutlass_layout(torch_layout) == cutlass_layout
603 def set_alignment(torch_layout, op_element) -> bool: argument
617 alignment = cutlass_utils.get_max_alignment(torch_layout)