diff --git a/csrc/python_frontend/python_bindings.cpp b/csrc/python_frontend/python_bindings.cpp index 1a1ba8099e0..bcdd4950377 100644 --- a/csrc/python_frontend/python_bindings.cpp +++ b/csrc/python_frontend/python_bindings.cpp @@ -55,9 +55,6 @@ Vector define_vector_fn( std::vector args; size_t idx = 0; for (const auto& item : values) { - NVF_CHECK( - idx < 8, - "The specified vector size exceeds the max tensor size for nvfuser."); if (py::isinstance(item)) { auto int_value = py::cast(item); NVF_CHECK( @@ -866,9 +863,6 @@ void initNvFuserPythonBindings(PyObject* module) { fusion_def.def( "define_vector", [](FusionDefinition& self, size_t size) -> Vector { - NVF_CHECK( - size < 8, - "The specified vector size exceeds the max tensor size for nvfuser."); std::vector args; args.reserve(size); for (size_t i = 0; i < size; ++i) { diff --git a/python_tests/pytest_input_generators.py b/python_tests/pytest_input_generators.py index a963620bb00..ac282584eea 100644 --- a/python_tests/pytest_input_generators.py +++ b/python_tests/pytest_input_generators.py @@ -26,7 +26,6 @@ MINIMUM_SYMBOLIC_SIZE = -1 INT64_MAX = 2**63 - 1 MAX_TENSOR_DIMS = 8 -MAX_VECTOR_SIZE = 8 # Determine if a number is with desired Domain [low, high) @@ -468,42 +467,11 @@ def define_vector_constant_error_generator( "The value -2 at index 0 was neither symbolic(-1), zero_element(0), broadcast(1), or static(>1)", ) - check_max_vector_size = ErrorSample( - { - "values": [-1 for _ in range(MAX_VECTOR_SIZE + 1)], - }, - "The specified vector size exceeds the max tensor size for nvfuser.", - ) - error_cases = [ # FIXME: The above_size_range case gives a non-sensical error message. # "Unable to cast Python instance to C++ type (#define PYBIND11_DETAILED_ER" # check_above_size_range, check_below_size_range, - check_max_vector_size, - ] - - for es in error_cases: - yield SampleInput(**es.kwargs), es.ex_type, es.ex_str - - -def define_vector_input_error_generator( - op: OpInfo, dtype: torch.dtype, requires_grad: bool = False, **kwargs -): - """ - "define_vector", - [](FusionDefinition& self, size_t size) -> Vector { - """ - - check_max_vector_size = ErrorSample( - { - "size": (MAX_VECTOR_SIZE + 1), - }, - "The specified vector size exceeds the max tensor size for nvfuser.", - ) - - error_cases = [ - check_max_vector_size, ] for es in error_cases: diff --git a/python_tests/pytest_opinfos.py b/python_tests/pytest_opinfos.py index 53d6f0f80b3..e52eb7fc183 100644 --- a/python_tests/pytest_opinfos.py +++ b/python_tests/pytest_opinfos.py @@ -22,7 +22,6 @@ define_tensor_generator, define_tensor_error_generator, define_vector_constant_error_generator, - define_vector_input_error_generator, elementwise_binary_generator, _elementwise_binary_torch, elementwise_unary_generator, @@ -90,15 +89,6 @@ ) fusion_input_ops.append(define_vector_constant_opinfo) -define_vector_input_opinfo = OpInfo( - lambda fd: fd.define_vector, - "define_vector_input", - sample_input_generator=None, - error_input_generator=define_vector_input_error_generator, - fd_error_input_fn=api_test_fd_fn, -) -fusion_input_ops.append(define_vector_input_opinfo) - """ End Fusion Input Operations """ """ Start Unary-Float Operations """ diff --git a/python_tests/test_python_frontend.py b/python_tests/test_python_frontend.py index 9e2da903913..c3163997f98 100644 --- a/python_tests/test_python_frontend.py +++ b/python_tests/test_python_frontend.py @@ -519,6 +519,26 @@ def nvfuser_fusion( self.assertEqual(eager_out, nvf_out[0]) + def test_tensor_ndim(self): + shape = [2 for i in range(12)] + new_shape = shape[:9] + new_shape.append(8) + + inputs = [torch.randn(shape, device="cuda"), new_shape] + + def fusion_func(fd: FusionDefinition): + t0 = fd.from_pytorch(inputs[0]) + n_shape = fd.define_vector(10) + + t1 = fd.ops.reshape(t0, n_shape) + t2 = fd.ops.sum(t1, axes=[3]) + + fd.add_output(t2) + + nvf_out, _ = self.exec_nvfuser(fusion_func, inputs) + eager_out = torch.sum(inputs[0].reshape(new_shape), dim=3) + self.assertEqual(eager_out, nvf_out[0]) + # Testing a scenario where a broadcast requires a symbolic output shape def test_tensor_shape(self): inputs = [ diff --git a/test/test_gpu_fused_reduction.cpp b/test/test_gpu_fused_reduction.cpp index a06d6d54f49..476da1981ba 100644 --- a/test/test_gpu_fused_reduction.cpp +++ b/test/test_gpu_fused_reduction.cpp @@ -2559,4 +2559,29 @@ TEST_F(NVFuserTest, FusionCrossEntropyGatherPattern_CUDA) { testValidate(&fusion, cg_outputs, inputs, {ref}, __LINE__, __FILE__); } +TEST_F(NVFuserTest, FusionTensorRankLimit) { + auto fusion = std::make_unique(); + FusionGuard fg(fusion.get()); + + std::vector input_shape; + for (__attribute__((unused)) auto i : c10::irange(12)) { + input_shape.push_back(3); + } + + auto tv0 = makeSymbolicTensor(input_shape.size()); + fusion->addInput(tv0); + auto tv1 = sum(tv0, {3}); + fusion->addOutput(tv1); + + auto options = at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0); + at::Tensor t0 = at::randn(input_shape, options); + std::vector aten_inputs({t0}); + + FusionExecutorCache executor_cache(std::move(fusion)); + auto cg_outputs = executor_cache.runFusionWithInputs(aten_inputs); + + testValidate( + executor_cache.fusion(), cg_outputs, aten_inputs, __LINE__, __FILE__); +} + } // namespace nvfuser