Skip to content

feat: factorizations #1234

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
wants to merge 9 commits into
base: main
Choose a base branch
from
Draft

feat: factorizations #1234

wants to merge 9 commits into from

Conversation

avik-pal
Copy link
Collaborator

@avik-pal avik-pal commented May 1, 2025

No description provided.

@avik-pal avik-pal marked this pull request as draft May 1, 2025 23:17
@avik-pal
Copy link
Collaborator Author

avik-pal commented May 1, 2025

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 avik-pal force-pushed the ap/fact branch 2 times, most recently from 539a485 to e48006a Compare May 14, 2025 20:15
@avik-pal
Copy link
Collaborator Author

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
Copy link
Member

wsmoses commented May 15, 2025

Another day another jll

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

3 participants