Add Tile config support to reduce op and validate reduction-axis constraint

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<std::vector<DimType>>());
+    if (unit_out_dims.ndims() != 4) {
+        ERR(PlanError, "Tile should have 4 elements");
+    }
     int axis = args_.at("Axis").value<int>();
     bool keep_dims = args_.at("KeepDim").value<bool>();
 
     // 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;
 }

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()