Add support for include-what-you-use (#71114)

Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/71114

`include-what-you-use` or `iwyu` is a clang-based tool that looks at
the code's AST to figure out which symbols need to be included and
with the help of user-defined mappings it suggests the include
files that are actually needed.

This is very nice for the per-operator headers build because it give
you a list of exactly the `ATen/ops` headers needed by the file. You
still need to manually write the include-guards etc. but at least this
automates the most tedious part.

The header mappings aren't perfect yet so it will still suggest you
include basic c10 components everywhere instead of taking it
transitively from `TensorBase.h`. However, this does provide some
useful mappings and removes bad include paths from the build system
that were causing bad suggestions.

Test Plan: Imported from OSS

Reviewed By: ngimel

Differential Revision: D33949901

Pulled By: malfet

fbshipit-source-id: d5b015ef9e168bee4b8717b8e87ccc0608da62a1
(cherry picked from commit ecb2ffb35a)

aten/src/ATen/CMakeLists.txt

@@ -230,8 +230,6 @@ if(USE_ROCM)
 endif()
 
 list(APPEND ATen_CPU_INCLUDE ${CMAKE_CURRENT_SOURCE_DIR}/..)
-# so the build can find the generated header files
-list(APPEND ATen_CPU_INCLUDE ${CMAKE_CURRENT_BINARY_DIR})
 
 if(USE_TBB)
   if(USE_SYSTEM_TBB)

aten/src/ATen/Parallel.h

@@ -150,11 +150,11 @@ TORCH_API int intraop_default_num_threads();
 } // namespace at
 
 #if AT_PARALLEL_OPENMP
-#include <ATen/ParallelOpenMP.h>
+#include <ATen/ParallelOpenMP.h> // IWYU pragma: keep
 #elif AT_PARALLEL_NATIVE
-#include <ATen/ParallelNative.h>
+#include <ATen/ParallelNative.h> // IWYU pragma: keep
 #elif AT_PARALLEL_NATIVE_TBB
-#include <ATen/ParallelNativeTBB.h>
+#include <ATen/ParallelNativeTBB.h> // IWYU pragma: keep
 #endif
 
-#include <ATen/Parallel-inl.h>
+#include <ATen/Parallel-inl.h> // IWYU pragma: keep

aten/src/ATen/core/Dict.h

@@ -387,4 +387,4 @@ namespace torch {
   template<class Key, class Value> using Dict = c10::Dict<Key, Value>;
 }
 
-#include <ATen/core/Dict_inl.h>
+#include <ATen/core/Dict_inl.h>  // IWYU pragma: keep

aten/src/ATen/core/List.h

@@ -475,4 +475,4 @@ namespace torch {
   template<class T> using List = c10::List<T>;
 }
 
-#include <ATen/core/List_inl.h>
+#include <ATen/core/List_inl.h>  // IWYU pragma: keep

aten/src/ATen/core/TensorBase.h

@@ -7,6 +7,7 @@
 #include <c10/core/ScalarTypeToTypeMeta.h>
 #include <c10/core/Storage.h>
 #include <c10/core/TensorImpl.h>
+#include <c10/core/TensorOptions.h>
 #include <c10/core/UndefinedTensorImpl.h>
 #include <c10/core/WrapDimMinimal.h>
 #include <c10/util/Exception.h>
@@ -19,7 +20,6 @@
 #include <ATen/core/TensorAccessor.h>
 
 namespace c10 {
-struct TensorOptions;
 class Scalar;
 }
 

aten/src/ATen/core/function_schema.h

@@ -510,4 +510,4 @@ inline std::string toString(const FunctionSchema& schema) {
 
 } // namespace c10
 
-#include <ATen/core/function_schema_inl.h>
+#include <ATen/core/function_schema_inl.h>  // IWYU pragma: keep

aten/src/ATen/core/ivalue.h

@@ -1269,4 +1269,4 @@ struct TORCH_API WeakOrStrongTypePtr {
 
 } // namespace c10
 
-#include <ATen/core/ivalue_inl.h>
+#include <ATen/core/ivalue_inl.h>  // IWYU pragma: keep

c10/util/BFloat16.h

@@ -98,4 +98,4 @@ struct alignas(2) BFloat16 {
 
 } // namespace c10
 
-#include <c10/util/BFloat16-inl.h>
+#include <c10/util/BFloat16-inl.h>  // IWYU pragma: keep

c10/util/Half.h

@@ -508,4 +508,4 @@ C10_API std::ostream& operator<<(std::ostream& out, const Half& value);
 
 } // namespace c10
 
