mlx 0.30.7.3 → 0.30.7.6

data/lib/mlx/onnx.rb

@@ -0,0 +1,250 @@
+# frozen_string_literal: true
+
+require "json"
+require_relative "../mlx-onnx/webgpu_harness"
+
+module MLX
+  module ONNX
+    ONNX_VERSION = 1
+
+    module_function
+
+    def export_onnx(
+      target_path,
+      fun,
+      *extras,
+      shapeless: false,
+      opset: 18,
+      model_name: "mlx_graph",
+      external_data: false,
+      external_data_file: nil,
+      external_data_size_threshold: 1024,
+      **trace_kwargs
+    )
+      ensure_native_graph_ir!
+      assert_boolean!(shapeless, "shapeless")
+      assert_boolean!(external_data, "external_data")
+      native_target_path = normalize_binary_target_path!(target_path, "export_onnx")
+      native_fun, native_extras, _positional_count, _keyword_names =
+        prepare_trace_invocation_without_native_kwargs(fun, extras, trace_kwargs)
+      translate_native_unsupported do
+        written = MLX::ONNX::Native.export_onnx(
+          native_target_path,
+          native_fun,
+          native_extras,
+          {},
+          shapeless,
+          opset,
+          model_name,
+          external_data,
+          external_data_file,
+          external_data_size_threshold
+        )
+        unless written.is_a?(String)
+          raise TypeError, "MLX::ONNX::Native.export_onnx must return String path"
+        end
+        written
+      end
+    end
+
+    def export_onnx_json(
+      fun,
+      *extras,
+      shapeless: false,
+      opset: 18,
+      model_name: "mlx_graph",
+      **trace_kwargs
+    )
+      ensure_native_graph_ir!
+      assert_boolean!(shapeless, "shapeless")
+      native_fun, native_extras, _positional_count, _keyword_names =
+        prepare_trace_invocation_without_native_kwargs(fun, extras, trace_kwargs)
+      translate_native_unsupported do
+        onnx_json = MLX::ONNX::Native.export_onnx_json(
+          native_fun,
+          native_extras,
+          {},
+          shapeless,
+          opset,
+          model_name
+        )
+        unless onnx_json.is_a?(String)
+          raise TypeError, "MLX::ONNX::Native.export_onnx_json must return String JSON"
+        end
+        onnx_json
+      end
+    end
+
+    def export_onnx_compatibility_report(fun, *extras, shapeless: false, **trace_kwargs)
+      ensure_native_graph_ir!
+      assert_boolean!(shapeless, "shapeless")
+      native_fun, native_extras, _positional_count, _keyword_names =
+        prepare_trace_invocation_without_native_kwargs(fun, extras, trace_kwargs)
+      report = MLX::ONNX::Native.export_onnx_compatibility_report(
+        native_fun,
+        native_extras,
+        {},
+        shapeless
+      )
+      unless report.is_a?(Hash)
+        raise TypeError, "MLX::ONNX::Native.export_onnx_compatibility_report must return Hash payload"
+      end
+
+      report
+    end
+
+    def export_graph_ir(fun, *extras, shapeless: false, **trace_kwargs)
+      ensure_native_graph_ir!
+      assert_boolean!(shapeless, "shapeless")
+      native_fun, native_extras, positional_count, keyword_names =
+        prepare_trace_invocation_without_native_kwargs(fun, extras, trace_kwargs)
+      payload = MLX::ONNX::Native.export_graph_ir(native_fun, native_extras, {}, shapeless)
+      unless payload.is_a?(Hash)
+        raise TypeError, "MLX::ONNX::Native.export_graph_ir must return Hash payload"
+      end
+      inject_keyword_inputs!(payload, positional_count, keyword_names)
+      payload
+    end
+
+    def export_graph_ir_json(fun, *extras, shapeless: false, **trace_kwargs)
+      ensure_native_graph_ir!
+      assert_boolean!(shapeless, "shapeless")
+      native_fun, native_extras, positional_count, keyword_names =
+        prepare_trace_invocation_without_native_kwargs(fun, extras, trace_kwargs)
+      content = MLX::ONNX::Native.export_graph_ir_json(native_fun, native_extras, {}, shapeless)
+      unless content.is_a?(String)
+        raise TypeError, "MLX::ONNX::Native.export_graph_ir_json must return String JSON"
+      end
+      return content if keyword_names.empty?
+
+      payload = JSON.parse(content)
+      inject_keyword_inputs!(payload, positional_count, keyword_names)
+      content = JSON.generate(payload)
+      content
+    end
+
+    def graph_ir_to_onnx(
+      target_path,
+      ir_source,
+      opset: 18,
+      model_name: "mlx_graph",
+      external_data: false,
+      external_data_file: nil,
+      external_data_size_threshold: 1024
+    )
+      ensure_native_graph_ir!
+      assert_boolean!(external_data, "external_data")
+      native_target_path = normalize_binary_target_path!(target_path, "graph_ir_to_onnx")
+      translate_native_unsupported do
+        written = MLX::ONNX::Native.graph_ir_to_onnx(
+          native_target_path,
+          ir_source,
+          opset,
+          model_name,
+          external_data,
+          external_data_file,
+          external_data_size_threshold
+        )
+        unless written.is_a?(String)
+          raise TypeError, "MLX::ONNX::Native.graph_ir_to_onnx must return String path"
+        end
+        written
+      end
+    end
+
+    def graph_ir_to_onnx_json(ir_source, opset: 18, model_name: "mlx_graph")
+      ensure_native_graph_ir!
+      translate_native_unsupported do
+        onnx_json = MLX::ONNX::Native.graph_ir_to_onnx_json(ir_source, opset, model_name)
+        unless onnx_json.is_a?(String)
+          raise TypeError, "MLX::ONNX::Native.graph_ir_to_onnx_json must return String JSON"
+        end
+        onnx_json
+      end
+    end
+
+    def ensure_native_graph_ir!
+      MLX::Core.ensure_native!
+      unless defined?(MLX::ONNX::Native)
+        raise RuntimeError, "MLX::ONNX::Native is unavailable"
+      end
+    end
+    private_class_method :ensure_native_graph_ir!
+
+    def assert_boolean!(value, label)
+      unless value == true || value == false
+        raise TypeError, "#{label} must be true or false"
+      end
+    end
+    private_class_method :assert_boolean!
+
+    def translate_native_unsupported
+      yield
+    rescue MLX::ONNX::Native::UnsupportedError => e
+      raise NotImplementedError, e.message
+    end
+    private_class_method :translate_native_unsupported
+
+    def normalize_binary_target_path!(target_path, method_name)
+      if target_path.respond_to?(:write) && !target_path.respond_to?(:to_path)
+        raise ArgumentError, "#{method_name} requires a path-like target, not an IO-like target"
+      end
+
+      path = if target_path.respond_to?(:to_path)
+        target_path.to_path
+      else
+        target_path
+      end
+
+      unless path.is_a?(String)
+        raise TypeError, "target_path must be a String or respond to #to_path"
+      end
+
+      path
+    end
+    private_class_method :normalize_binary_target_path!
+
+    def prepare_trace_invocation_without_native_kwargs(fun, extras, trace_kwargs)
+      positional_inputs = extras.dup
+      return [fun, positional_inputs, positional_inputs.length, []] if trace_kwargs.empty?
+
+      keyword_entries = trace_kwargs.to_a
+      positional_count = positional_inputs.length
+      native_fun = lambda do |*all_args|
+        positional_args = all_args.first(positional_count)
+        keyword_args = {}
+        keyword_entries.each_with_index do |(name, _value), idx|
+          keyword_args[name] = all_args[positional_count + idx]
+        end
+        fun.call(*positional_args, **keyword_args)
+      end
+
+      keyword_values = keyword_entries.map { |(_name, value)| value }
+      keyword_names = keyword_entries.map { |(name, _value)| name.to_s }
+      [native_fun, positional_inputs + keyword_values, positional_count, keyword_names]
+    end
+    private_class_method :prepare_trace_invocation_without_native_kwargs
+
+    def inject_keyword_inputs!(payload, positional_count, keyword_names)
+      return payload if keyword_names.empty?
+      return payload unless payload.is_a?(Hash)
+
+      inputs = payload["inputs"]
+      return payload unless inputs.is_a?(Array)
+
+      keyword_inputs = []
+      keyword_names.each_with_index do |name, idx|
+        input_entry = inputs[positional_count + idx]
+        next unless input_entry.is_a?(Hash)
+
+        tensor = input_entry["name"]
+        next unless tensor.is_a?(String) && !tensor.empty?
+
+        keyword_inputs << { "name" => name, "tensor" => tensor }
+      end
+      payload["keyword_inputs"] = keyword_inputs
+      payload
+    end
+    private_class_method :inject_keyword_inputs!
+  end
+end

