rich7420 commented on code in PR #1389: URL: https://github.com/apache/mahout/pull/1389#discussion_r3620193847
########## qdp/qdp-kernels/src/AdaptiveOzaki.cu: ########## @@ -0,0 +1,1382 @@ +// +// Licensed to the Apache Software Foundation (ASF) under one or more +// contributor license agreements. See the NOTICE file distributed with +// this work for additional information regarding copyright ownership. +// The ASF licenses this file to You under the Apache License, Version 2.0 +// (the "License"); you may not use this file except in compliance with +// the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "AdaptiveOzaki.h" +#include <iostream> +#include <cmath> +#include <mma.h> +#include <cuda_fp16.h> + +using namespace nvcuda; + +__device__ __forceinline__ void cp_async_16(void* smem_ptr, const void* global_ptr) { + uint32_t smem_addr = __cvta_generic_to_shared(smem_ptr); + asm volatile("cp.async.cg.shared.global [%0], [%1], 16;" :: "r"(smem_addr), "l"(global_ptr)); +} +__device__ __forceinline__ void cp_async_commit() { + asm volatile("cp.async.commit_group;\n" ::); +} +__device__ __forceinline__ void cp_async_wait_0() { + asm volatile("cp.async.wait_group 0;\n" ::); +} +__device__ __forceinline__ void cp_async_wait_1() { + asm volatile("cp.async.wait_group 1;\n" ::); +} + +__device__ __forceinline__ double pow2_int(int exp) { + return ldexp(1.0, exp); +} + +__device__ __forceinline__ size_t get_A8_offset(int m_idx, int k_idx, int m, int k) { + int tile_m = m_idx >> 7; + int tile_k = k_idx >> 5; + int local_m = m_idx & 127; + int local_k = k_idx & 31; + int num_tiles_k = (k + 31) >> 5; + return (size_t)(tile_m * num_tiles_k + tile_k) * 4096 + (local_m << 5) + local_k; +} + +__device__ __forceinline__ size_t get_B8_offset(int n_idx, int k_idx, int n, int k) { + int tile_n = n_idx >> 6; + int tile_k = k_idx >> 5; + int local_n = n_idx & 63; + int local_k = k_idx & 31; + int num_tiles_k = (k + 31) >> 5; + return (size_t)(tile_n * num_tiles_k + tile_k) * 2048 + (local_n << 5) + local_k; +} + +__device__ __forceinline__ double fp64_hi(double v, int split_bits) { + double scale = pow2_int(split_bits); + double scaled = v * scale; + double high_scaled = static_cast<double>(__double2ll_rn(scaled)); + return high_scaled / scale; +} + +__device__ __forceinline__ void mma_m16n8k32_s8( + int32_t* d, const uint32_t* a, const uint32_t* b, const int32_t* c) { + asm volatile( + "mma.sync.aligned.m16n8k32.row.col.s32.s8.s8.s32 " + "{%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%10, %11, %12, %13};" + : "=r"(d[0]), "=r"(d[1]), "=r"(d[2]), "=r"(d[3]) + : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), + "r"(b[0]), "r"(b[1]), + "r"(c[0]), "r"(c[1]), "r"(c[2]), "r"(c[3]) + ); +} + +namespace ozaki { + +__device__ __forceinline__ void ldmatrix_x4_int8(uint32_t* d, void* smem_ptr) { + uint32_t smem_addr = __cvta_generic_to_shared(smem_ptr); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];" + : "=r"(d[0]), "=r"(d[1]), "=r"(d[2]), "=r"(d[3]) : "r"(smem_addr)); +} + +__device__ __forceinline__ void ldmatrix_x2_int8(uint32_t* d, void* smem_ptr) { + uint32_t smem_addr = __cvta_generic_to_shared(smem_ptr); + asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0, %1}, [%2];" + : "=r"(d[0]), "=r"(d[1]) : "r"(smem_addr)); +} + +AdaptiveOzakiEngine::AdaptiveOzakiEngine(const OzakiConfig& config) : config_(config) { +} +AdaptiveOzakiEngine::~AdaptiveOzakiEngine() { + freeWorkspace(); +} + +void AdaptiveOzakiEngine::allocateWorkspace(int m, int n, int k) { + if (workspace_allocated_) freeWorkspace(); + int nm = (m + 127) / 128, nn = (n + 127) / 128; + cudaMalloc(&dmA_h, nm * 8); cudaMalloc(&dmA_l, nm * 8); Review Comment: None of these `cudaMalloc` calls check their return value — same pattern shows up in `ImplicitHadamardOzaki.cu` too. This is actually the first place in the kernels crate that does raw `cudaMalloc` at all, all the existing code goes through checked allocation paths, so it stands out. If one of these allocations fails (e.g. under VRAM pressure, which seems plausible given `vram_limit_mb` is a config knob here), the code just carries on and writes through a null/garbage pointer instead of surfacing a clean error. Worth adding checks before this ships, even if it's just a fail-fast `assert`/early return. ########## qdp/qdp-kernels/src/lib.rs: ########## @@ -402,6 +417,23 @@ unsafe extern "C" { num_qubits: u32, stream: *mut c_void, ) -> i32; + + /// Launch general (non-Hadamard) GEMM using AdaptiveOzakiEngine + /// Provides mixed-precision graded-ring Tensor Core acceleration for arbitrary matrices. + /// Complements launch_iqp_encode_tc which uses the specialized ImplicitHadamard path. + /// Returns CUDA error code (0 = success) + /// + /// # Safety + /// Requires valid GPU pointers, must sync before freeing + pub fn launch_adaptive_ozaki_gemm( Review Comment: This FFI signature takes a `stream` param, but looking at the `.cu` side, `AdaptiveOzakiEngine::execute()` opens and uses its own internal streams and just does a best-effort `cudaStreamSynchronize(stream)` at the end if one was passed in. So the caller's stream ordering guarantee doesn't actually hold — any async work the caller queued on that stream isn't ordered against this GEMM. Low risk right now since I don't see anything in Rust actually calling this yet (it's declared but unused), but the contract as written is misleading, so I'd flag it before something starts depending on it. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
