Speed up Value function Iterations in Julia - loops

Question

I am coding a simple value function iteration in Julia. I posted the code below.
# ................................
# Solving the model
function Solver(P,uf,hf)
#unpack r, σrisk, ρ, η, β, θ, nx, m, μ, fhat, ub, lb, tol, maxite, ne = P;
# ..................................................
# Tauchen discretization
ly, Π = mytauch(μ, ρ, η, nx, m); # it is a tauchen discretization
# ..................................................
# states
y = exp.(ly);
b = [LinRange(lb,ub, ne)...];
yb = b .+ y';
p0 = findfirst(x-> x ==0, b);
uhy = uf.(hf(y, fhat));
# ..................................................
# Initial guess for value functions
vr0 = 1 / (1 - β) * uf.((r / (1 + r)) * b .+ y');
vd0 = (1 / (1 - β) * uhy)';
myup(vr,vd) = updateVF(vr, vd, b, yb, uhy, Π, p0, uf, P);
# ..................................................
# Convergence algorithm
dif = 1
ite = 1
vr1, vd1, bopt = (nothing,nothing, nothing);
while dif>tol && ite < maxite
vr1, bopt, vd1 = myup(vr0,vd0)
dif = max(maximum(abs.(vr1-vr0)), maximum(abs.(vd1-vd0)));
ite+=1;
vr0, vd0 = (vr1, vd1);
end
# ..................................................
# Policies functions
vf1, D, δ, q = postVR(vr1,vd1,Π,r)
#bb = repeat(b,1,nx);
vd1 = repeat(vd1,ne,1);
bprime = b[bopt];
PolFun = (vf=vf1 , vr=vr1, vd=vd1, bprime = bprime, q=q,D=D, δ=δ);
out2 = (b= b, y= y,ydef=hf(y, fhat),Π=Π, uhy = uhy, p0 = p0);
return PolFun, out2;
end
function updateVF(vr0, vd0,b,yb,uhy,Π, p0, uf, P)
#unpack r , β, θ, nx,ne = P;
# ..................................................
# Preliminaries
vf0, D0, δ0, q0 = postVR(vr0,vd0,Π,r)
evf0 = vf0*Π';
evd0 = vd0*Π';
qb = q0 .* b;
βevf0 = β*evf0;
# ..................................................
# Value and debt policy function update
# Value of Default
vd1 = uhy' + β*(θ*evf0[p0,:]'+(1-θ)*evd0);
# Value of Repayment
# Method 2: by column
vr1 = Array{Float32,2}(undef, ne, nx);
bopt = Array{CartesianIndex{2},2}(undef, ne, nx);
upvr1!(vr1,bopt,yb,qb,βevf0,nx)
return vr1, map(ii -> ii[1],bopt), vd1
end
function postVR(vr,vd,Π,r)
vf = max.(vr,vd);
D = 1*(vd .> vr);
δ = D*Π';
q = (1/(1+r)) * (1 .- δ);
return vf, D, δ, q;
end
function upvr1!(vr1,bopt,yb,qb,βevf0,nx)
#fastmath #inbounds Threads.#threads for i in 1:nx
vr1[:,i], bopt[:,i] = #views findmax(uf.(max.(0,yb[:,i] .- qb[:,i]')) .+ βevf0[:,i], dims=1);
end
end
The master code is:
P = ( r = 0.017, σrisk = 2.0, ρ = 0.945, η = 0.025, β = 0.953,
θ = 0.282,nx = 21, m = 3, μ = 0.0,fhat = 0.969,
ub = 0, lb = -0.4, tol = 1e-6, maxite = 2000, ne = 251);
uf(x) = -1/x;
hf(y::Array, fhat::Float64) = min.(y, fhat * mean(y));
#time polfun, o2 = Solver(P,uf,hf);
The running tipe after first compilation is ~1 second. How can I speed-up this code? I am asking it because in previous versions of Matlab, in my experience, Julia was 10x faster but with the last version (Matlab 2020, that I tried last week) its performance was close to Julia but as I increase the number of grid-points (ne > 1000) Matlab reported faster. So I suppose that I am not taking advantage of all functionalities of Julia.
Note: The loop for upvr1! overperforms that a broadcasting version
yb = reshape(yb,(ne,nx,1));
βevf01 = PermutedDimsArray(reshape(βevf0,(ne,nx,1)),[3,2,1]);
qb1 = PermutedDimsArray(reshape(qb,(ne,nx,1)),[3,2,1]);
vr1, bopt = #views findmax(uf.(yb .- qb1) .+ βevf01, dims=3)
return vr1[:,:,1], map(ii -> ii[3], bopt)[:,:,1], vd1