use row major when building attributes #307

Pangoraw · 2024-11-26T17:41:09Z

No description provided.

mofeing · 2024-11-27T21:01:50Z

so discussing with @wsmoses about this in today's meeting, we have reached the conclusion that the column- vs row-major layout is not forced by MLIR but it's dialect specific (e.g. a Julia MLIR dialect could still use column-major layout).

furthermore, I just found that column-major layouts can be represented by affine maps https://mlir.llvm.org/docs/Dialects/Builtin/#affine-map-layout do we know if affine maps can be used in tensors? the only examples I found are with memrefs. also, is StableHLO compatible with them?

I still believe that this PR is useful, but I would suggest to make the conversion from column- to row-major optional. maybe a kwarg that defaults to false?

test/basic.jl

Pangoraw · 2024-11-28T11:10:40Z

we have reached the conclusion that the column- vs row-major layout is not forced by MLIR but it's dialect specific

I agree but here this is for building shaped builtin attributes through the C API. Row major is needed for the data and shape to be interpreted accordingly. A dialect is then free to reinterpret the attribute as suited within its ops.

but I would suggest to make the conversion from column- to row-major optional. maybe a kwarg that defaults to false?

I don't really see why one would not want to convert to row major here. We can add it later when/if a potential julia dialect needs transposed attributes.

mofeing · 2024-11-29T00:28:55Z

I agree but here this is for building shaped builtin attributes through the C API. Row major is needed for the data and shape to be interpreted accordingly. A dialect is then free to reinterpret the attribute as suited within its ops.

i see your point but I'm not sure of all the consequence. how about we discuss it the next meeting?

wsmoses · 2024-11-29T03:41:41Z

Yeah okay I can agree with that logic (since then the bultin attribute conversion is nice — which is distinct from the builtin op semantics).

im okay with this PR

wsmoses · 2024-11-29T15:22:24Z

src/mlir/IR/Attribute.jl

@@ -492,6 +492,9 @@ 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)
+to_row_major(x::AbstractVector) = x


From the error logs it looks like this also needs a 0-dim specialization

mofeing · 2024-11-29T18:52:26Z

Yeah okay I can agree with that logic (since then the bultin attribute conversion is nice — which is distinct from the builtin op semantics).

im okay with this PR

do you know if affine maps can be associated to dense array / dense elements attributes?

wsmoses · 2024-11-29T18:53:29Z

Not attributes directly, but as part of a type

mofeing · 2024-11-29T19:03:18Z

okay, and can be associated to a tensor? the only examples I've seen are with memrefs

mofeing · 2024-11-29T19:05:50Z

(since then the bultin attribute conversion is nice — which is distinct from the builtin op semantics)

yeah, but we pass the values to stablehlo.constant as an attribute. the only change I would like is to make this conversion optional (I don't mind if the default is true or false) so that we can choose if to make the permutedims on the Julia side or on MLIR side.

wsmoses · 2024-11-29T19:08:04Z

I don’t think tensors do, and also that would still be an issue here since it having a different type means that you couldn’t use a memref where you needed a regular affine map.

still cool to use, but I think this is necessary regardless

Pangoraw · 2024-11-29T21:34:10Z

Note that the fix here not only affects stablehlo.constant but also other ops which have dense attributes such as here:

Reactant.jl/ext/ReactantNNlibExt.jl

Lines 111 to 113 in 45ae14f

    
           padding = Reactant.MLIR.IR.DenseElementsAttribute( 
        
               reshape(collect(padding), (num_spatial_dims, 2)) 
        
           )

If we want to try not having the transpose in julia, it should be possible by doing it at the promotion level:

function promote_to(x)
    cst = stablehlo.constant(reshape(x, :))
    cst = stablehlo.reshape(cst, reverse(size(x)))
    cst = stablehlo.transpose(cst, ndims(x):-1:1)
    return cst
end

wsmoses · 2024-11-29T21:40:00Z

Note that the fix here not only affects stablehlo.constant but also other ops which have dense attributes such as here:

Reactant.jl/ext/ReactantNNlibExt.jl

Lines 111 to 113 in 45ae14f

padding = Reactant.MLIR.IR.DenseElementsAttribute(

reshape(collect(padding), (num_spatial_dims, 2))

)

If we want to try not having the transpose in julia, it should be possible by doing it at the promotion level:
function promote_to(x)
    cst = stablehlo.constant(reshape(x, :))
    cst = stablehlo.reshape(cst, reverse(size(x)))
    cst = stablehlo.transpose(cst, ndims(x):-1:1)
    return cst
end

Yeah but frankly that’s the more intuitive thing anyways

mofeing

okay, then let's merge it like this and revisit in the future if we need it

mofeing · 2024-11-30T09:32:36Z

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))'


