NNlib: handle kernel flip with reverse instead of window_reversal (#115)

* NNlib: handle kernel flip with reverse instead of window_reversal

* add conv flip test

Author: Pangoraw

ext/ReactantNNlibExt.jl

@@ -93,12 +93,20 @@ function NNlib.conv(
     )
     result_type = Reactant.MLIR.IR.TensorType(output_shape, Reactant.MLIR.IR.Type(T))
 
+    weight = W.mlir_data
+    if !flipkernel
+        weight = Reactant.MLIR.IR.result(
+            Reactant.MLIR.Dialects.stablehlo.reverse(
+                weight; dimensions=collect(kernel_spatial_dims .- 1)
+            ),
+        )
+    end
+
     conv = Reactant.MLIR.Dialects.stablehlo.convolution(
         x.mlir_data,
-        W.mlir_data;
+        weight;
         result_0=result_type,
         window_strides=collect(stride),
-        window_reversal=collect(fill(flipkernel, num_spatial_dims)),
         padding,
         dimension_numbers,
         lhs_dilation=1,

test/nn.jl

@@ -88,6 +88,22 @@ mean((out2[1, :] .> 0.5) .== truth)  # accuracy 94% so far!
     @test res_reactant ≈ nn_conv(img)
 end
 
+@testset "conv 1d: flip" begin
+    x = [1; 2; 3;;;]
+    W = [1; 2; 3;;;]
+
+    xx = Reactant.ConcreteRArray(x)
+    WW = Reactant.ConcreteRArray(W)
+
+    conv_noflip(x, W) = NNlib.conv(x, W; pad=1, flipped=true)
+    conv_flip(x, W) = NNlib.conv(x, W; pad=1, flipped=false)
+
+    @test Reactant.compile(conv_noflip, (xx, WW))(xx, WW) ==
+          [0*1+1*2+2*3; 1*1+2*2+3*3; 1*2+2*3+3*0;;;]
+    @test Reactant.compile(conv_flip, (xx, WW))(xx, WW) ==
+          [3*0+2*1+1*2; 3*1+2*2+1*3; 3*2+2*3+1*0;;;]
+end
+
 @testset "$f" for f in (NNlib.meanpool, NNlib.maxpool)
     img = randn(Float32, 224, 224, 3, 2)
     img_reactant = Reactant.ConcreteRArray(img)