data/lib/mlx/version.rb

@@ -1,5 +1,5 @@
 # frozen_string_literal: true
 
 module MLX
-  VERSION = "0.30.7.3"
+  VERSION = "0.30.7.6"
 end

data/lib/mlx-onnx/webgpu_harness.rb

@@ -0,0 +1,289 @@
+# frozen_string_literal: true
+
+require "json"
+require "open3"
+require "fileutils"
+
+module MLX
+  module ONNX
+    module WebGPUHarness
+      module_function
+
+      def export_onnx_webgpu_harness(
+        target_dir,
+        payload_or_source,
+        opset: 18,
+        model_name: "mlx_graph",
+        execution_providers: %w[webgpu wasm],
+        benchmark_warmup_runs: 2,
+        benchmark_measure_runs: 10,
+        external_data: false,
+        external_data_size_threshold: 1024,
+        external_data_file: nil
+      )
+        output_dir = file_path(target_dir)
+        raise ArgumentError, "target_dir must not be empty" if output_dir.empty?
+
+        providers = normalize_web_execution_providers(execution_providers)
+        warmup_runs = normalize_non_negative_integer(
+          benchmark_warmup_runs,
+          "benchmark_warmup_runs"
+        )
+        measure_runs = normalize_positive_integer(
+          benchmark_measure_runs,
+          "benchmark_measure_runs"
+        )
+
+        FileUtils.mkdir_p(output_dir)
+        model_filename = "model.onnx"
+        model_path = File.join(output_dir, model_filename)
+        onnx_json = MLX::ONNX.graph_ir_to_onnx_json(
+          payload_or_source,
+          opset: opset,
+          model_name: model_name
+        )
+        MLX::ONNX.graph_ir_to_onnx(
+          model_path,
+          payload_or_source,
+          opset: opset,
+          model_name: model_name,
+          external_data: external_data,
+          external_data_size_threshold: external_data_size_threshold,
+          external_data_file: external_data_file
+        )
+
+        stub = JSON.parse(onnx_json)
+        input_specs = stub.fetch("graph").fetch("inputs")
+        input_examples = build_input_examples(input_specs)
+
+        manifest = {
+          "format" => "onnx_webgpu_harness_v1",
+          "model" => model_filename,
+          "execution_providers" => providers,
+          "benchmark" => {
+            "warmup_runs" => warmup_runs,
+            "measure_runs" => measure_runs
+          },
+          "inputs" => input_specs.map do |spec|
+            {
+              "name" => spec.fetch("name"),
+              "shape" => spec.fetch("shape"),
+              "dtype" => spec.fetch("dtype")
+            }
+          end
+        }
+        if external_data
+          manifest["external_data"] = [
+            external_data_file.nil? ? "model.data" : external_data_file.to_s
+          ]
+        end
+
+        File.binwrite(
+          File.join(output_dir, "harness.manifest.json"),
+          JSON.pretty_generate(manifest)
+        )
+        File.binwrite(
+          File.join(output_dir, "inputs.example.json"),
+          JSON.pretty_generate(input_examples)
+        )
+        copy_assets!(output_dir)
+
+        manifest
+      end
+
+      def smoke_test_onnx_webgpu_harness(
+        harness_dir,
+        timeout_seconds: 30,
+        mock_ort: false,
+        local_ort: true,
+        node_bin: ENV.fetch("NODE", "node")
+      )
+        directory = file_path(harness_dir)
+        raise ArgumentError, "harness_dir must not be empty" if directory.empty?
+
+        directory = File.expand_path(directory)
+        unless Dir.exist?(directory)
+          raise ArgumentError, "harness_dir does not exist: #{directory}"
+        end
+
+        timeout = normalize_positive_integer(timeout_seconds, "timeout_seconds")
+        mock = normalize_boolean(mock_ort, "mock_ort")
+        local = normalize_boolean(local_ort, "local_ort")
+        node = node_bin.to_s
+        raise ArgumentError, "node_bin must not be empty" if node.empty?
+
+        smoke_script = web_harness_smoke_script_path
+        unless File.file?(smoke_script)
+          raise RuntimeError, "missing web harness smoke script: #{smoke_script}"
+        end
+
+        argv = [
+          node,
+          smoke_script,
+          "--harness-dir",
+          directory,
+          "--timeout-seconds",
+          timeout.to_s
+        ]
+        argv << "--mock-ort" if mock
+        argv << (local ? "--local-ort" : "--no-local-ort")
+
+        stdout, stderr, status = Open3.capture3(*argv, chdir: web_root_dir)
+        unless status.success?
+          raise RuntimeError, <<~MSG
+            web harness smoke test failed: #{argv.join(" ")}
+            stdout:
+            #{stdout}
+            stderr:
+            #{stderr}
+          MSG
+        end
+
+        telemetry = begin
+          JSON.parse(stdout)
+        rescue JSON::ParserError => e
+          raise RuntimeError, <<~MSG
+            web harness smoke test produced invalid JSON: #{e.message}
+            stdout:
+            #{stdout}
+            stderr:
+            #{stderr}
+          MSG
+        end
+
+        unless telemetry.is_a?(Hash)
+          raise RuntimeError, "web harness smoke test produced non-object telemetry"
+        end
+        unless telemetry.fetch("format", nil) == "onnx_webgpu_telemetry_v1"
+          raise RuntimeError, "unexpected web harness telemetry format: #{telemetry.fetch('format', nil).inspect}"
+        end
+
+        telemetry
+      end
+
+      def file_path(file)
+        if file.respond_to?(:to_path)
+          file.to_path.to_s
+        else
+          file.to_s
+        end
+      end
+      private_class_method :file_path
+
+      def normalize_web_execution_providers(value)
+        providers = if value.is_a?(::Array)
+          value
+        else
+          [value]
+        end
+        providers = providers.map(&:to_s)
+        raise ArgumentError, "execution_providers must contain at least one provider" if providers.empty?
+
+        allowed = %w[webgpu wasm]
+        providers.each do |provider|
+          unless allowed.include?(provider)
+            raise ArgumentError, "execution_providers contains unsupported provider #{provider.inspect}"
+          end
+        end
+        providers.uniq
+      end
+      private_class_method :normalize_web_execution_providers
+
+      def normalize_non_negative_integer(value, label)
+        integer = begin
+          Integer(value)
+        rescue ArgumentError, TypeError
+          raise ArgumentError, "#{label} must be a non-negative Integer"
+        end
+        raise ArgumentError, "#{label} must be a non-negative Integer" if integer.negative?
+
+        integer
+      end
+      private_class_method :normalize_non_negative_integer
+
+      def normalize_positive_integer(value, label)
+        integer = begin
+          Integer(value)
+        rescue ArgumentError, TypeError
+          raise ArgumentError, "#{label} must be a positive Integer"
+        end
+        raise ArgumentError, "#{label} must be a positive Integer" unless integer.positive?
+
+        integer
+      end
+      private_class_method :normalize_positive_integer
+
+      def normalize_boolean(value, label)
+        unless value == true || value == false
+          raise ArgumentError, "#{label} must be true or false"
+        end
+
+        value
+      end
+      private_class_method :normalize_boolean
+
+      def build_input_examples(input_specs)
+        input_specs.each_with_object({}) do |spec, out|
+          out[spec.fetch("name")] = build_zero_tensor_values(
+            spec.fetch("shape"),
+            spec.fetch("dtype")
+          )
+        end
+      end
+      private_class_method :build_input_examples
+
+      def build_zero_tensor_values(shape, dtype)
+        if shape.empty?
+          zero_leaf_value_for_dtype(dtype)
+        else
+          ::Array.new(shape.first) { build_zero_tensor_values(shape[1..], dtype) }
+        end
+      end
+      private_class_method :build_zero_tensor_values
+
+      def zero_leaf_value_for_dtype(dtype)
+        if dtype == "bool" || dtype == "bool_"
+          false
+        elsif dtype == "complex64"
+          { "__mlx_complex__" => [0.0, 0.0] }
+        elsif dtype.start_with?("float") || dtype == "bfloat16"
+          0.0
+        else
+          0
+        end
+      end
+      private_class_method :zero_leaf_value_for_dtype
+
+      def copy_assets!(output_dir)
+        template_dir = web_harness_template_dir
+        unless Dir.exist?(template_dir)
+          raise RuntimeError, "missing web harness template directory: #{template_dir}"
+        end
+
+        %w[index.html harness.js].each do |file_name|
+          source = File.join(template_dir, file_name)
+          unless File.file?(source)
+            raise RuntimeError, "missing web harness template file: #{source}"
+          end
+          FileUtils.cp(source, File.join(output_dir, file_name))
+        end
+      end
+      private_class_method :copy_assets!
+
+      def web_harness_template_dir
+        File.expand_path("../../web/onnx_webgpu_harness", __dir__)
+      end
+      private_class_method :web_harness_template_dir
+
+      def web_harness_smoke_script_path
+        File.join(web_harness_template_dir, "browser_smoke.mjs")
+      end
+      private_class_method :web_harness_smoke_script_path
+
+      def web_root_dir
+        File.expand_path("../../web", __dir__)
+      end
+      private_class_method :web_root_dir
+    end
+  end
+end

