diff --git a/ark/ops/ops_reduce.cpp b/ark/ops/ops_reduce.cpp index f2ae5e78..2386062c 100644 --- a/ark/ops/ops_reduce.cpp +++ b/ark/ops/ops_reduce.cpp @@ -49,7 +49,7 @@ ModelOpReduce::ModelOpReduce(const std::string &type_name, ModelTensorRef input, } std::string ModelOpReduce::impl_name(const Json &config) const { - check_fields_config(config, {"NumWarps", "SramBytes", "ImplType"}); + check_fields_config(config, {"NumWarps", "SramBytes", "ImplType", "Tile"}); check_fields_args(args_, {"Axis", "KeepDim"}); std::string red_type; @@ -66,12 +66,21 @@ std::string ModelOpReduce::impl_name(const Json &config) const { int num_warps = config.at("NumWarps"); int sram_bytes = config.at("SramBytes"); std::string impl_type = config.at("ImplType"); + Dims unit_out_dims(config.at("Tile").get>()); + if (unit_out_dims.ndims() != 4) { + ERR(PlanError, "Tile should have 4 elements"); + } int axis = args_.at("Axis").value(); bool keep_dims = args_.at("KeepDim").value(); // Translate the axis value into 4D representation. axis += 4 - read_tensors_[0]->shape().ndims(); + auto udims4 = unit_out_dims.dims4(); + if (udims4[axis] != 1) { + ERR(PlanError, "Tile dimension on the reduction axis must be 1"); + } + if (impl_type == "WarpWise") { impl_type = "w"; if (axis != 3) { @@ -99,7 +108,7 @@ std::string ModelOpReduce::impl_name(const Json &config) const { vec_string(read_tensors_[0]->shape().dims4()), vec_string(output_strides.dims4()), vec_string(output_shape.dims4()), - vec_string(Dims(1, 1, 1, 1)), + vec_string(udims4), std::to_string(num_warps), std::to_string(sram_bytes), std::to_string(axis), @@ -122,6 +131,7 @@ Json ModelOpReduce::default_config([[maybe_unused]] const ArchRef arch) const { config["ImplType"] = "ElementWise"; config["SramBytes"] = 0; } + config["Tile"] = {1, 1, 1, 1}; config["NumTasks"] = result_tensors_[0]->shape().nelems(); return config; } diff --git a/python/unittest/ops/test_reduce.py b/python/unittest/ops/test_reduce.py index 10353c12..4895a883 100644 --- a/python/unittest/ops/test_reduce.py +++ b/python/unittest/ops/test_reduce.py @@ -64,3 +64,65 @@ def test_reduce_mean(dtype): atol = 1e-4 if dtype == torch.float32 else 1e-2 rtol = 1e-4 if dtype == torch.float32 else 1e-2 assert torch.allclose(result, expected, atol=atol, rtol=rtol) + + +def test_reduce_sum_fused_tile(): + """WarpWise reduce_sum with a multi-row Tile. + + rows > 1, reduce-axis dim = 1. + """ + shape = [1, 1, 4, 1024] + a = torch.randn(shape, dtype=torch.float32, device=DEVICE) * 0.1 + with ark.PlannerContext( + config={ + "NumWarps": 1, + "SramBytes": 256, + "ImplType": "WarpWise", + "Tile": [1, 1, 2, 1], + } + ): + result = ark.reduce_sum(a, axis=3).eval() + expected = torch.sum(a, dim=3, keepdim=True) + atol = shape[3] * 1e-5 + assert torch.allclose( + result, expected, atol=atol, rtol=1e-4 + ), f"max_diff={(result - expected).abs().max()}" + + +def test_reduce_sum_fused_tile_elementwise(): + """ElementWise reduce_sum on axis 2 with a non-default Tile.""" + shape = [1, 2, 8, 512] + a = torch.randn(shape, dtype=torch.float32, device=DEVICE) * 0.1 + with ark.PlannerContext( + config={ + "NumWarps": 1, + "SramBytes": 0, + "ImplType": "ElementWise", + "Tile": [1, 1, 1, 64], + } + ): + result = ark.reduce_sum(a, axis=2).eval() + expected = torch.sum(a, dim=2, keepdim=True) + atol = shape[2] * 1e-5 + assert torch.allclose( + result, expected, atol=atol, rtol=1e-4 + ), f"max_diff={(result - expected).abs().max()}" + + +def test_reduce_tile_axis_validation(): + """Tile dimension on the reduce axis != 1 must raise ark.PlanError.""" + shape = [1, 1, 4, 1024] + a = torch.randn(shape, dtype=torch.float32, device=DEVICE) * 0.1 + with pytest.raises( + ark.PlanError, + match="Tile dimension on the reduction axis must be 1", + ): + with ark.PlannerContext( + config={ + "NumWarps": 1, + "SramBytes": 256, + "ImplType": "WarpWise", + "Tile": [1, 1, 1, 4], + } + ): + ark.reduce_sum(a, axis=3).eval()