How to compile a diagonal array efficiently?

Question

For example, if we use some functions defined prior to Compile, we usually have the main evaluators in the compiled codes, pointing to the definition of the function.

f[t_] := If[t <= 1., Cos[t]*Sin[t], 0.]

CompilePrint@
 Compile[{{t, _Real}}, IdentityMatrix[2] -  f[t], 
  CompilationOptions -> {"InlineCompiledFunctions" -> True, 
    "InlineExternalDefinitions" -> True}]

      1 argument
      1 Integer register
      3 Real registers
      2 Tensor registers
      Underflow checking off
      Overflow checking off
      Integer overflow checking on
      RuntimeAttributes -> {}

      R0 = A1
      I0 = 2
      Result = T(R2)1

1 T(I2)0 = MainEvaluate[ Hold[IdentityMatrix][ I0]]
2 R1 = MainEvaluate[ Hold[f][ R0]]
3 R2 = - R1
4 T(R2)1 = R2 + T(I2)0
5 Return

We can avoid main evaluator by using Evaluate in the function, which specifically substidue the definition of the function. However, sometimes this introduces repeated code in the compiled results. For instance, in the following example, Evaluate simplify expands a number into a matrix and repeatedly calculated this expression for two times. We can see that the 26-50 lines of the compiled code are essentially the same as 1-25 lines.

CompilePrint@
 Compile[{{t, _Real}}, Evaluate[IdentityMatrix[2] - f[t + 1./2.]], 
  CompilationOptions -> {"InlineCompiledFunctions" -> True, 
    "InlineExternalDefinitions" -> True}]

      1 argument
      1 Boolean register
      1 Integer register
      12 Real registers
      3 Tensor registers
      Underflow checking off
      Overflow checking off
      Integer overflow checking on
      RuntimeAttributes -> {}

      R0 = A1
      I0 = 1
      R3 = 1.
      R4 = 7.
      R1 = 0.5
      R7 = 0.
      Result = T(R2)2

1 R2 = R1 + R0
2 B0 = R2 <= R3 (tol R4)
3 if[ !B0] goto 10
4 R2 = R1 + R0
5 R5 = Cos[ R2]
6 R6 = Sin[ R2]
7 R5 = R5 * R6
8 R6 = R5
9 goto 11
10    R6 = R7
11    R5 = - R6
12    R6 = I0
13    R6 = R6 + R5
14    R5 = R1 + R0
15    B0 = R5 <= R3 (tol R4)
16    if[ !B0] goto 23
17    R5 = R1 + R0
18    R8 = Cos[ R5]
19    R9 = Sin[ R5]
20    R8 = R8 * R9
21    R9 = R8
22    goto 24
23    R9 = R7
24    R8 = - R9
25    T(R1)0 ={ R6, R8 }
26    R6 = R1 + R0
27    B0 = R6 <= R3 (tol R4)
28    if[ !B0] goto 35
29    R6 = R1 + R0
30    R8 = Cos[ R6]
31    R9 = Sin[ R6]
32    R8 = R8 * R9
33    R9 = R8
34    goto 36
35    R9 = R7
36    R8 = - R9
37    R9 = R1 + R0
38    B0 = R9 <= R3 (tol R4)
39    if[ !B0] goto 46
40    R9 = R1 + R0
41    R10 = Cos[ R9]
42    R11 = Sin[ R9]
43    R10 = R10 * R11
44    R11 = R10
45    goto 47
46    R11 = R7
47    R10 = - R11
48    R11 = I0
49    R11 = R11 + R10
50    T(R1)1 ={ R8, R11 }
51    T(R2)2 ={ T(R1)0, T(R1)1 }
52    Return

So is there a way to fix this repeating?

Note that change only the argument of the external function, the behavior changes, why?

CompilePrint@
 Compile[{{t, _Real}}, Evaluate[IdentityMatrix[2] - f[t]], 
  CompilationOptions -> {"InlineCompiledFunctions" -> True, 
    "InlineExternalDefinitions" -> True}]

      1 argument
      1 Boolean register
      1 Integer register
      6 Real registers
      3 Tensor registers
      Underflow checking off
      Overflow checking off
      Integer overflow checking on
      RuntimeAttributes -> {}

      R0 = A1
      I0 = 1
      R1 = 1.
      R2 = 7.
      R5 = 0.
      Result = T(R2)2

