Refactor instance_norm operator to reduce in PyTorch instead of using tl.atomic_add in Triton

Author: xu-zhengzhong

src/ntops/kernels/instance_norm.py

@@ -9,114 +9,77 @@
 
 def arrangement(
     input,
+    mean,
+    var,
     running_mean,
     running_var,
-    tmp_mean,
-    tmp_var,
     weight,
     bias,
-    momentum,
     eps,
     output,
     num_normalized_elements,
     use_input_stats,
-    tracking_running_stats,
     dims,
     block_size=None,
 ):
-    def _arrange_per_channel_tensor(tensor):
+    if block_size is None:
+        block_size = ninetoothed.block_size()
+
+    def _arrange_channel_tensor(tensor):
         arranged = tensor.tile((1,))
         arranged.dtype = arranged.dtype.squeeze(0)
         arranged = arranged.unsqueeze(0)
         arranged = arranged.expand((input.shape[0], -1))
 
         return arranged
 
+    def _arrange_mean_or_var(tensor):
+        arranged = tensor.tile((1, 1))
+        arranged.dtype = arranged.dtype.squeeze((0, 1))
+
+        return arranged
+
     input_arranged, output_arranged = reduction_arrangement(
         input, output, dim=dims, block_size=block_size
     )
-    running_mean_arranged = _arrange_per_channel_tensor(running_mean)
-    running_var_arranged = _arrange_per_channel_tensor(running_var)
-    tmp_mean_arranged = _arrange_per_channel_tensor(tmp_mean)
-    tmp_var_arranged = _arrange_per_channel_tensor(tmp_var)
-    weight_arranged = _arrange_per_channel_tensor(weight)
-    bias_arranged = _arrange_per_channel_tensor(bias)
-    momentum_arranged = momentum
+    mean_arranged = _arrange_mean_or_var(mean)
+    var_arranged = _arrange_mean_or_var(var)
+    running_mean_arranged = _arrange_channel_tensor(running_mean)
+    running_var_arranged = _arrange_channel_tensor(running_var)
+    weight_arranged = _arrange_channel_tensor(weight)
+    bias_arranged = _arrange_channel_tensor(bias)
     eps_arranged = eps
     num_normalized_elements_arranged = num_normalized_elements
 
     if use_input_stats:
