[BOO] Add an env variable for toggling backward boo convolutions (iree-org#865)

This is temporarily set as defaulting to "do not use backward boo
kernels" as our performance on many backward convs is not competitive.

---------

Signed-off-by: zjgarvey <[email protected]>

Author: zjgarvey

iree/turbine/kernel/boo/ops/conv.py

@@ -4,6 +4,7 @@
 # See https://llvm.org/LICENSE.txt for license information.
 # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
+import os
 from typing import Sequence, Tuple
 
 import torch
@@ -14,6 +15,8 @@
     "boo_conv",
 ]
 
+BOO_USE_BACKWARD_KERNELS = int(os.getenv("BOO_USE_BACKWARD_KERNELS", "0"))
+
 
 @torch.library.custom_op("iree_turbine::boo_convolution", mutates_args=())
 def boo_convolution(
@@ -163,6 +166,44 @@ def _b(
     return input_grad, weight_grad, bias_grad
 
 
+def pytorch_convolution_backward(ctx, grad_output):
+    x, w = ctx.saved_tensors
+
+    mask = tuple((ctx.needs_input_grad[i] for i in range(3)))
+
+    # return to NCHW if necessary
+    rank = len(x.shape)
+    perm = [0] + [rank - 1] + list(range(1, rank - 1))
+    inv_perm = [0] + list(range(2, rank)) + [1]
+    if ctx.input_layout.endswith("C"):
+        x = x.permute(perm)
+    if ctx.kernel_layout.endswith("C"):
+        w = w.permute(perm)
+    if ctx.output_layout.endswith("C"):
+        grad_output = grad_output.permute(perm)
+
+    input_grad, weight_grad, bias_grad = torch.ops.aten.convolution_backward(
+        grad_output,
+        x,
+        w,
+        None,
+        ctx.stride,
+        ctx.padding,
+        ctx.dilation,
+        False,
+        [0] * len(ctx.stride),
+        ctx.groups,
+        mask,
+    )
+
+    if ctx.input_layout.endswith("C"):
+        input_grad = input_grad.permute(inv_perm)
+    if ctx.kernel_layout.endswith("C"):
+        weight_grad = weight_grad.permute(inv_perm)
+    # return `None` for attribute args
+    return input_grad, weight_grad, bias_grad, None, None, None, None, None, None, None
+
+
 def boo_convolution_backward(ctx, grad_output):
     x, w = ctx.saved_tensors
 
@@ -216,8 +257,14 @@ def boo_convolution_context(
     ctx.use_bias = b is not None
 
 
+_backward_to_register = (
+    boo_convolution_backward
+    if (BOO_USE_BACKWARD_KERNELS)
+    else pytorch_convolution_backward
+)
+
 boo_convolution.register_autograd(
-    boo_convolution_backward, setup_context=boo_convolution_context
+    _backward_to_register, setup_context=boo_convolution_context
 )
 
 

tests/kernel/boo/modeling/boo_conv_2d_test.py

@@ -4,6 +4,11 @@
 # See https://llvm.org/LICENSE.txt for license information.
 # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
+import os
+
+# enable backward boo kernels for testing
+os.environ["BOO_USE_BACKWARD_KERNELS"] = "1"
+
 import unittest
 import tempfile
 from pathlib import Path

tests/kernel/boo/ops/boo_conv_test.py

@@ -1,3 +1,8 @@
+import os
+
+# enable backward boo kernels for testing
+os.environ["BOO_USE_BACKWARD_KERNELS"] = "1"
+
 import unittest
 import pytest
 import tempfile