-#include <c10/util/Half-inl.h>
+#include <c10/util/Half-inl.h>  // IWYU pragma: keep

c10/util/complex.h

@@ -607,7 +607,7 @@ C10_CLANG_DIAGNOSTIC_POP()
 
 #define C10_INTERNAL_INCLUDE_COMPLEX_REMAINING_H
 // math functions are included in a separate file
-#include <c10/util/complex_math.h>
+#include <c10/util/complex_math.h>  // IWYU pragma: keep
 // utilities for complex types
-#include <c10/util/complex_utils.h>
+#include <c10/util/complex_utils.h>  // IWYU pragma: keep
 #undef C10_INTERNAL_INCLUDE_COMPLEX_REMAINING_H

caffe2/CMakeLists.txt

@@ -1052,7 +1052,6 @@ if(NOT INTERN_BUILD_MOBILE OR NOT BUILD_CAFFE2_MOBILE)
   set(ATen_CPU_INCLUDE
     ${TORCH_ROOT}/aten/src
     ${CMAKE_CURRENT_BINARY_DIR}/../aten/src
-    ${CMAKE_CURRENT_BINARY_DIR}/../aten/src/ATen
     ${CMAKE_BINARY_DIR}/aten/src)
 
     if(CMAKE_CXX_COMPILER_ID MATCHES "Clang" OR CMAKE_CXX_COMPILER_ID STREQUAL "GNU")

tools/iwyu/all.imp

@@ -0,0 +1,8 @@
+[
+    { ref: "system.imp" },
+    { ref: "c10.imp" },
+    { ref: "aten.imp" },
+    { ref: "qnnpack.imp" },
+    { ref: "torch.imp" },
+    { ref: "gtest.imp" },
+]

tools/iwyu/aten.imp

@@ -0,0 +1,36 @@
+[
+    { include: ["@<ATen/cpu/vec/vec512/.*>", private, "<ATen/cpu/vec/vec.h>", public] },
+    { include: ["@<ATen/cpu/vec/vec256/.*>", private, "<ATen/cpu/vec/vec.h>", public] },
+    { include: ["<ATen/cpu/vec/vec_base.h>", private, "<ATen/cpu/vec/vec.h>", public] },
+    { include: ["<ATen/cpu/vec/functional_base.h>", private, "<ATen/cpu/vec/functional.h>", public] },
+    { include: ["<ATen/cpu/vec/functional_bfloat16.h>", private, "<ATen/cpu/vec/functional.h>", public] },
+
+    { include: ["<ATen/core/Dict_inl.h>", private, "<ATen/core/Dict.h>", public] },
+    { include: ["<ATen/core/ivalue_inl.h>", private, "<ATen/core/ivalue.h>", public] },
+    { include: ["<ATen/core/DimVector.h>", public, "<ATen/DimVector.h>", public] },
+    { include: ["<ATen/core/Dimname.h>", public, "<ATen/Dimname.h>", public] },
+    { include: ["<ATen/core/symbol.h>", public, "<ATen/core/Dimname.h>", public] },
+    { include: ["<ATen/TensorIterator.h>", public, "<ATen/core/TensorIterator.h>", public] },
+
+    { include: ["<c10/core/WrapDimMinimal.h>", public, "<ATen/WrapDimUtils.h>", public] },
+    { include: ["<c10/util/SmallVector.h>", public, "<ATen/core/DimVector.h>", public] },
+
+    { include: ["<c10/util/MaybeOwned.h>", public, "<ATen/core/TensorBase.h>", public] },
+    { include: ["<c10/core/TensorImpl.h>", public, "<ATen/core/TensorBase.h>", public] },
+    { include: ["<c10/core/UndefinedTensorImpl.h>", public, "<ATen/core/TensorBase.h>", public] },
+    { include: ["<c10/core/TensorOptions.h>", public, "<ATen/core/TensorBase.h>", public] },
+
+    { include: ["<ATen/core/TensorBase.h>", public, "<ATen/core/Tensor.h>", public] },
+    { include: ["<ATen/core/TensorBody.h>", private, "<ATen/core/Tensor.h>", public] },
+    { symbol: ["at::Tensor", private, "<ATen/core/Tensor.h>", public] },
+
+    { include: ["<ATen/Parallel-inl.h>", private, "<ATen/Parallel.h>", public] },
+    { include: ["<ATen/ParallelOpenMP.h>", private, "<ATen/Parallel.h>", public] },
+    { include: ["<ATen/ParallelNative.h>", private, "<ATen/Parallel.h>", public] },
+    { include: ["<ATen/ParallelNativeTBB.h>", private, "<ATen/Parallel.h>", public] },
+
+    { include: ["<qnnpack/common.h>", public, "<pack_block_sparse.h>", public] },
+    { include: ["<qnnpack/math.h>", public, "<pack_block_sparse.h>", public] },
+    { include: ["<qnnpack/operator.h>", public, "<qnnpack_func.h>", public] },
+    { include: ["<qnnpack/log.h>", public, "<pytorch_qnnpack.h>", public] },
+]

