Add Ops.hlo_call(::String, args...)

[Pangoraw]

No description provided.

[Pangoraw]

idea: we could hash the code and just the call the function if a function with the same hash is already present in the module (preventing duplicate functions).

So that:

for i in 1:N
  x = Ops.hlo_call(code, x)
end

has the function just once (but multiple calls).

---

Edit: implemented this by using the code hash in the function names.

[mofeing]

that can be done directly with a Dict{HLOIR,Function} (or sth like that) and then calling

get!(dict, code) do code
    compile(code)
end

[Pangoraw]

that can be done directly with a Dict{HLOIR,Function} (or sth like that) and then calling

ah right, do we have a place to store such state within a single trace ? Otherwise my idea is to name these function like $(orig_name)_$(hash(code)) and see if it exists within the module before adding it.

[mofeing]

I have a question... What is hlo_call returning? A function? Wouldn't it be better to just inject the MLIR code in the current MLIR block?

[Pangoraw]

No, it is returning a tuple of tracedarray returned by the call to the function.
i.e. Reactant.@jit Ops.hlo_call(code, x, y) is a tuple of concrete arrays.

[mofeing]

what if the first op is not main? like what if we the code was traced by us and we added some function barriers.

maybe we can add a kwarg for selecting the target function (and default it to main), so we just iterate over the ops doing first(Iterators.filter(op -> String(IR.attr(op, "sym_name")) == target_fn, OperationIterator(body))

[Pangoraw]

This is inside the code that was given by the caller. Currently, the expectation is that there is only one function inside the given module. We can surely revisit that with a keyword indeed, or a tuple as for Core.llvmcall

[wsmoses]

yeah I would instead ideally have a kwargument fn=main, and we extract that fn as the top level one

[Pangoraw]

added a func_name::String kwarg for this

[wsmoses]

we also need to mark all fn's as private, as well as make sure to move all fns in the module (e.g. the main function could call something)

[wsmoses]

we have some utilities here: https://github.com/EnzymeAD/Enzyme-JAX/blob/f6587e37ff7298f2a1a273b08c24d69fca7ff30f/src/enzyme_ad/jax/compile_with_xla.cc#L190 and https://github.com/EnzymeAD/Enzyme-JAX/blob/f6587e37ff7298f2a1a273b08c24d69fca7ff30f/src/enzyme_ad/jax/primitives.py#L811 in Enzyme-JaX for explicitly making we can do all the things

[Pangoraw]

Done, do you know if we can encounter ops other than func.func (maybe gpu.func in the future?) and what to do with them ?

[wsmoses]

I think it’s fine to assume func for now but if desired we could generalize to function interface or whatnot

[wsmoses]

Oh we should import all the ops, not just func, but it’s okay to be limited to just func as the entry function. Eg if main calls a gpu function that would be fine. Or if a global constant op

[wsmoses]

Can you add a test with multiple functions in the module.

and can you also add a test with two (different) hlo calls that happen to contain functions of the same name (to make sure we do the symbol rename properly)

[Pangoraw]

So hlo_call currently returns a tuple of arrays (one for each result), should we special case nresults() == 1 to return the only array (no tuple) ?

[wsmoses]

Can you also add a version of this where the two definitions are different.

just because if we fix caching then we might not actually not emit it twice (and thus not check things)

[Pangoraw]

Added a test with the same name but different definitions:

Reactant.jl/test/ops.jl

Lines 945 to 970 in 8eb71cc

	function f_multiple_hlo_calls(x, y)
	x, = Ops.hlo_call(
	"""
	module {
	func.func @main(%arg0: tensor<3xf32>, %arg1: tensor<3xf32>) -> tensor<3xf32> {
	%0 = stablehlo.add %arg0, %arg1 : tensor<3xf32>
	return %0 : tensor<3xf32>
	}
	}
	""",
	x,
	y,
	)
	return Ops.hlo_call(
	"""
	module {
	func.func @main(%arg0: tensor<3xf32>, %arg1: tensor<3xf32>) -> tensor<3xf32> {
	%0 = stablehlo.multiply %arg0, %arg1 : tensor<3xf32>
	return %0 : tensor<3xf32>
	}
	}
	""",
	x,
	y,
	)
	end

[wsmoses]

So hlo_call currently returns a tuple of arrays (one for each result), should we special case nresults() == 1 to return the only array (no tuple) ?

Honestly I think it’s better to always return the tuple. That way folks using it don’t need to special case if there are multiple returns or not

[wsmoses]

Awesome stuff!

[mofeing]

this is amazing and opens a way to do things like #354

[Pangoraw]

Other potential applications:

• Pytorch has pretty good stablehlo export with torch_xla nowadays.
• ONNX via onnx-mlir

Strings really are the universal model format (we can add hlo_call(::Vector{UInt8}, ...) which takes mlir bytecode).