Skip to content

Transpose scheduler small dim sizes better support #1910

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 13 commits into from
Aug 23, 2022
Merged

Transpose scheduler small dim sizes better support #1910

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
13 commits
Select commit Hold shift + click to select a range
c7096b2
splitDims and mergeDims
zasdfgbnm Aug 12, 2022
b94a765
disable view
zasdfgbnm Aug 12, 2022
cc7b4ad
scheduler
zasdfgbnm Aug 12, 2022
55a4665
small tile sizes for small tensors
zasdfgbnm Aug 13, 2022
ba27a92
maybeBuildVirtualInnerDims
zasdfgbnm Aug 13, 2022
ea3fa00
fix
zasdfgbnm Aug 13, 2022
9e56452
save
zasdfgbnm Aug 13, 2022
79d0e8d
fix
zasdfgbnm Aug 13, 2022
97f16f0
Merge branch 'devel' of github.com:csarofeen/pytorch into merge-small…
zasdfgbnm Aug 22, 2022
fea8d1f
Note: [Supporting small transpose dimensions]
zasdfgbnm Aug 22, 2022
fdbb30e
shape_in_ref1
zasdfgbnm Aug 22, 2022
2a85f9f
comment
zasdfgbnm Aug 22, 2022
4938af2
save
zasdfgbnm Aug 23, 2022
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 8 additions & 0 deletions c10/util/hash.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -304,6 +304,14 @@ struct hash<std::tuple<Types...>> {
}
};

template <typename T1, typename T2>
Copy link
Owner

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Warning to @jjsjann123

struct hash<std::pair<T1, T2>> {
size_t operator()(const std::pair<T1, T2>& pair) const {
std::tuple<T1, T2> tuple = std::make_tuple(pair.first, pair.second);
return _hash_detail::simple_get_hash(tuple);
};
};

template <typename T>
struct hash<c10::ArrayRef<T>> {
size_t operator()(c10::ArrayRef<T> v) const {
Expand Down
1 change: 1 addition & 0 deletions test/cpp/jit/CMakeLists.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -103,6 +103,7 @@ if(USE_CUDA)
list(APPEND JIT_TEST_SRCS ${TORCH_SRC_DIR}/csrc/jit/codegen/cuda/test/test_gpu_view.cpp)
list(APPEND JIT_TEST_SRCS ${TORCH_SRC_DIR}/csrc/jit/codegen/cuda/test/test_gpu_transpose.cpp)
list(APPEND JIT_TEST_SRCS ${TORCH_SRC_DIR}/csrc/jit/codegen/cuda/test/test_gpu_rng.cu)
list(APPEND JIT_TEST_SRCS ${TORCH_SRC_DIR}/csrc/jit/codegen/cuda/test/test_gpu_scheduler_utils.cpp)
endif()

add_executable(test_jit
Expand Down
2 changes: 1 addition & 1 deletion torch/csrc/jit/codegen/cuda/ir_nodes.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1765,7 +1765,7 @@ void TensorDomain::split(
resetDomains();
}

// Merge "axis" and "axis+1" into 1 dimension
// Merge "axis_o" and "axis_i" into 1 dimension
void TensorDomain::merge(int axis_o, int axis_i) {
TORCH_INTERNAL_ASSERT(nDims() > 0, "Tried to do merge on a 0-dim domain");
if (axis_o < 0)
Expand Down
11 changes: 10 additions & 1 deletion torch/csrc/jit/codegen/cuda/scheduler/registry.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -829,7 +829,7 @@ class ReductionScheduler : public SchedulerEntry {

//! Check if the reduction heuristics apply in given fusion
static bool canScheduleCompileTime(Fusion* fusion) {
// Temporarily allow view in reduction scheduler
// Temporarily disallow view in reduction scheduler
// TODO Add more testing before enabling
auto view_tvs = scheduler_utils::getViewTVs(fusion);
if (view_tvs.size() > 0) {
Expand Down Expand Up @@ -1259,6 +1259,15 @@ class TransposeScheduler : public SchedulerEntry {
// Not enabling this yet. Needs more validation.
return false;
#if 0
// Temporarily disallow view in transpose scheduler
// TODO Add more testing before enabling
auto view_tvs = scheduler_utils::getViewTVs(fusion);
if (view_tvs.size() > 0) {
scheduler_debug_utils::canScheduleRejectReason(
ScheduleHeuristic::Reduction, "No support for view op");
return false;
}

if (!hasAtLeastTwoValidGroups(fusion)) {
scheduler_debug_utils::canScheduleRejectReason(
ScheduleHeuristic::Transpose,
Expand Down
Loading