tools/iwyu/c10.imp

@@ -0,0 +1,40 @@
+[
+    { include: [ "<c10/macros/cmake_macros.h>", private, "<c10/macros/Macros.h>", public ] },
+    { include: [ "<c10/macros/Export.h>", public, "<c10/macros/Macros.h>", public ] },
+
+    { include: [ "<c10/util/BFloat16-inl.h>", private, "<c10/util/BFloat16.h>", public ] },
+    { include: [ "<c10/util/Half-inl.h>", private, "<c10/util/Half.h>", public ] },
+
+    { include: [ "<c10/util/complex_math.h>", private, "<c10/util/complex.h>", public ] },
+    { include: [ "<c10/util/complex_utils.h>", private, "<c10/util/complex.h>", public ] },
+
+    { include: [ "<c10/util/logging_is_google_glog.h>", private, "<c10/util/Logging.h>", public ] },
+    { include: [ "<c10/util/logging_is_not_google_glog.h>", private, "<c10/util/Logging.h>", public ] },
+
+    { include: [ "<c10/util/qint32.h>", public, "<c10/core/ScalarType.h>", public ] },
+    { include: [ "<c10/util/qint8.h>", public, "<c10/core/ScalarType.h>", public ] },
+    { include: [ "<c10/util/quint2x4.h>", public, "<c10/core/ScalarType.h>", public ] },
+    { include: [ "<c10/util/quint4x4.h>", public, "<c10/core/ScalarType.h>", public ] },
+    { include: [ "<c10/util/quint8.h>", public, "<c10/core/ScalarType.h>", public ] },
+
+    { include: ["<c10/core/Backend.h>", public, "<c10/core/TensorOptions.h>", public] },
+    { include: ["<c10/core/Device.h>", public, "<c10/core/TensorOptions.h>", public] },
+    { include: ["<c10/core/DefaultDtype.h>", public, "<c10/core/TensorOptions.h>", public] },
+    { include: ["<c10/core/Layout.h>", public, "<c10/core/TensorOptions.h>", public] },
+    { include: ["<c10/core/MemoryFormat.h>", public, "<c10/core/TensorOptions.h>", public] },
+    { include: ["<c10/core/ScalarType.h>", public, "<c10/core/TensorOptions.h>", public] },
+    { include: ["<c10/util/Optional.h>", public, "<c10/core/TensorOptions.h>", public] },
+
+    { include: ["<c10/core/Allocator.h>", public, "<c10/core/StorageImpl.h>", public] },
+    { include: ["<c10/core/Device.h>", public, "<c10/core/StorageImpl.h>", public] },
+    { include: ["<c10/core/ScalarType.h>", public, "<c10/core/StorageImpl.h>", public] },
+    { include: ["<c10/core/StorageImpl.h>", public, "<c10/core/Storage.h>", public] },
+
+    { include: ["<c10/util/ArrayRef.h>", public, "<c10/core/TensorImpl.h>", public] },
+    { include: ["<c10/core/DispatchKeySet.h>", public, "<c10/core/TensorImpl.h>", public] },
+    { include: ["<c10/core/Layout.h>", public, "<c10/core/TensorImpl.h>", public] },
+    { include: ["<c10/core/Storage.h>", public, "<c10/core/TensorImpl.h>", public] },
+    { include: ["<c10/core/MemoryFormat.h>", public, "<c10/core/TensorImpl.h>", public] },
+
+    { symbol: [ "c10::complex", private, "<c10/util/complex.h>", public ] },
+]

tools/iwyu/fixup.py

