Skip to content

Commit

Permalink
Temporarily disable test requring large shared memory. (csarofeen#1802)
Browse files Browse the repository at this point in the history
  • Loading branch information
shmsong authored Jul 5, 2022
1 parent 5f375d0 commit 37c579e
Showing 1 changed file with 14 additions and 0 deletions.
14 changes: 14 additions & 0 deletions torch/csrc/jit/codegen/cuda/test/test_gpu.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -9432,6 +9432,13 @@ TEST_F(NVFuserTest, FusionPersistentSoftmaxLocalSmem_CUDA) {
const int64_t dimx = 1024;
const int64_t dimy = 16384;

auto properties = at::cuda::getDeviceProperties(0);
// Require 70KB of smem to run test
const size_t required_smem_size = 70 << 10;
if (properties->sharedMemPerBlockOptin < required_smem_size) {
GTEST_SKIP() << "not enough shared memory space on device to run test";
}

auto options = at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0);
at::Tensor aten_input = at::randn({dimx, dimy}, options);
at::Tensor aten_static_in = aten_input.narrow(1, 0, static_size);
Expand Down Expand Up @@ -9631,6 +9638,13 @@ TEST_F(NVFuserTest, FusionPersistentNormLocalShared_CUDA) {
fe.compileFusion(&fusion, aten_inputs);
fe.runFusion(aten_inputs, {cg_static_out, cg_dynamic_out});

auto properties = at::cuda::getDeviceProperties(0);
// Require 70KB of smem to run test
const size_t required_smem_size = 70 << 10;
if (properties->sharedMemPerBlockOptin < required_smem_size) {
GTEST_SKIP() << "not enough shared memory space on device to run test";
}

auto at_mu = at::mean(aten_input.to(at::kDouble), -1).unsqueeze(1);
auto at_var = at::var(aten_input.to(at::kDouble), -1, false).unsqueeze(1);
auto at_rvar = at::rsqrt(at::add(at_var, kEps));
Expand Down

0 comments on commit 37c579e

Please sign in to comment.