JuliaLabs / MLIR.jl

MIT License
56 stars 9 forks source link

MLIRValueTrait to support user defined Value types #64

Open jumerckx opened 7 months ago

jumerckx commented 7 months ago

What this enables:

# Setup #
using MLIR
using MLIR: IR, API
using MLIR.IR: Value, Operation, Block, push_argument!, Context, MLIRValueTrait, Convertible
using MLIR.Dialects: arith
ctx = Context()

# Define a Julia type and mapping to MLIR #
struct MLIRInteger{N} <: Integer
    value::Value
    MLIRInteger{N}(i::Value) where {N} = new(i)
end
const i64 = MLIRInteger{64}
IR.Type(::Type{MLIRInteger{N}}) where {N} = API.mlirIntegerTypeGet(ctx, N)
MLIRValueTrait(::Type{<:MLIRInteger}) = Convertible()

# Specialize or write function #
function Base.:+(a::T, b::T)::T where {T<:MLIRInteger}
    T(IR.result(arith.addi(a, b)))
end

# Execute # 
block = Block()
a = i64(push_argument!(block, IR.Type(i64)))
b = i64(push_argument!(block, IR.Type(i64)))

a+b

Of course, using this approach you lose direct access to a handle to the generated operation (arith.addi) so it is of dubious value as of now. In Jojo.jl this system is used in combination with CassetteOverlay to make sure that all created operations are collected. This system is still a bit finicky and I'm still working on it so I wouldn't try upstream this yet.

While collecting things to push here, it's clear some cleaning up is still required.

I'm not even sure if it's necessary to have this trait system. For now, the conversion to IR.Value is so simple that you might as well just specialize get_value instead of letting it happen through the trait. On the other hand, this might help in the future as people want to add extra functionality to Value-like types.

jumerckx commented 7 months ago

Thanks for the changes @mofeing! What do you think of using Base.convert(Value, ...) instead of value? Also, if we don't go with the convert, wouldn't Value (capitalized) make more sense since we're returning an object of that type?

mofeing commented 7 months ago

Thanks for the changes @mofeing! What do you think of using Base.convert(Value, ...) instead of value? Also, if we don't go with the convert, wouldn't Value (capitalized) make more sense since we're returning an object of that type?

I like to use constructor methods (i.e. Value) whenever the passing argument has a field from which it can be easily extracted.

The problem with Base.convert is that it has some clear semantics and can be implicitly called on some cases: https://docs.julialang.org/en/v1/manual/conversion-and-promotion/#When-is-convert-called?

I'm not sure in which cases you will need to convert from some other type to Value, but check out the link I posted above and decide based on that.

EDIT: Can you give me some example where you need to convert some other type to Value please?

jumerckx commented 7 months ago

I like to use constructor methods (i.e. Value) whenever the passing argument has a field from which it can be easily extracted.

Isn't this the case here then? The standard assumption is that these value-like types have a field holding an actual IR.Value. And if wouldn't hold for a type, the user would have to implement the Value-getter (whatever it's called) themselves.

EDIT: Can you give me some example where you need to convert some other type to Value please?

Actually, the one example I had one was to get rid of explicit conversions when pushing to the operands array in generated dialect wrappers. e.g.:

function addf(lhs, rhs; result=nothing::Union{Nothing, IR.Type}, fastmath=nothing, location=Location())
    results = IR.Type[]
    operands = IR.Value[lhs, rhs] # instead of IR.Value[value(lhs), value(rhs)]

But upon thinking more about this, that really isn't a good use-case as this code is auto-generated in the first place and the implicit conversion might cause more confusion elsewhere.

Conversion in the other direction could be more interesting, though. For example, an explicit return type for a function would convert whatever result that that function returns automatically. But again, this might cause more confusion, so maybe left best for later down the line in a different pr.

Do we merge this pr once the naming has decided or are there other considerations?

mofeing commented 7 months ago

Do we merge this pr once the naming has decided or are there other considerations?

I'm ok with merging as it is but you just need to update the other value method definitions, like https://github.com/JuliaLabs/MLIR.jl/blob/d438d41f529ba32b926c55d2bf59acff1da3b857/src/IR/AffineExpr.jl#L124

codecov-commenter commented 7 months ago

Codecov Report

Attention: Patch coverage is 0% with 6 lines in your changes are missing coverage. Please review.

Project coverage is 2.28%. Comparing base (f561fa2) to head (1638254). Report is 3 commits behind head on main.

Files Patch % Lines
src/IR/Value.jl 0.00% 5 Missing :warning:
src/IR/AffineExpr.jl 0.00% 1 Missing :warning:

:exclamation: Your organization needs to install the Codecov GitHub app to enable full functionality.

Additional details and impacted files ```diff @@ Coverage Diff @@ ## main #64 +/- ## ======================================== - Coverage 2.28% 2.28% -0.01% ======================================== Files 125 125 Lines 28527 28532 +5 ======================================== Hits 653 653 - Misses 27874 27879 +5 ```

:umbrella: View full report in Codecov by Sentry.
:loudspeaker: Have feedback on the report? Share it here.

mofeing commented 5 months ago

I've updated MLIR_jll to also accept v17, so the CI failing on Julia nightly should no longer fail.

Anyway, you're free to merge whenever you want to.

jumerckx commented 5 months ago

I've updated MLIR_jll to also accept v17, so the CI failing on Julia nightly should no longer fail.

Thanks for the heads up, and for your CI work! The error seems to be an Pkg environment conflict so not sure what could be causing this as project.toml is untouched?

I kind of lost the plot with this PR. I'm currently prepping a branch with my work on generating MLIR from Julia code that will supersede this one so this shouldn't be merged anymore.