-        if tracking_running_stats:
-            return (
-                input_arranged,
-                running_mean_arranged,
-                running_var_arranged,
-                tmp_mean_arranged,
-                tmp_var_arranged,
-                weight_arranged,
-                bias_arranged,
-                momentum_arranged,
-                eps_arranged,
-                output_arranged,
-                num_normalized_elements_arranged,
-            )
-        else:
-            return (
-                input_arranged,
-                weight_arranged,
-                bias_arranged,
-                eps_arranged,
-                output_arranged,
-                num_normalized_elements_arranged,
-            )
-
-    return (
-        input_arranged,
-        running_mean_arranged,
-        running_var_arranged,
-        weight_arranged,
-        bias_arranged,
-        eps_arranged,
-        output_arranged,
-    )
-
-
-def application_without_tracking(
-    input,
-    weight,
-    bias,
-    eps,
-    output,
-    num_normalized_elements,
-):
-    _mean = ntl.zeros(input.dtype.shape, dtype=ntl.float32)
-
-    for i in range(input.shape[0]):
-        _mean += ntl.cast(input[i], ntl.float32)
-
-    mean = ntl.sum(_mean, 0) / num_normalized_elements
-
-    _var = ntl.zeros(input.dtype.shape, dtype=ntl.float32)
-
-    for i in range(input.shape[0]):
-        diff = ntl.cast(input[i], ntl.float32) - mean
-        diff = ntl.where(input[i].offsets(-1) < input.source.shape[-1], diff, 0)
-        _var += diff * diff
-
-    var = ntl.sum(_var, 0) / num_normalized_elements
-
-    application_with_mean_var(input, mean, var, weight, bias, eps, output)
-
-
-def application_with_tracking(
+        return (
+            input_arranged,
+            mean_arranged,
+            var_arranged,
+            weight_arranged,
+            bias_arranged,
+            eps_arranged,
+            output_arranged,
+            num_normalized_elements_arranged,
+        )
+    else:
+        return (
+            input_arranged,
+            running_mean_arranged,
+            running_var_arranged,
+            weight_arranged,
+            bias_arranged,
+            eps_arranged,
+            output_arranged,
+        )
+
+
+def application_using_input_stats(
     input,
-    running_mean,
-    running_var,
-    tmp_mean,
-    tmp_var,
+    mean,
+    var,
     weight,
     bias,
-    momentum,
     eps,
     output,
     num_normalized_elements,
@@ -137,22 +100,6 @@ def application_with_tracking(
 
     var = ntl.sum(_var, 0) / num_normalized_elements
 
-    ntl.atomic_add(
-        tmp_mean.source.data_ptr() + tmp_mean.offsets(0), ntl.cast(mean, ntl.float32)
-    )
-    ntl.atomic_add(
-        tmp_var.source.data_ptr() + tmp_mean.offsets(0), ntl.cast(var, ntl.float32)
-    )
-
-    ntl.debug_barrier()
-
-    if input[0].offsets(0) == 0:
-        tmp_mean = tmp_mean / input.source.shape[0]
-        tmp_var = tmp_var / input.source.shape[0]
-
-        running_mean = running_mean * (1 - momentum) + tmp_mean * momentum
-        running_var = running_var * (1 - momentum) + tmp_var * momentum
-
     application_with_mean_var(input, mean, var, weight, bias, eps, output)
 
 
@@ -174,7 +121,6 @@ def application_with_mean_var(
 def premake(
     ndim,
     use_input_stats,
-    tracking_running_stats,
     num_normalized_elements,
     dtype=None,
     block_size=None,
@@ -184,36 +130,30 @@ def premake(
     arrangement_ = functools.partial(
         arrangement,
         use_input_stats=use_input_stats,
-        tracking_running_stats=tracking_running_stats,
         dims=dims,
         block_size=block_size,
     )
 
     input = Tensor(ndim, other=0, dtype=dtype)
-    running_mean, running_var, tmp_mean, tmp_var, weight, bias = (
-        Tensor(1, dtype=dtype) for _ in range(6)
-    )
-    momentum, eps = (Tensor(0, dtype=ninetoothed.float64) for _ in range(2))
+    mean, var = (Tensor(2, dtype=dtype) for _ in range(2))
+    running_mean, running_var, weight, bias = (Tensor(1, dtype=dtype) for _ in range(4))
+    eps = Tensor(0, dtype=ninetoothed.float64)
     output = Tensor(ndim, dtype=dtype)
     num_normalized_elements = Tensor(0, constexpr=True, value=num_normalized_elements)
 
     if use_input_stats:
-        if tracking_running_stats:
-            application = application_with_tracking
-        else:
-            application = application_without_tracking
+        application = application_using_input_stats
     else:
         application = application_with_mean_var
 
     tensors = (
         input,
+        mean,
+        var,
         running_mean,
         running_var,
-        tmp_mean,
-        tmp_var,
         weight,
         bias,
-        momentum,
         eps,
         output,
         num_normalized_elements,

src/ntops/torch/instance_norm.py

@@ -22,23 +22,11 @@ def instance_norm(
     if bias is None:
         bias = torch.zeros(input.shape[1], device=input.device, dtype=input.dtype)
 
-    tracking_running_stats = False
+    has_running_stats = running_mean is not None and running_var is not None
 
-    if not use_input_stats:
-        assert running_mean is not None and running_var is not None, (
-            "`running_mean` and `running_var` must be provided when `use_input_stats=False`."
-        )
-        assert running_mean.shape == (input.shape[1],) and running_var.shape == (
-            input.shape[1],
-        ), "`running_mean` and `running_var` must have shape (C,)"
-    else:
-        if running_mean is not None and running_var is not None:
-            assert running_mean.shape == (input.shape[1],) and running_var.shape == (
-                input.shape[1],
-            ), "`running_mean` and `running_var` must have shape (C,)"
-            tracking_running_stats = True
-            tmp_mean = torch.zeros_like(running_mean)
-            tmp_var = torch.zeros_like(running_var)
+    if use_input_stats:
+        mean = torch.empty(input.shape[:2], device=input.device, dtype=input.dtype)
+        var = torch.empty(input.shape[:2], device=input.device, dtype=input.dtype)
 
     output = torch.empty_like(input)
 
@@ -47,35 +35,37 @@ def instance_norm(
         ntops.kernels.instance_norm.premake,
         input.ndim,
         use_input_stats,
-        tracking_running_stats,
         num_normalized_elements,
-        block_size=32,
+        dtype=input.dtype,
     )
 
     if use_input_stats:
-        if tracking_running_stats:
-            kernel(
-                input,
-                running_mean,
-                running_var,
-                tmp_mean,
-                tmp_var,
-                weight,
-                bias,
-                momentum,
-                eps,
-                output,
-                num_normalized_elements,
-            )
-        else:
-            kernel(
-                input,
-                weight,
-                bias,
-                eps,
-                output,
-                num_normalized_elements,
+        kernel(
+            input,
+            mean,
+            var,
+            weight,
+            bias,
+            eps,
+            output,
+            num_normalized_elements,
+        )
+
+        # We reduce in PyTorch instead of using tl.atomic_add in Triton because:
+        # 1. Triton blocks cannot synchronize to safely apply the momentum update after all additions finish.
+        # 2. N blocks atomically adding to the same C addresses creates severe memory contention.
+        if use_input_stats and has_running_stats:
+            batch_mean = mean.mean(0)
+            avg_vars = var.mean(0)
+
+            unbiased_var = (
+                (avg_vars) * num_normalized_elements / (num_normalized_elements - 1)
+                if num_normalized_elements > 1
+                else avg_vars
             )
+
+            running_mean.mul_(1 - momentum).add_(momentum * batch_mean)
+            running_var.mul_(1 - momentum).add_(momentum * unbiased_var)
     else:
         kernel(input, running_mean, running_var, weight, bias, eps, output)
 

tests/test_instance_norm.py

@@ -81,5 +81,6 @@ def test_instance_norm(
         assert torch.allclose(
             ninetoothed_running_mean, reference_running_mean, rtol=rtol, atol=atol
         )
-        # TODO: The running var is not close.
-        # assert torch.allclose(ninetoothed_running_var, reference_running_var, rtol=rtol, atol=atol)
+        assert torch.allclose(
+            ninetoothed_running_var, reference_running_var, rtol=rtol, atol=atol
+        )