1 B0 = R0 <= R1 (tol R2)
2 if[ !B0] goto 8
3 R3 = Cos[ R0]
4 R4 = Sin[ R0]
5 R3 = R3 * R4
6 R4 = R3
7 goto 9
8 R4 = R5
9 R3 = - R4
10    R4 = I0
11    R4 = R4 + R3
12    T(R1)0 ={ R4, R3 }
13    T(R1)1 ={ R3, R4 }
14    T(R2)2 ={ T(R1)0, T(R1)1 }
15    Return

Answer 1

For nicer inlining techniques, see my answer here

In what IdentityMatrix[2] - f[t] evaluates to, there is pretty much four times the same code.

Clear[t];
IdentityMatrix[2] - f[t]

{{1 - If[t <= 1., Cos[t] Sin[t], 0.], -If[t <= 1., Cos[t] Sin[t], 0.]}, {-If[t <= 1., Cos[t] Sin[t], 0.], 1 - If[t <= 1., Cos[t] Sin[t], 0.]}}

I will inline f using DownValues. Throughout this answer, you should read a construction like g@@(Hold[...f[t]...]/.DownValues[f]) as g[...f[t]...], only realising that now f has been inlined. Sadly the syntax highlighter now makes the colours of the variables a bit confusing, but what can you do.

Note that a call to MainEvaluate to compute a big IdentityMatrix is not bad. It is also not so bad to calculate f[t] using MainEvaluate only once if the matrix is very large. But I am not sure if you have large matrices in mind, or if you are interested in generating small matrices quickly (a lot of times?).

Big matrices

The following code makes one call to MainEvaluate to get the IdentitiyMatrix. It only calculates the Sin and the Cos once.

f[t_] := If[t <= 1., Cos[t]*Sin[t], 0.]

