MLIR bindings change: improve Attribute Julia types #660

glou-nes · 2025-01-30T11:21:36Z

I close #447, and apply change to Reactant code.

renaming result_ to result
apply the typing defining in args, for instance: result_layouts=nothing::Union{Attribute, Nothing} doesn't constraint result_layouts so change to : result_layouts::Union{Attribute, Nothing} =nothing
create Julia enum for EnumAttr : using Attr parsing (MLIR doesn't expose methods in C to easily construct arbitrary attribute) and discriminant.
typing Attribute and use Julia type for easy ones (Int, Float, Vector).
more complex Attributes ("struct" layout) such as conv

glou-nes · 2025-01-30T11:35:40Z

Attributes are now typed using Julia type system, reducing the usage of direct C API. There are currently limitations from Tablegen files (custom C++ parser/printer for instance). My next idea is to use "Dialect" IRDL dialect to do the code generation directly in Julia.

struct DotAlgorithm
	lhs_precision_type::IR.Type
	rhs_precision_type::IR.Type
	accumulation_type::IR.Type
	lhs_component_count::Int64
	rhs_component_count::Int64
	num_primitive_operations::Int64
	allow_imprecise_accumulation::Bool
end

[...]

function dot_general(lhs::Value, rhs::Value; result::IR.Type, dot_dimension_numbers::Dot, precision_config::Union{Vector{Precision.T}, Nothing}=nothing, algorithm::Union{DotAlgorithm, Nothing}=nothing, location::Location=Location())
    op_ty_results = IR.Type[result, ]
    operands = Value[lhs, rhs, ]
    owned_regions = Region[]
    successors = Block[]
    attributes = NamedAttribute[namedattribute("dot_dimension_numbers", dot_dimension_numbers), ]
    !isnothing(precision_config) && push!(attributes, namedattribute("precision_config", precision_config))
    !isnothing(algorithm) && push!(attributes, namedattribute("algorithm", algorithm))
    
    create_operation(
        "stablehlo.dot_general", location;
        operands, owned_regions, successors, attributes,
        results=op_ty_results,
        result_inference=false
    )
end

mofeing · 2025-01-30T17:23:06Z

src/mlir/IR/Attribute.jl

+struct DenseElements{T}
+    attr::Attribute
+    DenseElements(a::AbstractArray{T}) where {T} = new{T}(IR.DenseElementsAttribute(a))
+end
+Attribute(d::DenseElements) = d.attr
+
+function DenseArrayAttribute(values::Vector{<:Enum})
+    return Attribute([Attribute(value) for value in values])
+end


~~Why not just create an Attribute when you pass a DenseArray to Attribute?~~

What's the difference between DenseElementsAttribute and DenseElements?

I force DenseElements type in MLIR calls such as Convolution:

function convolution( [...] padding::Union{IR.DenseElements{Int64},Nothing}=nothing, [...] location::Location=Location(), )

Because I cannot get Attribute name from the Attribute type.

mofeing · 2025-01-30T17:30:25Z

src/mlir/Dialects.jl

+c(a::AbstractArray) = isempty(a) ? "[]" : a
+c(x) = x


mmm would you mind giving a more descriptive name?

For the moment, I add this function to all struct attributes; I will prefer a one character to not challenging Julia formatter too much. I will add a comment for this.

mofeing

in general i like this PR and anything that makes the MLIR interface more idiomatic but I have some doubts:

I don't think that adding so many type constraints to methods is a good idea (specially adding the ::Any because it's redudant).
i guess that the type contraints are there to avoid C-side crashes by adding a type check (and i'm ok with that), but that can turn impractical already with things like "accept a Vector or a Tuple"
What is the purpose of FlatSymbol and how it's different to FlatSymbolRefAttribute? IMHO MLIR.IR types should have a annotation in the name to remind that it's another type of object (so if it's an attribute, it should have Attribute in the name to remind you it's a MLIR Attribute wrapper).

mofeing · 2025-01-30T17:32:37Z

src/mlir/Dialects/Affine.jl

@@ -353,11 +363,11 @@ func.func @pad_edges(%I : memref<10x10xf32>) -> (memref<12x12xf32) {
 """
 function if_(
    operand_0::Vector{Value};
-    results::Vector{IR.Type},
-    condition,
+    results::Union{Vector{IR.Type},Tuple{Vararg{IR.Type}}},


I think you can use AbstractVecOrTuple, but at this point I think is better to maybe not use type constraints?

I don't know about that one! I don't see why we would want to remove type constraints, these calls can be easily misused.

mofeing · 2025-01-30T17:33:12Z

src/mlir/Dialects/Affine.jl

-    map,
-    location=Location(),
+    result::Union{Nothing,IR.Type}=nothing,
+    map::Any,


I prefer if we put no type constraint here because is more legible and it's completely equivalent to Any.

I use Any as a fallback when an attribute Julia type cannot be created

yeah, but maybe we could just don't put anything? or is it more complicated?

mofeing · 2025-01-30T17:57:41Z

My next idea is to use "Dialect" IRDL dialect to do the code generation directly in Julia.

sounds cool but what do you mean?

glou-nes · 2025-01-31T08:13:10Z

My next idea is to use "Dialect" IRDL dialect to do the code generation directly in Julia.

sounds cool but what do you mean?

Basically, get Dialect definitions as a .mlir. Process this file in Julia with MLIR.jl capabilities and emitting files such as 'Affine.jl' and friends directly there. Several benefits from this: simplify generator code, "easy" runtime dialect bindings generation.

glou-nes · 2025-01-31T08:27:36Z

in general i like this PR and anything that makes the MLIR interface more idiomatic but I have some doubts:

I don't think that adding so many type constraints to methods is a good idea (specially adding the ::Any because it's redudant).
i guess that the type contraints are there to avoid C-side crashes by adding a type check (and i'm ok with that), but that can turn impractical already with things like "accept a Vector or a Tuple"

What is the purpose of FlatSymbol and how it's different to FlatSymbolRefAttribute? IMHO MLIR.IR types should have a annotation in the name to remind that it's another type of object (so if it's an attribute, it should have Attribute in the name to remind you it's a MLIR Attribute wrapper).

Thank for the feedback!
For 1), I will register Attributes that cannot be converted to Julia types ones and remove the useless bound. For 2), Julia Attribute type (and C MLIR Attribute API) doesn't give much information: for instance, I cannot get the mnemonic from it. So It's hard to get useful runtime verification here. So I decide to wrap these Attribute in a Julia Type, and let dispatch deals with the verification. I will fuse these wrapper types with the classic ones when the CI is green!

mofeing · 2025-02-01T10:21:36Z

Basically, get Dialect definitions as a .mlir. Process this file in Julia with MLIR.jl capabilities and emitting files such as 'Affine.jl' and friends directly there. Several benefits from this: simplify generator code, "easy" runtime dialect bindings generation.

do you mean to bypass the limitations of the C++ tablegen @jumerckx was facing? sounds super cool, didn't know this could be done.

does IRDL have all we need for this?

For 1), I will register Attributes that cannot be converted to Julia types ones and remove the useless bound. For 2), Julia Attribute type (and C MLIR Attribute API) doesn't give much information: for instance, I cannot get the mnemonic from it. So It's hard to get useful runtime verification here. So I decide to wrap these Attribute in a Julia Type, and let dispatch deals with the verification.

i like it, but how about this:

change Attribute to be an abstract type so that we can dispatch on all concrete subtypes to the common MlirAttribute methods
make one concrete type per attribute type with the Attribute suffix (i'm open to not add the suffix but i think is good to know that it's an attribute and it's more in concordance with MLIR naming; e.g. https://mlir.llvm.org/doxygen/classmlir_1_1DenseElementsAttr.html)

i believe that the reason why we put only one Attribute type is because MLIR's C-API only has MlirAttribute type.

wdyt?

glou-nes · 2025-02-01T10:38:56Z

Basically, get Dialect definitions as a .mlir. Process this file in Julia with MLIR.jl capabilities and emitting files such as 'Affine.jl' and friends directly there. Several benefits from this: simplify generator code, "easy" runtime dialect bindings generation.

do you mean to bypass the limitations of the C++ tablegen @jumerckx was facing? sounds super cool, didn't know this could be done.

Mathieu Fehr has designed a dialect for this! It's still wip but it can give us much more information than tablegen. And it has really great benefit to inspect dialect for analysis or fuzzing.

For 1), I will register Attributes that cannot be converted to Julia types ones and remove the useless bound. For 2), Julia Attribute type (and C MLIR Attribute API) doesn't give much information: for instance, I cannot get the mnemonic from it. So It's hard to get useful runtime verification here. So I decide to wrap these Attribute in a Julia Type, and let dispatch deals with the verification.

i like it, but how about this:
* change `Attribute` to be an abstract type so that we can dispatch on all concrete subtypes to the common `MlirAttribute` methods

* make one concrete type per attribute type with the `Attribute` suffix (i'm open to not add the suffix but i think is good to know that it's an attribute and it's more in concordance with MLIR naming; e.g. https://mlir.llvm.org/doxygen/classmlir_1_1DenseElementsAttr.html)
i believe that the reason why we put only one Attribute type is because MLIR's C-API only has MlirAttribute type.

wdyt?

Really nice! We discussed this design yesterday with Jules, adding the 'Attribute' suffix seems logical to limit cognitive overhead between Julia and C-API. There is still an open question regarding verification of Attributes that cannot be created from Julia: for now, I will use the raw MlirAttribute type.

mofeing · 2025-02-01T11:06:45Z

Mathieu Fehr has designed a dialect for this! It's still wip but it can give us much more information than tablegen. And it has really great benefit to inspect dialect for analysis or fuzzing.

Cool! do you have a link to it?

There is still an open question regarding verification of Attributes that cannot be created from Julia: for now, I will use the raw MlirAttribute type.

yeah, the MLIR C-API is quite limited and not very actively developed... we discussed this with @makslevental. using MlirAttribute is good for now. i guess we (or someone else) will end up developing our own MLIR C-API, but for now is ok.

glou-nes · 2025-02-01T11:10:26Z

Mathieu Fehr has designed a dialect for this! It's still wip but it can give us much more information than tablegen. And it has really great benefit to inspect dialect for analysis or fuzzing.

Cool! do you have a link to it?

https://mlir.llvm.org/docs/Dialects/IRDL/#irdltype-irdltypeop

There is still an open question regarding verification of Attributes that cannot be created from Julia: for now, I will use the raw MlirAttribute type.

yeah, the MLIR C-API is quite limited and not very actively developed... we discussed this with @makslevental. using MlirAttribute is good for now. i guess we (or someone else) will end up developing our own MLIR C-API.

Right!

mofeing · 2025-02-01T12:24:16Z

ups sorry, i mean how you get a mlir of a dialect. like irdl manages the tablegen or?

glou-nes · 2025-02-02T14:38:37Z

ups sorry, i mean how you get a mlir of a dialect. like irdl manages the tablegen or?

Yes like that! I haven't test the tool, but a wip .td -> .mlir seems to exist.

mofeing · 2025-02-03T07:27:08Z

src/mlir/IR/Attribute.jl

 end

 function DenseElementsAttribute(values::AbstractArray)
    shaped_type = TensorType(size(values), Type(eltype(values)))
-    return Attribute(
+    return DenseElementsAttribute{:unknown}(Attribute(


mmm i think it may be better to put here eltype(values)

Suggested change

return DenseElementsAttribute{:unknown}(Attribute(

return DenseElementsAttribute{eltype(values)}(Attribute(

mofeing · 2025-02-03T07:27:27Z

src/mlir/IR/Attribute.jl

-function Base.reshape(attr::Attribute, shape)
-    @assert isdenseelements(attr) "attribute $(attr) is not a dense elements attribute"
+function Base.reshape(attr::DenseElementsAttribute, shape)
+    attr = attr.attr


Suggested change

attr = attr.attr

attr = Attribute(attr)

mofeing · 2025-02-03T07:27:45Z

src/mlir/IR/Attribute.jl

@@ -746,6 +770,38 @@ function Base.length(attr::Attribute)
    end
 end

+function Base.getindex(attr::DenseElementsAttribute, i)
+    attr = attr.attr


Suggested change

attr = attr.attr

attr = Attribute(attr)

glou-nes · 2025-02-04T12:45:19Z

@mofeing I'm running into segmentation fault in CI. And I cannot reproduce them locally because of cuda. Do you have ideas to debug this? Last working commit is bc846e3.

[1103121] signal (11.1): Segmentation fault
in expression starting at /var/lib/buildkite-agent/builds/gpuci-5/julialang/reactant-dot-jl/test/integration/cuda.jl:22
_ZNK17ReadOnlyKernelArg15matchAndRewriteEN4mlir9enzymexla12KernelCallOpERNS0_15PatternRewriterE at /root/.cache/julia-buildkite-plugin/depots/0190ad31-dfb1-4a4c-ac5b-02d05277ba6c/artifacts/f07b6664aaacf8139f3c3461dd9842dfc2152904/lib/libReactantExtra.so (unknown line)
_ZZN4mlir17PatternApplicator15matchAndRewriteEPNS_9OperationERNS_15PatternRewriterEN4llvm12function_refIFbRKNS_7PatternEEEENS6_IFvS9_EEENS6_IFNS5_13LogicalResultES9_EEEENKUlvE_clEv at /root/.cache/julia-buildkite-plugin/depots/0190ad31-dfb1-4a4c-ac5b-02d05277ba6c/artifacts/f07b6664aaacf8139f3c3461dd9842dfc2152904/lib/libReactantExtra.so (unknown line)
_ZN4mlir17PatternApplicator15matchAndRewriteEPNS_9OperationERNS_15PatternRewriterEN4llvm12function_refIFbRKNS_7PatternEEEENS6_IFvS9_EEENS6_IFNS5_13LogicalResultES9_EEE at /root/.cache/julia-buildkite-plugin/depots/0190ad31-dfb1-4a4c-ac5b-02d05277ba6c/artifacts/f07b6664aaacf8139f3c3461dd9842dfc2152904/lib/libReactantExtra.so (unknown line)

(https://buildkite.com/julialang/reactant-dot-jl/builds/4214#0194c925-1ac0-46a8-a16a-c134be86ad6b)

mofeing · 2025-02-04T13:34:37Z

do you mind finding which is the first test that fails?

also, once you find it, do you mind showing the MLIR if @code_hlo optimize=false f(x...) works? the DenseElements and DenseArray attributes have given us problems before to @avik-pal and me, and probably there is sth wrong in the generated IR

glou-nes · 2025-02-04T16:15:13Z

do you mind finding which is the first test that fails?

It's the first one in cuda.jl!

also, once you find it, do you mind showing the MLIR if @code_hlo optimize=false f(x...) works? the DenseElements and DenseArray attributes have given us problems before to @avik-pal and me, and probably there is sth wrong in the generated IR

Using stacktrace, I get :

@static if !Sys.isapple()
    @testset "Square Kernel" begin
        oA = collect(1:1:64)
        A = Reactant.to_rarray(oA)
        B = Reactant.to_rarray(100 .* oA)
        if CUDA.functional()
            @jit square!(A, B) <- This line
            @test all(Array(A) .≈ (oA .* oA .* 100))
            @test all(Array(B) .≈ (oA .* 100))
        else
            @static if skip_non_cuda_tests
                @test false broken = true
            else
                @code_hlo optimize = :before_kernel square!(A, B)
            end
        end
    end
end

mofeing · 2025-02-04T18:18:59Z

i'm not familiar with the CUDA lowering, sorry
you will need to ask @wsmoses

glou-nes · 2025-02-04T18:28:41Z

Thanks a lot for the feebacks, I think I've found it! CI debugging & code_hlo helps!

codecov · 2025-02-05T13:29:29Z

Codecov Report

Attention: Patch coverage is 17.85714% with 575 lines in your changes missing coverage. Please review.

Project coverage is 33.93%. Comparing base (b6ffc96) to head (8186305).
Report is 388 commits behind head on main.

Files with missing lines	Patch %	Lines
src/mlir/Dialects/Nvvm.jl	0.00%	151 Missing ⚠️
src/mlir/Dialects/VHLO.jl	0.00%	58 Missing ⚠️
src/mlir/Dialects/TPU.jl	0.00%	51 Missing ⚠️
src/mlir/Dialects/StableHLO.jl	53.33%	49 Missing ⚠️
src/mlir/IR/Attribute.jl	35.82%	43 Missing ⚠️
src/mlir/Dialects/Llvm.jl	4.76%	40 Missing ⚠️
src/mlir/Dialects/Triton.jl	0.00%	38 Missing ⚠️
src/mlir/Dialects/CHLO.jl	28.26%	33 Missing ⚠️
src/mlir/Dialects/Gpu.jl	0.00%	33 Missing ⚠️
src/mlir/Dialects/Arith.jl	0.00%	19 Missing ⚠️
... and 10 more

Additional details and impacted files

@@             Coverage Diff             @@
##             main     #660       +/-   ##
===========================================
+ Coverage   21.66%   33.93%   +12.26%     
===========================================
  Files          46       73       +27     
  Lines        8048    14723     +6675     
===========================================
+ Hits         1744     4996     +3252     
- Misses       6304     9727     +3423

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

glou-nes changed the title ~~MLIR bindings change: improve Attribute Types~~ MLIR bindings change: improve Attribute Julia types Jan 30, 2025

mofeing reviewed Jan 30, 2025

View reviewed changes

mofeing reviewed Feb 3, 2025

View reviewed changes

This comment was marked as resolved.

Sign in to view

glou-nes marked this pull request as ready for review February 5, 2025 14:28

glou-nes force-pushed the MLIRgen branch from 8186305 to 894ea96 Compare February 5, 2025 14:37

glou-nes closed this Feb 5, 2025

glou-nes force-pushed the MLIRgen branch from 894ea96 to d15980a Compare February 5, 2025 14:44

glou-nes mentioned this pull request Feb 5, 2025

change to dialect generator #691

Open

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

MLIR bindings change: improve Attribute Julia types #660

MLIR bindings change: improve Attribute Julia types #660

glou-nes commented Jan 30, 2025 •

edited

Loading

glou-nes commented Jan 30, 2025 •

edited

Loading

mofeing Jan 30, 2025 •

edited

Loading

glou-nes Jan 31, 2025 •

edited

Loading

mofeing Jan 30, 2025

glou-nes Jan 31, 2025

mofeing left a comment

mofeing Jan 30, 2025

glou-nes Jan 31, 2025

mofeing Jan 30, 2025

glou-nes Jan 31, 2025

mofeing Jan 31, 2025

mofeing commented Jan 30, 2025

glou-nes commented Jan 31, 2025 •

edited

Loading

glou-nes commented Jan 31, 2025 •

edited

Loading

mofeing commented Feb 1, 2025 •

edited

Loading

glou-nes commented Feb 1, 2025

mofeing commented Feb 1, 2025 •

edited

Loading

glou-nes commented Feb 1, 2025 •

edited

Loading

mofeing commented Feb 1, 2025

glou-nes commented Feb 2, 2025

mofeing Feb 3, 2025

mofeing Feb 3, 2025

mofeing Feb 3, 2025

glou-nes commented Feb 4, 2025

mofeing commented Feb 4, 2025

glou-nes commented Feb 4, 2025

This comment was marked as resolved.

This comment was marked as resolved.

mofeing commented Feb 4, 2025

glou-nes commented Feb 4, 2025

codecov bot commented Feb 5, 2025 •

edited

Loading

	return DenseElementsAttribute{:unknown}(Attribute(
	return DenseElementsAttribute{eltype(values)}(Attribute(

		c(a::AbstractArray) = isempty(a) ? "[]" : a
		c(x) = x

MLIR bindings change: improve Attribute Julia types #660

MLIR bindings change: improve Attribute Julia types #660

Conversation

glou-nes commented Jan 30, 2025 • edited Loading

glou-nes commented Jan 30, 2025 • edited Loading

mofeing Jan 30, 2025 • edited Loading

Choose a reason for hiding this comment

glou-nes Jan 31, 2025 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

mofeing left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

mofeing commented Jan 30, 2025

glou-nes commented Jan 31, 2025 • edited Loading

glou-nes commented Jan 31, 2025 • edited Loading

mofeing commented Feb 1, 2025 • edited Loading

glou-nes commented Feb 1, 2025

mofeing commented Feb 1, 2025 • edited Loading

glou-nes commented Feb 1, 2025 • edited Loading

mofeing commented Feb 1, 2025

glou-nes commented Feb 2, 2025

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

glou-nes commented Feb 4, 2025

mofeing commented Feb 4, 2025

glou-nes commented Feb 4, 2025

This comment was marked as resolved.

This comment was marked as resolved.

mofeing commented Feb 4, 2025

glou-nes commented Feb 4, 2025

codecov bot commented Feb 5, 2025 • edited Loading

Codecov Report

glou-nes commented Jan 30, 2025 •

edited

Loading

glou-nes commented Jan 30, 2025 •

edited

Loading

mofeing Jan 30, 2025 •

edited

Loading

glou-nes Jan 31, 2025 •

edited

Loading

glou-nes commented Jan 31, 2025 •

edited

Loading

glou-nes commented Jan 31, 2025 •

edited

Loading

mofeing commented Feb 1, 2025 •

edited

Loading

mofeing commented Feb 1, 2025 •

edited

Loading

glou-nes commented Feb 1, 2025 •

edited

Loading

codecov bot commented Feb 5, 2025 •

edited

Loading