feat: factorizations

[avik-pal]

No description provided.

[avik-pal]

current status

julia> @jit lu(Reactant.to_rarray(rand(4, 4)))
malloc(): invalid size (unsorted)

[29206] signal 6 (-6): Aborted
in expression starting at REPL[13]:1

[avik-pal]

well crap I forgot to link the custom_call stuff from jaxlib

┌ Error: Compilation failed, MLIR module written to /tmp/reactant_iKiI73/module_000_reactant_lu_post_xla_compile.mlir
└ @ Reactant.MLIR.IR /mnt/software/lux/Reactant.jl/src/mlir/IR/Pass.jl:116
ERROR: module @reactant_lu attributes {mhlo.num_partitions = 1 : i64, mhlo.num_replicas = 1 : i64} {
  func.func @main(%arg0: tensor<3x3x4xf64>) -> (tensor<3x3x4xf64>, tensor<3x4xi32>, tensor<3x4xi32>, tensor<4xi32>) {
    %0 = mhlo.constant dense<1> : tensor<4x3xi32>
    %1 = "mhlo.transpose"(%arg0) <{permutation = dense<[2, 1, 0]> : tensor<3xi64>}> : (tensor<3x3x4xf64>) -> tensor<4x3x3xf64>
    %2:3 = mhlo.custom_call @cusolver_getrf_ffi(%1) {operand_layouts = [dense<[1, 2, 0]> : tensor<3xindex>], output_operand_aliases = [#mhlo.output_operand_alias<output_tuple_indices = [0], operand_index = 0, operand_tuple_indices = []>], result_layouts = [dense<[1, 2, 0]> : tensor<3xindex>, dense<[1, 0]> : tensor<2xindex>, dense<0> : tensor<1xindex>]} : (tensor<4x3x3xf64>) -> (tensor<4x3x3xf64>, tensor<4x3xi32>, tensor<4xi32>)
    %3 = mhlo.subtract %2#1, %0 : tensor<4x3xi32>
    %4 = mhlo.custom_call @cu_lu_pivots_to_permutation(%3) : (tensor<4x3xi32>) -> tensor<4x3xi32>
    %5 = mhlo.add %4, %0 : tensor<4x3xi32>
    %6 = "mhlo.transpose"(%2#0) <{permutation = dense<[2, 1, 0]> : tensor<3xi64>}> : (tensor<4x3x3xf64>) -> tensor<3x3x4xf64>
    %7 = "mhlo.transpose"(%2#1) <{permutation = dense<[1, 0]> : tensor<2xi64>}> : (tensor<4x3xi32>) -> tensor<3x4xi32>
    %8 = "mhlo.transpose"(%5) <{permutation = dense<[1, 0]> : tensor<2xi64>}> : (tensor<4x3xi32>) -> tensor<3x4xi32>
    return %6, %7, %8, %2#2 : tensor<3x3x4xf64>, tensor<3x4xi32>, tensor<3x4xi32>, tensor<4xi32>
  }
}
UNIMPLEMENTED: No registered implementation for custom call to cusolver_getrf_ffi for platform CUDA

[wsmoses]

Another day another jll

[avik-pal]

diff --git a/deps/ReactantExtra/BUILD b/deps/ReactantExtra/BUILD
index 413619bf..621dce2b 100644
--- a/deps/ReactantExtra/BUILD
+++ b/deps/ReactantExtra/BUILD
@@ -1040,6 +1040,7 @@ cc_library(
             "@xla//xla/backends/profiler/gpu:device_tracer",
             "@xla//xla/backends/profiler:profiler_backends",
             "@xla//xla/pjrt/c:pjrt_c_api_gpu_internal",
+           "@jax//jaxlib/cuda:cuda_gpu_kernels",
         ],
         "//conditions:default": [
         ],

doesn't seem to fix this

[avik-pal]

cuda backend is producing incorrrect permutations, will take a look into that now

[avik-pal]

Tests are now passing!