Internal Compiler Error: Type Error and Unknown Variable Usage

[ka-weihe]

This code:

import "lib/github.com/diku-dk/linalg/linalg"
import "lib/github.com/diku-dk/linalg/lu"

module linalg = mk_linalg f64
module lu = mk_lu f64

def grad [n] (f: [n]f64 -> f64) (x: [n]f64) =
    map (\i -> 
        let X = linalg.veczeros n
        let X[i] = 1.0
        in jvp f x X
    ) (iota n)

def hess [n] (f: [n]f64 -> f64) (x: [n]f64) =
    map (\i ->
        map (\j ->
            let X = linalg.veczeros n
            let X[i] = 1.0
            let Y = linalg.veczeros n
            let Y[j] = 1.0
            in jvp (\w -> jvp f w X) x Y
        ) (iota n)
    ) (iota n)

def newton_equality [n][m] (f: [n]f64 -> f64) 
                           (A: [m][n]f64) 
                           (b: [m]f64) 
                           (x0: [n]f64) 
                           (epsilon: f64) 
                           : [n]f64 =
    let f_grad = grad f
    let f_hess = hess f
    let x = loop x = x0 for i < 10 do
        let grad_val = f_grad x
        let hessian_val = f_hess x
        let KKT = linalg.matzeros (n + m) (n + m)
        let rhs = linalg.veczeros (n + m)
        let KKT[:n, :n] = hessian_val
        let KKT[:n, n:] = transpose A
        let KKT[n:, :n] = A
        let rhs[:n] = map f64.neg grad_val
        let rhs[n:] = map2 (-) b (linalg.matvecmul_row A x)
        let delta_x_lam = lu.ols 4 KKT rhs
        let delta_x = delta_x_lam[:n]
        let t = 1.0
        let x_new = map2 (+) x (map (*t) delta_x)
        in if linalg.vecnorm delta_x < epsilon then
            x
        else
            x_new
    in x

    let main =
        let A = [[1.0, 1.0, 1.0], [1.0, 2.0, 3.0]]
        let b = [6.0, 14.0]
        let x0 = [1.0, 1.0, 1.0]
        let epsilon = 0.0001
        let mu = [0.05, 0.1, 0.15]
        let esg = [0.5, 0.6, 0.7]
        let h (exp: f64) (x: [3]f64): [5]f64 =
            let a = linalg.dotprod x mu
            let b = linalg.dotprod x esg
            in ([       
                exp - a,
                0.8 - b             
            ] ++ map (f64.neg <-< (+1e-6)) x) :> [5]f64
        let phi (exp: f64) (x: [3]f64): f64 = 
            f64.neg (h exp x |> map (f64.log <-< f64.neg) |> reduce (+) 0.0)
        let x = map (\e -> newton_equality (phi e) A b x0 epsilon) [0.1, 0.2, 0.3]
        in x

Produces this internal compiler error (with futhark cuda)

Internal compiler error.  Please report this:
 https://github.com/diku-dk/futhark/issues
Type error after pass 'tile loops':
In expression of statement
 {defunc_2_map_res_30456 : ({}, [3i64][3i64]f64)}
in body of last case
In expression of statement
 {defunc_2_map_res_47235 : ({}, [3i64][3i64]f64)}
Inside the loop body
In expression of statement
 {x_47240 : ({}, [3i64][3i64]f64)}
in body of case {true}
In expression of statement
 {x_47245 : ({}, [3i64][3i64]f64)}
In expression of statement
 {tiled_inside_loop_57581 : ({}, [tile_size_52178][5i64]f64)}
Inside the loop body
In expression of statement
 {accs_55773 : ({}, [tile_size_52178]f64)}
Inside the loop body
In expression of statement
 {full_tile_54587 : ({}, [tile_size_52178]f64)}
In expression of statement
 {dotprod_arg_54599 : ({defunc_2_map_res_47795}, [i_47826]f64)}
Use of unknown variable defunc_2_map_res_47795.

[athas]

Condensed:

def dotprod [n] (a: [n]f64) (b: [n]f64): f64 =
  map2 (*) a b |> reduce (+) 0

def newton_equality [n][m] (_: [m][n]f64) (x: [n]f64) hessian_val =
    let KKT = replicate (n + m) (replicate (n + m) 0)
    let KKT[:n, :n] = hessian_val
    in loop y=replicate n 0 for i in 0..<n do
       let sum = dotprod KKT[i,:i] y[:i]
       let y[i] = copy(0 - sum) / KKT[i,i]
       in y

entry main A x0s hessian_vals =
  map2 (newton_equality A) x0s hessian_vals

[athas]

Looks like this was caused by the fix to #1933.