cfu =
 Function[Null, Compile[{{t, _Real}}, #], HoldAll] @@
  (
   Hold[
     IdentityMatrix[2] - f[t - 0.5]
     ] /. DownValues[f]
   )

You could also write this like this if you prefer

specialReleaseHold[expr_] := Delete[expr, {0, 0}]

cfu2 =
 specialReleaseHold@
  Hold[Compile][
   Unevaluated[{{t, _Real}}]
   ,
   Unevaluated @@
    (
     Hold[
       IdentityMatrix[2] - f[t - 0.5]
       ] /. DownValues[f]
     )
   ]

And we have

CompilePrint@cfu == CompilePrint@cfu2

True

ReleaseHold would actually do the same thing as specialReleaseHold here (in fact, everywhere in this answer), but in similar cases using ReleaseHold could be bad.

Another alternative (focussed on large matrices), using ConstantArray

cfu3 =
 specialReleaseHold@
  Hold[Compile][
   Unevaluated[{{t, _Real}}],

   Unevaluated @@
     Hold[
      Block[
       {res, i},
       res =
        ConstantArray[
         -f[t - 0.5], {2, 2}
         ];
       For[i = 1, i <= 2, i++,
        res[[i, i]] += 1

        ];
       res
       ]
      ] /. DownValues[f]
   ]

Small matrices

For small matrices we want to avoid MainEvaluate altogether. We could do

cfu4 =
  specialReleaseHold@
   Hold[Compile][
    Unevaluated[{{t, _Real}}],
    Unevaluated @@
     (
      Hold[
        Block[{res, term},
         term = -f[t - 0.5];
         res = {{1., 0.}, {0., 1.}};
         res + term
         ]
        ] /. DownValues[f]
      )

    ];

Which somehow turns out to be faster than

cfu5 =
  specialReleaseHold@
   Hold[Compile][
    Unevaluated[{{t, _Real}}],
    Unevaluated @@
     (
      Hold[
        Block[{term},
         term = -f[t - 0.5];
         {{1. + term, 0. + term}, {0. + term, 1. + term}}
         ]
        ] /. DownValues[f]
      )
    ];

As we will see further below, neither of these functions use MainEvaluate.

Uncompiled (small matrices)

We will see indeed that compiling pays off. For reference, we define the following function

fu =
  ReleaseHold@
   Hold[Function][
    Hold[t],
    Hold[
      Block[{term},
       term = -f[t - 0.5];
       {{1. + term, 0. + term}, {0. + term, 1. + term}}
       ]
      ] /. DownValues[f]
    ];

Comparison

cfu[2.] == cfu2[2.] == cfu3[2.] == cfu4[2.] == cfu5[2.] == fu[2.]==
 IdentityMatrix[2]

True

cfu[0.1] == cfu2[0.1] == cfu3[0.1] == cfu4[0.1] == cfu5[0.1] == fu[0.1]

True

StringFreeQ[CompilePrint@#, "MainEvaluate"] & /@
 {cfu, cfu2, cfu3, 
  cfu4, cfu5}

{False, False, False, True, True}

Function[Do[#[0.1], {10000}] // Timing // First] /@ {cfu, cfu2, cfu3, 
  cfu4, cfu5, fu}

{0.023010, 0.020006, 0.024175, 0.009206, 0.011438, 0.077503}

Answer 2

This answer deals only with f, not IdentityMatrix (which is sadly not compilable).

Compile is much better at extracting expressions form pure functions. If you are going to use InlineExternalDefinitions, consider making the external functions pure Functions.

In[10]:= Clear[f]

f = Function[t, If[t <= 1., Cos[t]*Sin[t], 0.]];

CompilePrint@
 Compile[{{t, _Real}}, IdentityMatrix[2] - f[t], 
  CompilationOptions -> {"InlineCompiledFunctions" -> True, 
    "InlineExternalDefinitions" -> True}]

Out[12]= "
        1 argument
        1 Boolean register
        1 Integer register
        7 Real registers
        2 Tensor registers
        Underflow checking off
        Overflow checking off
        Integer overflow checking on
        RuntimeAttributes -> {}

        R0 = A1
        R2 = 1.
        R3 = 7.
        I0 = 2
        R6 = 0.
        Result = T(R2)1

1   T(I2)0 = MainEvaluate[ Hold[IdentityMatrix][ I0]]
2   R1 = R0
3   B0 = R1 <= R2 (tol R3)
4   if[ !B0] goto 10
5   R4 = Cos[ R1]
6   R5 = Sin[ R1]
7   R4 = R4 * R5
8   R5 = R4
9   goto 11
10  R5 = R6
11  R1 = - R5
12  T(R2)1 = R1 + T(I2)0
13  Return
"

Answer 3

The reason about change only the argument of the external function, the behavior changes may be related to the setting of "ExpressionOptimization" under CompilationOptions.

There is a related function called OptimizeExpression under Experimental` context, which I believe, according to this post by Daniel Lichtblau, is the one used by Compile for expression optimization.

OptimizeExpression can detect some common expressions in code:

Needs["Experimental`"]

OptimizeExpression[  IdentityMatrix[2] - f[t]  ]

OptimizedExpression[
    Block[{Compile`$8, Compile`$9}, 
      Compile`$8 = -If[t <= 1., Cos[t] Sin[t], 0.]; 
          Compile`$9 = 1 + Compile`$8;
      {{Compile`$9, Compile`$8}, {Compile`$8, Compile`$9}}
         ]]

But somehow it fails to detect the common expression in IdentityMatrix[2] - f[t + 1./2.]:

OptimizeExpression[   IdentityMatrix[2] - f[t + 1./2.]   ]

OptimizedExpression[
  Block[{Compile`$5, Compile`$13, Compile`$18, Compile`$23},
  {{1 - If[0.5` + t <= 1.`, Compile`$5 = 0.5` + t; Cos[Compile`$5] Sin[Compile`$5], 0.`],

-If[0.5+ t <= 1., Compile$13 = 0.5 + t; Cos[Compile$13] Sin[Compile$13], 0.`]}, {-If[0.5` + t <= 1.`, Compile`$18 = 0.5` + t; Cos[Compile`$18] Sin[Compile`$18], 0.`], 1 - If[0.5` + t <= 1.`, Compile`$23 = 0.5` + t; Cos[Compile`$23] Sin[Compile`$23], 0.`] }} ]]

One of the possible (and hopefully temporary) solution would be substituting definition of f after the OptimizeExpression being applied:

compfunc = Compile[{{t, _Real}},
        Evaluate[
            OptimizeExpression[
                    Hold[
                                IdentityMatrix[2] - f[t + 1./2.]
                                ] /. f -> fTemp // ReleaseHold
                    ] /. fTemp[x_] :> With[{val = f[x]}, val /; True]
            ]
        ]

Needs["CompiledFunctionTools`"]
compfunc // CompilePrint

 "(*omitted*)

 1    R2 = R1 + R0
 2    B0 = R2 <= R3 (tol R4)
 3    if[ !B0] goto 9
 4    R5 = Cos[ R2]
 5    R6 = Sin[ R2]
 6    R5 = R5 * R6
 7    R6 = R5
 8    goto 10
 9    R6 = R7
 10   R5 = - R6
 11   R8 = I0
 12   R8 = R8 + R5
 13   T(R1)0 ={ R8, R5 }
 14   T(R1)1 ={ R5, R8 }
 15   T(R2)2 ={ T(R1)0, T(R1)1 }
 16   Return
 "