@@ -837,6 +837,53 @@ function stream2token(source::Value; result::IR.Type, location=Location())
837837 )
838838end
839839
840+ function triton_call (
841+ gridx:: Value ,
842+ gridy:: Value ,
843+ gridz:: Value ,
844+ shmem:: Value ,
845+ inputs:: Vector{Value} ;
846+ result_0:: Vector{IR.Type} ,
847+ fn,
848+ backend_config= nothing ,
849+ operand_layouts= nothing ,
850+ result_layouts= nothing ,
851+ arg_attrs= nothing ,
852+ res_attrs= nothing ,
853+ output_operand_aliases= nothing ,
854+ xla_side_effect_free= nothing ,
855+ location= Location (),
856+ )
857+ op_ty_results = IR. Type[result_0... ,]
858+ operands = Value[gridx, gridy, gridz, shmem, inputs... ]
859+ owned_regions = Region[]
860+ successors = Block[]
861+ attributes = NamedAttribute[namedattribute (" fn" , fn),]
862+ ! isnothing (backend_config) &&
863+ push! (attributes, namedattribute (" backend_config" , backend_config))
864+ ! isnothing (operand_layouts) &&
865+ push! (attributes, namedattribute (" operand_layouts" , operand_layouts))
866+ ! isnothing (result_layouts) &&
867+ push! (attributes, namedattribute (" result_layouts" , result_layouts))
868+ ! isnothing (arg_attrs) && push! (attributes, namedattribute (" arg_attrs" , arg_attrs))
869+ ! isnothing (res_attrs) && push! (attributes, namedattribute (" res_attrs" , res_attrs))
870+ ! isnothing (output_operand_aliases) &&
871+ push! (attributes, namedattribute (" output_operand_aliases" , output_operand_aliases))
872+ ! isnothing (xla_side_effect_free) &&
873+ push! (attributes, namedattribute (" xla_side_effect_free" , xla_side_effect_free))
874+
875+ return create_operation (
876+ " enzymexla.triton_call" ,
877+ location;
878+ operands,
879+ owned_regions,
880+ successors,
881+ attributes,
882+ results= op_ty_results,
883+ result_inference= false ,
884+ )
885+ end
886+
840887function wrap (
841888 operand:: Value ;
842889 result= nothing :: Union{Nothing,IR.Type} ,
0 commit comments