@@ -0,0 +1,58 @@
+import sys
+import re
+
+QUOTE_INCLUDE_RE = re.compile(r'^#include "(.*)"')
+ANGLE_INCLUDE_RE = re.compile(r'^#include <(.*)>')
+
+# By default iwyu will pick the C include, but we prefer the C++ headers
+STD_C_HEADER_MAP = {
+    "<assert.h>": "<cassert>",
+    "<complex.h>": "<ccomplex>",
+    "<ctype.h>": "<cctype>",
+    "<errno.h>": "<cerrno>",
+    "<fenv.h>": "<cfenv>",
+    "<float.h>": "<cfloat>",
+    "<inttypes.h>": "<cinttypes>",
+    "<iso646.h>": "<ciso646>",
+    "<limits.h>": "<climits>",
+    "<locale.h>": "<clocale>",
+    "<math.h>": "<cmath>",
+    "<setjmp.h>": "<csetjmp>",
+    "<signal.h>": "<csignal>",
+    "<stdalign.h>": "<cstdalign>",
+    "<stdarg.h>": "<cstdarg>",
+    "<stdbool.h>": "<cstdbool>",
+    "<stddef.h>": "<cstddef>",
+    "<stdint.h>": "<cstdint>",
+    "<stdio.h>": "<cstdio>",
+    "<stdlib.h>": "<cstdlib>",
+    "<string.h>": "<cstring>",
+    "<tgmath.h>": "<ctgmath>",
+    "<time.h>": "<ctime>",
+    "<uchar.h>": "<cuchar>",
+    "<wchar.h>": "<cwchar>",
+    "<wctype.h>": "<cwctype>",
+}
+
+def main() -> None:
+    for line in sys.stdin:
+        # Convert all quoted includes to angle brackets
+        match = QUOTE_INCLUDE_RE.match(line)
+        if match is not None:
+            print(f"#include <{match.group(1)}>{line[match.end(0):]}", end='')
+            continue
+
+        match = ANGLE_INCLUDE_RE.match(line)
+        if match is not None:
+            path = f"<{match.group(1)}>"
+            new_path = STD_C_HEADER_MAP.get(path, path)
+            tail = line[match.end(0):]
+            if len(tail) > 1:
+                tail = ' ' + tail
+            print(f"#include {new_path}{tail}", end='')
+            continue
+
+        print(line, end='')
+
+if __name__ == "__main__":
+    main()

tools/iwyu/gtest.imp

@@ -0,0 +1,4 @@
+[
+    { include: [ "@<gtest/internal/.*>", private, "<gtest/gtest.h>", public ] },
+    { include: [ "@<gtest/[^/]*>", public, "<gtest/gtest.h>", public ] },
+]

tools/iwyu/qnnpack.imp

@@ -0,0 +1,2 @@
+[
+]

tools/iwyu/run.sh

@@ -0,0 +1,14 @@
+# Run include-what-you-use on a file or folder
+# e.g. tools/iwyu/run.sh aten/src/ATen/native/sparse/SparseBlas.cpp
+# Which will print suggested changes to the console
+#
+# Currently the include mappings aren't good enough to trust iwyu's
+# output e.g. we probably just want to include Tensor.h and trust it
+# brings in the c10 headers. So, for now, use iwyu as a guide and
+# update includes manually.
+
+TORCH_ROOT=$(dirname $(dirname $(dirname $(readlink -f $0))))
+
+iwyu_tool -p $TORCH_ROOT/build $@ -- -Wno-unknown-warning-option -Xiwyu \
+          --no_fwd_decls -Xiwyu --mapping_file=$TORCH_ROOT/tools/iwyu/all.imp \
+    | python $TORCH_ROOT/tools/iwyu/fixup.py

tools/iwyu/system.imp

@@ -0,0 +1,14 @@
+# Toolchain-specific mappings. These ones seem to be missing from iwyu defaults
+[
+    { include: [ "<bits/struct_stat.h>", private, "<sys/stat.h>", public ] },
+    { include: [ "<bits/stdint-uintn.h>", private, "<cstdint>", public ] },
+    { include: [ "<bits/stdint-intn.h>", private, "<cstdint>", public ] },
+    { include: [ "<bits/pthreadtypes.h>", private, "<pthread.h>", public ] },
+    { include: [ "<bits/exception.h>", private, "<stdexcept>", public ] },
+    { include: [ "<bits/refwrap.h>", private, "<reference_wrapper>", public ] },
+    { include: [ "<bits/shared_ptr.h>", private, "<memory>", public ] },
+    { include: [ "<bits/std_function.h>", private, "<functional>", public ] },
+
+    { symbol: [ "size_t", private, "<cstddef>", public ] },
+    { symbol: [ "ptrdiff_t", private, "<cstddef>", public ] },
+]

tools/iwyu/torch.imp

@@ -0,0 +1,2 @@
+[
+]