data/{mlx → submodules/mlx}/mlx/backend/cuda/cublas_utils.cpp

@@ -105,13 +105,6 @@ void CublasMatmulBase::init_base(
   CHECK_CUBLAS_ERROR(
       cublasLtMatmulDescCreate(&matmul_desc_, compute_type, scale_type));
 
-  int32_t pointer_mode = CUBLASLT_POINTER_MODE_HOST;
-  CHECK_CUBLAS_ERROR(cublasLtMatmulDescSetAttribute(
-      matmul_desc_,
-      CUBLASLT_MATMUL_DESC_POINTER_MODE,
-      &pointer_mode,
-      sizeof(int32_t)));
-
   // In cublasLt matrices use column-major layout, while it is possible to use
   // the CUBLASLT_ORDER_ROW option to switch to row-major layout, the bias
   // epilogue does not work with the option. So instead we swap A and B to make

data/{mlx → submodules/mlx}/mlx/backend/cuda/gemms/cublas_gemm.cpp

@@ -73,6 +73,14 @@ CublasGemm::CublasGemm(
       batch_count,
       a_batch_stride,
       b_batch_stride);
+
+  // alpha and beta are both host pointers
+  cublasLtPointerMode_t pointer_mode = CUBLASLT_POINTER_MODE_HOST;
+  CHECK_CUBLAS_ERROR(cublasLtMatmulDescSetAttribute(
+      matmul_desc_,
+      CUBLASLT_MATMUL_DESC_POINTER_MODE,
+      &pointer_mode,
+      sizeof(pointer_mode)));
 }
 
 CublasGemm::CublasGemm(
@@ -215,8 +223,8 @@ void CublasGemm::execute(
     const void* a,
     const void* b,
     const void* c,
-    float alpha /* = 1 */,
-    float beta /* = 0 */) {
+    const float alpha /* = 1 */,
+    const float beta /* = 0 */) {
   const void* alpha_ptr = α
   const void* beta_ptr = β
   complex64_t alpha_c, beta_c;