don't call '/adjoint because it will conjugate complex matrices

Suggested change

reshape(collect(padding), (2, num_spatial_dims))'

transpose(reshape(collect(padding), (2, num_spatial_dims)))

I did not know about that. As you said, it does not apply here but I will be cautious in the future 👍

mofeing · 2024-11-30T09:33:32Z

ext/ReactantNNlibExt.jl

@@ -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))'


Suggested change

reshape([padding..., 0, 0, 0, 0], (2, N))'

transpose(reshape([padding..., 0, 0, 0, 0], (2, N)))

mofeing

ahh, NNlib.padding just returns a tuple of ints... so it's alright if you call adjoint

wsmoses reviewed Nov 27, 2024

View reviewed changes

test/basic.jl Outdated Show resolved Hide resolved

Pangoraw added 4 commits November 28, 2024 09:06

use row major when building attributes

2328706

format

42c0080

opt for d=1

ba6bb68

reproducable test

cae90dd

Pangoraw force-pushed the row-major branch from e407436 to cae90dd Compare November 28, 2024 09:07

wsmoses reviewed Nov 29, 2024

View reviewed changes

workaround for 0 dim array

2b89698

This comment was marked as resolved.

Sign in to view

Pangoraw added 3 commits November 29, 2024 21:59

transpose padding

8ed8c99

2,N' -> N,2

2188ca3

make_causal_mask

a7f82fd

mofeing requested changes Nov 30, 2024

View reviewed changes

mofeing approved these changes Nov 30, 2024

View reviewed changes

Pangoraw merged commit 5731c0b into EnzymeAD:main Dec 1, 2024
24 of 38 checks passed

Pangoraw deleted the row-major branch December 1, 2024 09:32

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

use row major when building attributes #307

use row major when building attributes #307

Pangoraw commented Nov 26, 2024

mofeing commented Nov 27, 2024 •

edited

Loading

Pangoraw commented Nov 28, 2024

mofeing commented Nov 29, 2024

wsmoses commented Nov 29, 2024

wsmoses Nov 29, 2024

mofeing commented Nov 29, 2024

wsmoses commented Nov 29, 2024

mofeing commented Nov 29, 2024

mofeing commented Nov 29, 2024

wsmoses commented Nov 29, 2024

Pangoraw commented Nov 29, 2024

wsmoses commented Nov 29, 2024

This comment was marked as resolved.

mofeing left a comment

mofeing Nov 30, 2024

Pangoraw Dec 1, 2024

mofeing Nov 30, 2024

mofeing left a comment

	reshape(collect(padding), (2, num_spatial_dims))'
	transpose(reshape(collect(padding), (2, num_spatial_dims)))

	reshape([padding..., 0, 0, 0, 0], (2, N))'
	transpose(reshape([padding..., 0, 0, 0, 0], (2, N)))

use row major when building attributes #307

use row major when building attributes #307

Conversation

Pangoraw commented Nov 26, 2024

mofeing commented Nov 27, 2024 • edited Loading

Pangoraw commented Nov 28, 2024

mofeing commented Nov 29, 2024

wsmoses commented Nov 29, 2024

wsmoses Nov 29, 2024

Choose a reason for hiding this comment

mofeing commented Nov 29, 2024

wsmoses commented Nov 29, 2024

mofeing commented Nov 29, 2024

mofeing commented Nov 29, 2024

wsmoses commented Nov 29, 2024

Pangoraw commented Nov 29, 2024

wsmoses commented Nov 29, 2024

This comment was marked as resolved.

mofeing left a comment

Choose a reason for hiding this comment

mofeing Nov 30, 2024

Choose a reason for hiding this comment

Pangoraw Dec 1, 2024

Choose a reason for hiding this comment

mofeing Nov 30, 2024

Choose a reason for hiding this comment

mofeing left a comment

Choose a reason for hiding this comment

mofeing commented Nov 27, 2024 •

edited

Loading