Migrate reduce dim + More JIT instructions + Major wgpu reduce dim refactor

[louisfd]

• NaiveReduceDim migrated to gpu representation

• SharedReduceDim migrated to gpu representation

• Instruction set needed for SharedReduceDim was added:
  -workgroup size
  -local invocation id
  -num workgroups
  -workgroup barrier
  -loop (not while because we found out it's error-prone, better loop with a break)
  -shared memory

• Major refactor of reduce dim so kernels are generic over the algorithm type (sum, mean, argmin, argmax). Makes it much cleaner and scalable for autotune

Fix #1175

[codecov]

Codecov Report

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

Project coverage is 85.72%. Comparing base (3d93e6d) to head (cae9f04).
Report is 8 commits behind head on main.

Files	Patch %	Lines
...burn-wgpu/src/codegen/dialect/gpu/vectorization.rs	0.00%	25 Missing ⚠️
crates/burn-wgpu/src/codegen/dialect/wgsl/base.rs	25.80%	23 Missing ⚠️
...ates/burn-wgpu/src/codegen/dialect/gpu/variable.rs	4.54%	21 Missing ⚠️
crates/burn-wgpu/src/kernel/reduce/base.rs	57.57%	14 Missing ⚠️
...tes/burn-wgpu/src/codegen/dialect/wgsl/compiler.rs	84.44%	7 Missing ⚠️
crates/burn-wgpu/src/codegen/dialect/gpu/shader.rs	0.00%	3 Missing ⚠️
crates/burn-wgpu/src/fusion/tracing/builder.rs	0.00%	3 Missing ⚠️
crates/burn-wgpu/src/codegen/dialect/gpu/branch.rs	88.88%	1 Missing ⚠️
...rn-wgpu/src/codegen/dialect/gpu/synchronization.rs	0.00%	1 Missing ⚠️
crates/burn-wgpu/src/compute/tune_key.rs	0.00%	1 Missing ⚠️
... and 1 more

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1396      +/-   ##
==========================================
+ Coverage   85.52%   85.72%   +0.20%     
==========================================
  Files         594      608      +14     
  Lines       68226    69745    +1519     
==========================================
+ Hits        58348    59792    +1444     
- Misses       9878     9953      +75     

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

[nathanielsimard]

LGTM

[nathanielsimard]

Is it a while loop or just a loop?

[louisfd]

just a loop, forgot to change the comment