Compile NNlib.maxpool and NNlib.meanpool (#102)

Author: Pangoraw

ext/ReactantNNlibExt.jl

@@ -109,4 +109,84 @@ function NNlib.conv(
     return Reactant.TracedRArray{T,N}((), Reactant.MLIR.IR.result(conv), output_shape)
 end
 
+function reduce_window(f, x::Reactant.TracedRArray{T,N}, pdims; init) where {T,N}
+    num_spatial_dims = N - 2
+    input_spatial_dims = 1:num_spatial_dims
+
+    dilation = NNlib.dilation(pdims)
+    kernel_size = NNlib.kernel_size(pdims)
+    stride = NNlib.stride(pdims)
+    padding = NNlib.padding(pdims)
+
+    window_dimensions = [kernel_size..., 1, 1]
+    window_strides = [stride..., 1, 1]
+    window_dilations = [dilation..., 1, 1]
+
+    output_spatial_shapes = map(input_spatial_dims) do i
+        K = kernel_size[i]
+        pl, pr = padding[2i - 1], padding[2i]
+        d = dilation[i]
+        s = stride[i]
+
+        (size(x, i) + pl + pr - d * (K - 1) - 1) ÷ s + 1
+    end
+
+    padding = Reactant.MLIR.IR.DenseElementsAttribute(
+        reshape([padding..., 0, 0, 0, 0], (N, 2))
+    )
+
+    output_shape = (output_spatial_shapes..., size(x, N - 1), size(x, N))
+    result_type = Reactant.MLIR.IR.TensorType(output_shape, Reactant.MLIR.IR.Type(T))
+
+    unranked = Reactant.MLIR.IR.TensorType((), eltype(Reactant.MLIR.IR.type(x.mlir_data)))
+    body =
+        let body = Reactant.MLIR.IR.Region(),
+            loc = Reactant.MLIR.IR.Location(),
+            block = Reactant.MLIR.IR.Block([unranked, unranked], [loc, loc])
+
+            Reactant.MLIR.IR.block!(block) do
+                red = f(
+                    Reactant.MLIR.IR.argument(block, 1),
+                    Reactant.MLIR.IR.argument(block, 2);
+                    result=nothing,
+                )
+                Reactant.MLIR.Dialects.stablehlo.return_([Reactant.MLIR.IR.result(red)])
+            end
+            push!(body, block)
+
+            body
+        end
+
+    attr = fill(Reactant.MLIR.IR.Attribute(init), unranked)
+    init_value = Reactant.MLIR.IR.result(
+        Reactant.MLIR.Dialects.stablehlo.constant(; value=attr)
+    )
+    reduction = Reactant.MLIR.Dialects.stablehlo.reduce_window(
+        [x.mlir_data],
+        [init_value];
+        result_0=[result_type],
+        window_dimensions,
+        window_strides,
+        window_dilations,
+        padding,
+        body,
+    )
+
+    return Reactant.TracedRArray{T,N}(
+        (), Reactant.MLIR.IR.result(reduction), size(result_type)
+    )
+end
+
+function NNlib.maxpool(x::Reactant.TracedRArray{T}, pdims::NNlib.PoolDims) where {T}
+    return reduce_window(
+        Reactant.MLIR.Dialects.stablehlo.maximum, x, pdims; init=typemin(T)
+    )
 end
+
+function NNlib.meanpool(x::Reactant.TracedRArray{T}, pdims::NNlib.PoolDims) where {T}
+    numel = prod(NNlib.kernel_size(pdims))
+    return reduce_window(Reactant.MLIR.Dialects.stablehlo.add, x, pdims; init=zero(T)) ./
+           T(numel)
+end
+
+end # module ReactantNNlibExt

test/nn.jl

@@ -84,3 +84,15 @@ mean((out2[1, :] .> 0.5) .== truth)  # accuracy 94% so far!
 
     @test res_reactant ≈ res
 end
+
+@testset "$f" for f in (NNlib.meanpool, NNlib.maxpool)
+    img = randn(Float32, 224, 224, 3, 2)
+    img_reactant = Reactant.ConcreteRArray(img)
+
+    f_reactant = Reactant.compile(f, (img_reactant, (3, 3)))
+
+    res_reactant = f_reactant(img_reactant, (3, 3))
+    res = f(img, (3, 3))
+
+    @test res_reactant ≈ res
+end