pytorch/tools
Mengwei Liu eebe0ee141 [Executorch][codegen] Add ETKernelIndex for aggregating all kernels for kernel (#102874)
Summary:
keys and change codegen to take ETKernelIndex

We are adding support for dtype and dim order specialized kernel registration. This requires us to reorganize `BackendIndex` (which is a `Dict[DispatchKey, Dict[OperatorName, BackendMetadata]]`) to be `Dict[OperatorName, Dict[ETKernelKey, BackendMetadata]]`. This PR adds new data structures in order to support this change:

* `ETKernelKey` to retrieve a certain kernel from the registry.
* `ETKernelIndex`, the dictionary from operator name to kernel key to kernel mapping.

Note that the codegen logic is not changed yet, we need subsequent diffs to actually generate code for different kernel keys.

Test Plan: Added tests

Reviewed By: Jack-Khuu

Differential Revision: D46407096

Pull Request resolved: https://github.com/pytorch/pytorch/pull/102874
Approved by: https://github.com/Jack-Khuu, https://github.com/kirklandsign
2023-06-03 17:23:42 +00:00
..
alerts Revert "fix alert upload action (#102840)" 2023-06-02 18:24:29 +00:00
amd_build
autograd fix bug in trace model when out-operator has more than one output (#101563) 2023-05-31 19:39:52 +00:00
bazel_tools
build/bazel Bump requests from 2.30.0 to 2.31.0 in /tools/build/bazel (#102059) 2023-05-26 07:01:22 -07:00
build_defs
code_analyzer
code_coverage
config
coverage_plugins_package
dynamo
gdb
iwyu
jit
linter improve diagnostics from bazel_linter.py (#101445) 2023-05-21 19:02:35 +00:00
lite_interpreter
lldb
onnx
pyi Hint Tensor._make_subclass as a staticmethod (#101961) 2023-05-22 12:42:50 +00:00
rules
rules_cc
setup_helpers
shared Test Reordering: Run previously failing tests first (#101123) 2023-05-16 19:57:54 +00:00
stats Update list of bots in upload_external_contrib_stats.py (#102786) 2023-06-02 18:34:22 +00:00
test [Executorch][codegen] Add ETKernelIndex for aggregating all kernels for kernel (#102874) 2023-06-03 17:23:42 +00:00
testing Add print statements to debug sharding error (#102713) 2023-06-01 22:38:28 +00:00
__init__.py
bazel.bzl update rules_python and let bazel install its own pip dependencies (#101405) 2023-05-23 06:20:33 +00:00
BUCK.bzl
BUCK.oss
build_libtorch.py
build_pytorch_libs.py
download_mnist.py
extract_scripts.py
gen_flatbuffers.sh
gen_vulkan_spv.py
generate_torch_version.py
generated_dirs.txt
git_add_generated_dirs.sh
git_reset_generated_dirs.sh
nightly.py
nvcc_fix_deps.py
pytorch.version
README.md
render_junit.py
substitute.py
update_masked_docs.py
vscode_settings.py

This folder contains a number of scripts which are used as part of the PyTorch build process. This directory also doubles as a Python module hierarchy (thus the __init__.py).

Overview

Modern infrastructure:

  • autograd - Code generation for autograd. This includes definitions of all our derivatives.
  • jit - Code generation for JIT
  • shared - Generic infrastructure that scripts in tools may find useful.
    • module_loader.py - Makes it easier to import arbitrary Python files in a script, without having to add them to the PYTHONPATH first.

Build system pieces:

  • setup_helpers - Helper code for searching for third-party dependencies on the user system.
  • build_pytorch_libs.py - cross-platform script that builds all of the constituent libraries of PyTorch, but not the PyTorch Python extension itself.
  • build_libtorch.py - Script for building libtorch, a standalone C++ library without Python support. This build script is tested in CI.

Developer tools which you might find useful:

Important if you want to run on AMD GPU:

  • amd_build - HIPify scripts, for transpiling CUDA into AMD HIP. Right now, PyTorch and Caffe2 share logic for how to do this transpilation, but have separate entry-points for transpiling either PyTorch or Caffe2 code.
    • build_amd.py - Top-level entry point for HIPifying our codebase.

Tools which are only situationally useful: