	ifargs.trace_size>0:
	trace.set_trace_size(args.trace_size)
	trace.start_trace()

	# JIT-compile the kernel then launches the kernel with the given arguments
	vector_scalar_mul(input0,scalar,output)

	# Stop tracing and save results if tracing was enabled
	ifargs.trace_size>0:
	trace_filename=f"trace_output_{args.num_elements}_{args.data_type}.json"
	trace.stop_trace(trace_filename)
	print(f"Tracing completed and saved to{trace_filename}")

And this:

mlir-aie/programming_examples/basic/vector_scalar_mul/vector_scalar_mul_jit.py

Lines 115 to 126 in612af49

	worker=Worker(
	core_body,
	fn_args=[of_in.cons(),of_factor.cons(),of_out.prod(),scale],
	trace=1iftrace.get_trace_size()>0else0,
	)

	# Runtime operations to move data to/from the AIE-array
	rt=Runtime()

	withrt.sequence(tensor_ty,scalar_ty,tensor_ty)as (A,F,C):
	iftrace.get_trace_size()>0:
	rt.enable_trace(trace.get_trace_size())

I would have liked to hide this away too but it's not straight forward at the moment because of how the default workflow works (i.e., I can't simply callenable_trace for the runtime sequence if tracing is enabled in the backend). There are also other complications for tracing that Jack has brought up.

Output trace visualization:

Add trace with JIT example

612af49

Copy link

CollaboratorAuthor

mawad-amd commentedAug 27, 2025

@jackl-xilinx @AndraBisca @jgmelber This is an initial draft for trace support in JIT.

mawad-amd mentioned this pull request

Aug 28, 2025

Experiment with trace#2251

Closed

mawad-amdand others added4 commits

August 28, 2025 12:22

Merge branch 'main' into muhaawad/trace-jit

62e3409

Combine set trace and start tracing

5921a2f

Autoset the tracing for worker and runtime

e14428d

Merge branch 'main' into muhaawad/trace-jit

74a6b88

mawad-amd mentioned this pull request

Oct 2, 2025

Improvements to JIT workflow#2206

Open

6 tasks

hunhoffe added this to theJIT/Runtime/Compilation Refactor milestone

Nov 10, 2025

hunhoffe mentioned this pull request

Nov 19, 2025

Refine JIT Runtime/Compilation/Build Capabilities#2721

Open

21 tasks

hunhoffe removed this from theJIT/Runtime/Compilation Refactor milestone

Nov 19, 2025

Reviewers

AndraBiscaAwaiting requested review from AndraBiscaAndraBisca will be requested when the pull request is marked ready for reviewAndraBisca is a code owner

jackl-xilinxAwaiting requested review from jackl-xilinxjackl-xilinx will be requested when the pull request is marked ready for reviewjackl-xilinx is a code owner

andrejAwaiting requested review from andrejandrej will be requested when the pull request is marked ready for reviewandrej is a code owner

pvasireddy-amdAwaiting requested review from pvasireddy-amdpvasireddy-amd will be requested when the pull request is marked ready for reviewpvasireddy-amd is a code owner

jgmelberAwaiting requested review from jgmelberjgmelber will be requested when the pull request is marked ready for reviewjgmelber is a code owner

hunhoffeAwaiting requested review from hunhoffehunhoffe will be requested when the pull request is marked ready for reviewhunhoffe is a code owner

denolfAwaiting requested review from denolfdenolf will be requested when the pull request is marked ready for reviewdenolf is a code owner

fifieldAwaiting requested review from fifieldfifield will be requested when the pull request is marked ready for reviewfifield is a code owner

Copilot code review CopilotAwaiting requested review from CopilotCopilot will automatically review once the pull request is marked ready for review

Labels

None yet

Movatterモバイル変換

Navigation Menu

Search code, repositories, users, issues, pull requests...

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Add trace with JIT example#2541

Are you sure you want to change the base?

Add trace with JIT example#2541

Uh oh!

Conversation

mawad-amd commentedAug 27, 2025•
edited
Loading

Uh oh!