From 2328706ce0da7cccb79db7fab1f340a5142cd61d Mon Sep 17 00:00:00 2001 From: Paul Date: Tue, 26 Nov 2024 17:39:12 +0000 Subject: [PATCH 1/8] use row major when building attributes --- src/mlir/IR/Attribute.jl | 45 ++++++++++++++++++++-------------------- test/basic.jl | 15 ++++++++++++++ 2 files changed, 38 insertions(+), 22 deletions(-) diff --git a/src/mlir/IR/Attribute.jl b/src/mlir/IR/Attribute.jl index 69bd9b99a1..2b27e69657 100644 --- a/src/mlir/IR/Attribute.jl +++ b/src/mlir/IR/Attribute.jl @@ -492,6 +492,8 @@ function Base.fill(::Core.Type{Attribute}, value, shape) return Base.fill(value, shaped_type) end +to_row_major(x) = permutedims(x, ndims(x):-1:1) + """ DenseElementsAttribute(array::AbstractArray) @@ -501,66 +503,66 @@ function DenseElementsAttribute(values::AbstractArray{Bool}) shaped_type = TensorType(size(values), Type(Bool)) return Attribute( API.mlirDenseElementsAttrBoolGet( - shaped_type, length(values), AbstractArray{Cint}(values) + shaped_type, length(values), AbstractArray{Cint}(to_row_major(values)) ), ) end function DenseElementsAttribute(values::AbstractArray{UInt8}) shaped_type = TensorType(size(values), Type(UInt8)) - return Attribute(API.mlirDenseElementsAttrUInt8Get(shaped_type, length(values), values)) + return Attribute(API.mlirDenseElementsAttrUInt8Get(shaped_type, length(values), to_row_major(values))) end function DenseElementsAttribute(values::AbstractArray{Int8}) shaped_type = TensorType(size(values), Type(Int8)) - return Attribute(API.mlirDenseElementsAttrInt8Get(shaped_type, length(values), values)) + return Attribute(API.mlirDenseElementsAttrInt8Get(shaped_type, length(values), to_row_major(values))) end function DenseElementsAttribute(values::AbstractArray{UInt16}) shaped_type = TensorType(size(values), Type(UInt16)) return Attribute( - API.mlirDenseElementsAttrUInt16Get(shaped_type, length(values), values) + API.mlirDenseElementsAttrUInt16Get(shaped_type, length(values), to_row_major(values)) ) end function DenseElementsAttribute(values::AbstractArray{Int16}) shaped_type = TensorType(size(values), Type(Int16)) - return Attribute(API.mlirDenseElementsAttrInt16Get(shaped_type, length(values), values)) + return Attribute(API.mlirDenseElementsAttrInt16Get(shaped_type, length(values), to_row_major(values))) end function DenseElementsAttribute(values::AbstractArray{UInt32}) shaped_type = TensorType(size(values), Type(UInt32)) return Attribute( - API.mlirDenseElementsAttrUInt32Get(shaped_type, length(values), values) + API.mlirDenseElementsAttrUInt32Get(shaped_type, length(values), to_row_major(values)) ) end function DenseElementsAttribute(values::AbstractArray{Int32}) shaped_type = TensorType(size(values), Type(Int32)) - return Attribute(API.mlirDenseElementsAttrInt32Get(shaped_type, length(values), values)) + return Attribute(API.mlirDenseElementsAttrInt32Get(shaped_type, length(values), to_row_major(values))) end function DenseElementsAttribute(values::AbstractArray{UInt64}) shaped_type = TensorType(size(values), Type(UInt64)) return Attribute( - API.mlirDenseElementsAttrUInt64Get(shaped_type, length(values), values) + API.mlirDenseElementsAttrUInt64Get(shaped_type, length(values), to_row_major(values)) ) end function DenseElementsAttribute(values::AbstractArray{Int64}) shaped_type = TensorType(size(values), Type(Int64)) - return Attribute(API.mlirDenseElementsAttrInt64Get(shaped_type, length(values), values)) + return Attribute(API.mlirDenseElementsAttrInt64Get(shaped_type, length(values), to_row_major(values))) end function DenseElementsAttribute(values::AbstractArray{Float32}) shaped_type = TensorType(size(values), Type(Float32)) - return Attribute(API.mlirDenseElementsAttrFloatGet(shaped_type, length(values), values)) + return Attribute(API.mlirDenseElementsAttrFloatGet(shaped_type, length(values), to_row_major(values))) end function DenseElementsAttribute(values::AbstractArray{Float64}) shaped_type = TensorType(size(values), Type(Float64)) return Attribute( - API.mlirDenseElementsAttrDoubleGet(shaped_type, length(values), values) + API.mlirDenseElementsAttrDoubleGet(shaped_type, length(values), to_row_major(values)) ) end @@ -569,16 +571,15 @@ end function DenseElementsAttribute(values::AbstractArray{Float16}) shaped_type = TensorType(size(values), Type(Float16)) return Attribute( - API.mlirDenseElementsAttrFloat16Get(shaped_type, length(values), values) + API.mlirDenseElementsAttrFloat16Get(shaped_type, length(values), to_row_major(values)) ) end function DenseElementsAttribute(values::AbstractArray{<:Complex}) shaped_type = TensorType(size(values), Type(eltype(values))) - # TODO: row major return Attribute( API.mlirDenseElementsAttrRawBufferGet( - shaped_type, length(values) * Base.elsize(values), values + shaped_type, length(values) * Base.elsize(values), to_row_major(values) ), ) end @@ -592,7 +593,7 @@ function DenseElementsAttribute(values::AbstractArray{String}) # TODO may fail because `Type(String)` is not defined shaped_type = TensorType(size(values), Type(String)) return Attribute( - API.mlirDenseElementsAttrStringGet(shaped_type, length(values), values) + API.mlirDenseElementsAttrStringGet(shaped_type, length(values), to_row_major(values)) ) end @@ -663,25 +664,25 @@ function DenseArrayAttribute end @llvmversioned min = v"16" DenseArrayAttribute( values::AbstractArray{Bool}; context::Context=context() -) = Attribute(API.mlirDenseBoolArrayGet(context, length(values), values)) +) = Attribute(API.mlirDenseBoolArrayGet(context, length(values), AbstractArray{Cint}(to_row_major(values)))) @llvmversioned min = v"16" DenseArrayAttribute( values::AbstractArray{Int8}; context::Context=context() -) = Attribute(API.mlirDenseI8ArrayGet(context, length(values), values)) +) = Attribute(API.mlirDenseI8ArrayGet(context, length(values), to_row_major(values))) @llvmversioned min = v"16" DenseArrayAttribute( values::AbstractArray{Int16}; context::Context=context() -) = Attribute(API.mlirDenseI16ArrayGet(context, length(values), values)) +) = Attribute(API.mlirDenseI16ArrayGet(context, length(values), to_row_major(values))) @llvmversioned min = v"16" DenseArrayAttribute( values::AbstractArray{Int32}; context::Context=context() -) = Attribute(API.mlirDenseI32ArrayGet(context, length(values), values)) +) = Attribute(API.mlirDenseI32ArrayGet(context, length(values), to_row_major(values))) @llvmversioned min = v"16" DenseArrayAttribute( values::AbstractArray{Int64}; context::Context=context() -) = Attribute(API.mlirDenseI64ArrayGet(context, length(values), values)) +) = Attribute(API.mlirDenseI64ArrayGet(context, length(values), to_row_major(values))) @llvmversioned min = v"16" DenseArrayAttribute( values::AbstractArray{Float32}; context::Context=context() -) = Attribute(API.mlirDenseF32ArrayGet(context, length(values), values)) +) = Attribute(API.mlirDenseF32ArrayGet(context, length(values), to_row_major(values))) @llvmversioned min = v"16" DenseArrayAttribute( values::AbstractArray{Float64}; context::Context=context() -) = Attribute(API.mlirDenseF64ArrayGet(context, length(values), values)) +) = Attribute(API.mlirDenseF64ArrayGet(context, length(values), to_row_major(values))) @llvmversioned min = v"16" Attribute(values::AbstractArray) = DenseArrayAttribute(values) diff --git a/test/basic.jl b/test/basic.jl index d188c276e9..1e2ae57d5c 100644 --- a/test/basic.jl +++ b/test/basic.jl @@ -625,3 +625,18 @@ end @test y == x end end + +function f_row_major(x) + y = [1 2; 3 4] + if x isa Reactant.TracedRArray + y = Reactant.promote_to(Reactant.TracedRArray{eltype(x),2}, y) + end + return x .+ y +end + +@testset "array attributes: row major" begin + x = rand(Int, 2, 2) + x_ra = Reactant.to_rarray(x) + + @test @jit(f_row_major(x_ra)) ≈ f_row_major(x) +end From 42c0080f2a8881d40eaa42ed0cc9f8211f6bc4a4 Mon Sep 17 00:00:00 2001 From: Paul Date: Tue, 26 Nov 2024 17:42:10 +0000 Subject: [PATCH 2/8] format --- src/mlir/IR/Attribute.jl | 54 ++++++++++++++++++++++++++++++---------- 1 file changed, 41 insertions(+), 13 deletions(-) diff --git a/src/mlir/IR/Attribute.jl b/src/mlir/IR/Attribute.jl index 2b27e69657..82d800df3a 100644 --- a/src/mlir/IR/Attribute.jl +++ b/src/mlir/IR/Attribute.jl @@ -510,59 +510,79 @@ end function DenseElementsAttribute(values::AbstractArray{UInt8}) shaped_type = TensorType(size(values), Type(UInt8)) - return Attribute(API.mlirDenseElementsAttrUInt8Get(shaped_type, length(values), to_row_major(values))) + return Attribute( + API.mlirDenseElementsAttrUInt8Get(shaped_type, length(values), to_row_major(values)) + ) end function DenseElementsAttribute(values::AbstractArray{Int8}) shaped_type = TensorType(size(values), Type(Int8)) - return Attribute(API.mlirDenseElementsAttrInt8Get(shaped_type, length(values), to_row_major(values))) + return Attribute( + API.mlirDenseElementsAttrInt8Get(shaped_type, length(values), to_row_major(values)) + ) end function DenseElementsAttribute(values::AbstractArray{UInt16}) shaped_type = TensorType(size(values), Type(UInt16)) return Attribute( - API.mlirDenseElementsAttrUInt16Get(shaped_type, length(values), to_row_major(values)) + API.mlirDenseElementsAttrUInt16Get( + shaped_type, length(values), to_row_major(values) + ), ) end function DenseElementsAttribute(values::AbstractArray{Int16}) shaped_type = TensorType(size(values), Type(Int16)) - return Attribute(API.mlirDenseElementsAttrInt16Get(shaped_type, length(values), to_row_major(values))) + return Attribute( + API.mlirDenseElementsAttrInt16Get(shaped_type, length(values), to_row_major(values)) + ) end function DenseElementsAttribute(values::AbstractArray{UInt32}) shaped_type = TensorType(size(values), Type(UInt32)) return Attribute( - API.mlirDenseElementsAttrUInt32Get(shaped_type, length(values), to_row_major(values)) + API.mlirDenseElementsAttrUInt32Get( + shaped_type, length(values), to_row_major(values) + ), ) end function DenseElementsAttribute(values::AbstractArray{Int32}) shaped_type = TensorType(size(values), Type(Int32)) - return Attribute(API.mlirDenseElementsAttrInt32Get(shaped_type, length(values), to_row_major(values))) + return Attribute( + API.mlirDenseElementsAttrInt32Get(shaped_type, length(values), to_row_major(values)) + ) end function DenseElementsAttribute(values::AbstractArray{UInt64}) shaped_type = TensorType(size(values), Type(UInt64)) return Attribute( - API.mlirDenseElementsAttrUInt64Get(shaped_type, length(values), to_row_major(values)) + API.mlirDenseElementsAttrUInt64Get( + shaped_type, length(values), to_row_major(values) + ), ) end function DenseElementsAttribute(values::AbstractArray{Int64}) shaped_type = TensorType(size(values), Type(Int64)) - return Attribute(API.mlirDenseElementsAttrInt64Get(shaped_type, length(values), to_row_major(values))) + return Attribute( + API.mlirDenseElementsAttrInt64Get(shaped_type, length(values), to_row_major(values)) + ) end function DenseElementsAttribute(values::AbstractArray{Float32}) shaped_type = TensorType(size(values), Type(Float32)) - return Attribute(API.mlirDenseElementsAttrFloatGet(shaped_type, length(values), to_row_major(values))) + return Attribute( + API.mlirDenseElementsAttrFloatGet(shaped_type, length(values), to_row_major(values)) + ) end function DenseElementsAttribute(values::AbstractArray{Float64}) shaped_type = TensorType(size(values), Type(Float64)) return Attribute( - API.mlirDenseElementsAttrDoubleGet(shaped_type, length(values), to_row_major(values)) + API.mlirDenseElementsAttrDoubleGet( + shaped_type, length(values), to_row_major(values) + ), ) end @@ -571,7 +591,9 @@ end function DenseElementsAttribute(values::AbstractArray{Float16}) shaped_type = TensorType(size(values), Type(Float16)) return Attribute( - API.mlirDenseElementsAttrFloat16Get(shaped_type, length(values), to_row_major(values)) + API.mlirDenseElementsAttrFloat16Get( + shaped_type, length(values), to_row_major(values) + ), ) end @@ -593,7 +615,9 @@ function DenseElementsAttribute(values::AbstractArray{String}) # TODO may fail because `Type(String)` is not defined shaped_type = TensorType(size(values), Type(String)) return Attribute( - API.mlirDenseElementsAttrStringGet(shaped_type, length(values), to_row_major(values)) + API.mlirDenseElementsAttrStringGet( + shaped_type, length(values), to_row_major(values) + ), ) end @@ -664,7 +688,11 @@ function DenseArrayAttribute end @llvmversioned min = v"16" DenseArrayAttribute( values::AbstractArray{Bool}; context::Context=context() -) = Attribute(API.mlirDenseBoolArrayGet(context, length(values), AbstractArray{Cint}(to_row_major(values)))) +) = Attribute( + API.mlirDenseBoolArrayGet( + context, length(values), AbstractArray{Cint}(to_row_major(values)) + ), +) @llvmversioned min = v"16" DenseArrayAttribute( values::AbstractArray{Int8}; context::Context=context() ) = Attribute(API.mlirDenseI8ArrayGet(context, length(values), to_row_major(values))) From ba6bb6886f389fad03db759ccd4f9587627d392f Mon Sep 17 00:00:00 2001 From: Paul Date: Tue, 26 Nov 2024 17:51:51 +0000 Subject: [PATCH 3/8] opt for d=1 --- src/mlir/IR/Attribute.jl | 1 + 1 file changed, 1 insertion(+) diff --git a/src/mlir/IR/Attribute.jl b/src/mlir/IR/Attribute.jl index 82d800df3a..81355ccf77 100644 --- a/src/mlir/IR/Attribute.jl +++ b/src/mlir/IR/Attribute.jl @@ -493,6 +493,7 @@ function Base.fill(::Core.Type{Attribute}, value, shape) end to_row_major(x) = permutedims(x, ndims(x):-1:1) +to_row_major(x::AbstractVector) = x """ DenseElementsAttribute(array::AbstractArray) From cae90dd7efb4a9d60657a20dc787ef47fff772ee Mon Sep 17 00:00:00 2001 From: Paul Date: Thu, 28 Nov 2024 09:06:01 +0000 Subject: [PATCH 4/8] reproducable test --- test/basic.jl | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/test/basic.jl b/test/basic.jl index 1e2ae57d5c..3796b7275c 100644 --- a/test/basic.jl +++ b/test/basic.jl @@ -627,7 +627,7 @@ end end function f_row_major(x) - y = [1 2; 3 4] + y = [1 2; 3 4; 5 6] if x isa Reactant.TracedRArray y = Reactant.promote_to(Reactant.TracedRArray{eltype(x),2}, y) end @@ -635,7 +635,7 @@ function f_row_major(x) end @testset "array attributes: row major" begin - x = rand(Int, 2, 2) + x = zeros(Int, 3, 2) x_ra = Reactant.to_rarray(x) @test @jit(f_row_major(x_ra)) ≈ f_row_major(x) From 2b896981f4349c46a5626213f96d4a71014254ff Mon Sep 17 00:00:00 2001 From: Paul Date: Fri, 29 Nov 2024 16:19:32 +0000 Subject: [PATCH 5/8] workaround for 0 dim array --- src/mlir/IR/Attribute.jl | 1 + 1 file changed, 1 insertion(+) diff --git a/src/mlir/IR/Attribute.jl b/src/mlir/IR/Attribute.jl index 81355ccf77..d37e7c9862 100644 --- a/src/mlir/IR/Attribute.jl +++ b/src/mlir/IR/Attribute.jl @@ -494,6 +494,7 @@ end to_row_major(x) = permutedims(x, ndims(x):-1:1) to_row_major(x::AbstractVector) = x +to_row_major(x::AbstractArray{T,0}) where {T} = x """ DenseElementsAttribute(array::AbstractArray) From 8ed8c99baf6ebbd0f43d804d3e8c642e0ed1a76a Mon Sep 17 00:00:00 2001 From: Paul Date: Fri, 29 Nov 2024 21:59:51 +0000 Subject: [PATCH 6/8] transpose padding --- ext/ReactantNNlibExt.jl | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/ext/ReactantNNlibExt.jl b/ext/ReactantNNlibExt.jl index 07c8a8e83c..81b9e4f563 100644 --- a/ext/ReactantNNlibExt.jl +++ b/ext/ReactantNNlibExt.jl @@ -109,7 +109,7 @@ function NNlib.conv!( #! format: on padding = Reactant.MLIR.IR.DenseElementsAttribute( - reshape(collect(padding), (num_spatial_dims, 2)) + reshape(collect(padding), (num_spatial_dims, 2))' ) result_type = Reactant.MLIR.IR.TensorType(size(y), Reactant.MLIR.IR.Type(T)) @@ -163,7 +163,7 @@ function reduce_window(f, x::AnyTracedRArray{T,N}, pdims; init) where {T,N} end padding = Reactant.MLIR.IR.DenseElementsAttribute( - reshape([padding..., 0, 0, 0, 0], (N, 2)) + reshape([padding..., 0, 0, 0, 0], (N, 2))' ) output_shape = (output_spatial_shapes..., size(x, N - 1), size(x, N)) From 2188ca3abd07018391ab357f68f3fff11395e6d2 Mon Sep 17 00:00:00 2001 From: Paul Date: Fri, 29 Nov 2024 22:12:13 +0000 Subject: [PATCH 7/8] 2,N' -> N,2 --- ext/ReactantNNlibExt.jl | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/ext/ReactantNNlibExt.jl b/ext/ReactantNNlibExt.jl index 81b9e4f563..5d0c0f04b9 100644 --- a/ext/ReactantNNlibExt.jl +++ b/ext/ReactantNNlibExt.jl @@ -109,7 +109,7 @@ function NNlib.conv!( #! format: on padding = Reactant.MLIR.IR.DenseElementsAttribute( - reshape(collect(padding), (num_spatial_dims, 2))' + reshape(collect(padding), (2, num_spatial_dims))' ) result_type = Reactant.MLIR.IR.TensorType(size(y), Reactant.MLIR.IR.Type(T)) @@ -163,7 +163,7 @@ function reduce_window(f, x::AnyTracedRArray{T,N}, pdims; init) where {T,N} end padding = Reactant.MLIR.IR.DenseElementsAttribute( - reshape([padding..., 0, 0, 0, 0], (N, 2))' + reshape([padding..., 0, 0, 0, 0], (2, N))' ) output_shape = (output_spatial_shapes..., size(x, N - 1), size(x, N)) From a7f82fdf3f7c4ee6666dba9bc9b4c4e66d59b41b Mon Sep 17 00:00:00 2001 From: Paul Date: Fri, 29 Nov 2024 22:34:50 +0000 Subject: [PATCH 8/8] make_causal_mask --- ext/ReactantNNlibExt.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ext/ReactantNNlibExt.jl b/ext/ReactantNNlibExt.jl index 5d0c0f04b9..90b1cd2ccc 100644 --- a/ext/ReactantNNlibExt.jl +++ b/ext/ReactantNNlibExt.jl @@ -306,7 +306,7 @@ function NNlib.make_causal_mask(x::AnyTracedRArray; dims::Int=2) len = size(x, dims) # directly generating booleans were causing an incorrect constant attribute generation # but the optimized IR removes the type case so we are probably ok - mask = MLIR.IR.DenseElementsAttribute(collect(triu(fill(1, (len, len)))')) + mask = MLIR.IR.DenseElementsAttribute(collect(triu(fill(1, (len, len))))) return Reactant.promote_to( TracedRArray{Bool,2}, TracedRArray{Int,2}(