feat(qualcomm): Qnn AOT Lowering Passes #584
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,69 @@ | ||
| // Copyright (c) MLLM Team. | ||
| // Licensed under the MIT License. | ||
|
|
||
| #include "mllm/backends/cpu/ops/GatherOp.hpp" | ||
| #include "mllm/core/Tensor.hpp" | ||
|
|
||
| namespace mllm::cpu { | ||
|
|
||
| CPUGatherOp::CPUGatherOp(const aops::GatherOpOptions& options) : aops::GatherOp(options) {} | ||
|
|
||
| void CPUGatherOp::forward(const std::vector<Tensor>& inputs, std::vector<Tensor>& outputs) { | ||
| auto& table = inputs[0]; | ||
| auto& indices = inputs[1]; | ||
| auto& output = outputs[0]; | ||
|
|
||
| int dim = options_.dim; | ||
| if (dim < 0) dim += table.shape().size(); | ||
|
|
||
| int64_t outer_size = 1; | ||
| for (int i = 0; i < dim; ++i) outer_size *= table.shape()[i]; | ||
|
|
||
| int64_t inner_size = 1; | ||
| for (int i = dim + 1; i < table.shape().size(); ++i) inner_size *= table.shape()[i]; | ||
|
|
||
| int64_t dim_size = table.shape()[dim]; | ||
| int64_t indices_count = indices.numel(); | ||
|
|
||
| size_t data_type_size = 4; | ||
| switch (table.dtype()) { | ||
| case MLLM_TYPE_F32: data_type_size = sizeof(float); break; | ||
| case MLLM_TYPE_F16: data_type_size = sizeof(mllm_fp16_t); break; | ||
| case MLLM_TYPE_I32: data_type_size = sizeof(int32_t); break; | ||
| default: MLLM_ERROR("GatherOp table type not supported: {}", (int)table.dtype()); | ||
| } | ||
|
|
||
| const uint8_t* table_ptr = table.ptr<uint8_t>(); | ||
| uint8_t* output_ptr = output.ptr<uint8_t>(); | ||
|
|
||
| const int32_t* indices_i32 = indices.dtype() == MLLM_TYPE_I32 ? indices.ptr<int32_t>() : nullptr; | ||
| const float* indices_f32 = !indices_i32 && indices.dtype() == MLLM_TYPE_F32 ? indices.ptr<float>() : nullptr; | ||
|
|
||
| if (!indices_i32 && !indices_f32) { | ||
| MLLM_ERROR("GatherOp indices type not supported: {}", (int)indices.dtype()); | ||
| return; | ||
| } | ||
|
|
||
| // FIXME: parallel | ||
| for (int64_t o = 0; o < outer_size; ++o) { | ||
| for (int64_t i = 0; i < indices_count; ++i) { | ||
| int64_t idx = 0; | ||
| if (indices_i32) { | ||
| idx = indices_i32[i]; | ||
| } else if (indices_f32) { | ||
| idx = (int64_t)indices_f32[i]; | ||
| } | ||
|
|
||
| if (idx < 0) idx += dim_size; | ||
|
|
||
| if (idx < 0 || idx >= dim_size) { continue; } | ||
|
|
||
| int64_t src_offset = (o * dim_size + idx) * inner_size * data_type_size; | ||
| int64_t dst_offset = (o * indices_count + i) * inner_size * data_type_size; | ||
|
|
||
| std::memcpy(output_ptr + dst_offset, table_ptr + src_offset, inner_size * data_type_size); | ||
| } | ||
| } | ||
| } | ||
|
|
||
| } // namespace mllm::cpu | ||
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,25 @@ | ||
| // Copyright (c) MLLM Team. | ||
| // Licensed under the MIT License. | ||
|
|
||
| #pragma once | ||
|
|
||
| #include "mllm/core/BaseOp.hpp" | ||
| #include "mllm/core/aops/GatherOp.hpp" | ||
|
|
||
| namespace mllm::cpu { | ||
|
|
||
| class CPUGatherOp final : public aops::GatherOp { | ||
| public: | ||
| explicit CPUGatherOp(const aops::GatherOpOptions& options); | ||
|
|
||
| void forward(const std::vector<Tensor>& inputs, std::vector<Tensor>& outputs) override; | ||
| }; | ||
|
|
||
| class CPUGatherOpFactory : public TypedOpFactory<OpTypes::kGather, aops::GatherOpOptions> { | ||
| public: | ||
| std::shared_ptr<BaseOp> createOpImpl(const aops::GatherOpOptions& options) override { | ||
| return std::make_shared<CPUGatherOp>(options); | ||
| } | ||
| }; | ||
|
|
||
| } // namespace mllm::cpu |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -31,6 +31,7 @@ QnnAOTNodeTensor::QnnAOTNodeTensor(const ir::tensor::TensorValue::ptr_t& v, bool | |
| } else { | ||
| tensor_wrapper_ = mllm::qnn::QNNTensorWrapper::create(name, type, v->tensor_, quant); | ||
| } | ||
| setupComplexTensorQuantization(v); // per-channel and LPBQ cases | ||
| } | ||
|
|
||
| Qnn_TensorType_t QnnAOTNodeTensor::parseQnnTensorTypeFromIR(const ir::tensor::TensorValue::ptr_t& v) { | ||
|
|
@@ -90,7 +91,7 @@ Qnn_TensorType_t QnnAOTNodeTensor::parseQnnTensorTypeFromIR(const ir::tensor::Te | |
|
|
||
| // Check Attribute. The Attribute priority is higher than tensor type | ||
| if (v->getAttr("qnn_graph_outputs")) { ret_qnn_tensor_type = QNN_TENSOR_TYPE_APP_READ; } | ||
| if (v->getAttr("qnn_graph_inputs")) { ret_qnn_tensor_type = QNN_TENSOR_TYPE_APP_WRITE; } | ||
| if (v->getAttr("constant")) { ret_qnn_tensor_type = QNN_TENSOR_TYPE_STATIC; } | ||
|
|
||
| return ret_qnn_tensor_type; | ||
|
|
@@ -109,7 +110,16 @@ Qnn_QuantizeParams_t QnnAOTNodeTensor::parseQnnQuantizeParamFromIR(const ir::ten | |
| auto quant_spec = v->getAttr("quant_recipe")->cast_<ir::linalg::LinalgIRQuantizatonSpecAttr>()->spec_; | ||
|
|
||
| switch (quant_spec->type) { | ||
| case ir::linalg::QuantizationSpecType::kRaw: | ||
| case ir::linalg::QuantizationSpecType::kSymPerChannel: | ||
| case ir::linalg::QuantizationSpecType::kLPBQ: { | ||
| break; | ||
| } | ||
| case ir::linalg::QuantizationSpecType::kAsymPerTensor: { | ||
| auto cfg = std::static_pointer_cast<ir::linalg::QuantizationSpecAsymPerTensor>(quant_spec); | ||
| ret.encodingDefinition = QNN_DEFINITION_DEFINED; | ||
| ret.quantizationEncoding = QNN_QUANTIZATION_ENCODING_SCALE_OFFSET; | ||
| ret.scaleOffsetEncoding = Qnn_ScaleOffset_t{.scale = cfg->scale.item<float>(), .offset = cfg->zero_point.item<int32_t>()}; | ||
| break; | ||
| } | ||
| case ir::linalg::QuantizationSpecType::kSymPerTensor: { | ||
|
|
@@ -119,6 +129,19 @@ Qnn_QuantizeParams_t QnnAOTNodeTensor::parseQnnQuantizeParamFromIR(const ir::ten | |
| ret.scaleOffsetEncoding = Qnn_ScaleOffset_t{.scale = cfg->scale.item<float>(), .offset = 0}; | ||
| break; | ||
| } | ||
| default: { | ||
| MLLM_ERROR_EXIT(ExitCode::kCoreError, "Can't handle kNone type"); | ||
| } | ||
| } | ||
|
|
||
| return ret; | ||
| } | ||
|
|
||
| void QnnAOTNodeTensor::setupComplexTensorQuantization(const ir::tensor::TensorValue::ptr_t& v) { | ||
| MLLM_RT_ASSERT(v->getAttr("quant_recipe")); | ||
| auto quant_spec = v->getAttr("quant_recipe")->cast_<ir::linalg::LinalgIRQuantizatonSpecAttr>()->spec_; | ||
|
|
||
| switch (quant_spec->type) { | ||
| case ir::linalg::QuantizationSpecType::kSymPerChannel: { | ||
| auto cfg = std::static_pointer_cast<ir::linalg::QuantizationSpecSymPerChannel>(quant_spec); | ||
|
|
||
|
|
@@ -135,12 +158,6 @@ Qnn_QuantizeParams_t QnnAOTNodeTensor::parseQnnQuantizeParamFromIR(const ir::ten | |
| tensor_wrapper_->setScaleOffsetQuantization(scale_offsets, cfg->ch_axis); | ||
| break; | ||
| } | ||
| case ir::linalg::QuantizationSpecType::kLPBQ: { | ||
| auto cfg = std::static_pointer_cast<ir::linalg::QuantizationSpecLPBQ>(quant_spec); | ||
|
|
||
|
|
@@ -150,28 +167,23 @@ Qnn_QuantizeParams_t QnnAOTNodeTensor::parseQnnQuantizeParamFromIR(const ir::ten | |
| MLLM_RT_ASSERT_EQ(num_scale_offsets, cfg->scale_level_1_fp.size(0)); | ||
| MLLM_RT_ASSERT_EQ(cfg->scale_level_0_int.dtype(), kUInt8); | ||
| for (int i = 0; i < num_scale_offsets; ++i) { | ||
| scale_offsets[i].scale = cfg->scale_level_1_fp.at<float>({i, 0, 0}); | ||
| scale_offsets[i].offset = 0; | ||
| } | ||
|
|
||
| Qnn_BlockwiseExpansion_t blockwise_expansion; | ||
| blockwise_expansion.axis = cfg->ch_axis; | ||
| blockwise_expansion.scaleOffsets = nullptr; // Will be set by setBlockwiseQuantization | ||
| blockwise_expansion.numBlocksPerAxis = v->tensor_.size(1) / cfg->block_size; | ||
| blockwise_expansion.blockScaleBitwidth = 12; // 12 bits for 4 to 16 expansion | ||
|
Comment on lines
174
to
178
Contributor
There was a problem hiding this comment. 🧩 Analysis chain🏁 Script executed: # First, understand the structure of QuantizationSpecLPBQ and ch_axis
rg -n "ch_axis" --type cpp -B 2 -A 2 | head -80Repository: UbiquitousLearning/mllm Length of output: 6583 🏁 Script executed: # Check how blockwise_expansion is used and what axis/numBlocksPerAxis mean
rg -n "blockwise_expansion" --type cpp -B 3 -A 3 | head -100Repository: UbiquitousLearning/mllm Length of output: 1458 🏁 Script executed: # Look for LPBQ configuration and how ch_axis is set
rg -n "QuantizationSpecLPBQ\|kLPBQ" --type cpp -B 2 -A 5 | head -150Repository: UbiquitousLearning/mllm Length of output: 49 🏁 Script executed: # Check the tensor dimensions in context of weight tensors
rg -n "v->tensor_\.size\(" --type cpp -B 2 -A 2 | head -100Repository: UbiquitousLearning/mllm Length of output: 1627 🏁 Script executed: # Search for where LPBQ ch_axis is set in LLMQuantRecipePass
rg -n "QuantizationSpecLPBQ::create" --type cpp -B 5 -A 15 | head -150Repository: UbiquitousLearning/mllm Length of output: 4470 🏁 Script executed: # Look for Linear layer quantization setup - typically ch_axis would be set there
rg -n "Linear\|MatMul" mllm/backends/qnn/aot/passes/LLMQuantRecipePass.cpp -B 2 -A 8 | head -200Repository: UbiquitousLearning/mllm Length of output: 49 🏁 Script executed: # Check what values ch_axis can have - look at initialization and usage patterns
rg -n "ch_axis\s*=" --type cpp -B 1 -A 1 | grep -v ".hpp" | head -100Repository: UbiquitousLearning/mllm Length of output: 81 🏁 Script executed: # Found ch_axis = 0 at line 797. Let's verify if this is the only place LPBQ is created
rg -n "QuantizationSpecLPBQ::create" --type cpp -A 2 | head -50Repository: UbiquitousLearning/mllm Length of output: 831 🏁 Script executed: # Check weight tensor dimensions - for Linear layers, shapes should be [out_features, in_features]
rg -n "shape\|size\(0\)|size\(1\)" mllm/backends/qnn/aot/QnnWrappersAPI.cpp | head -50Repository: UbiquitousLearning/mllm Length of output: 153 🏁 Script executed: # Verify: does the code document what axis the blockwise quantization should be along?
rg -n "blockwise\|block_size" mllm/backends/qnn/aot/QnnWrappersAPI.cpp -B 2 -A 2 | head -80Repository: UbiquitousLearning/mllm Length of output: 49 🏁 Script executed: # Let's check if there are any tests or comments explaining the relationship between ch_axis and tensor dims
rg -n "blockwise\|block_size" mllm/compile/ir/linalg/ -B 2 -A 2 | head -100Repository: UbiquitousLearning/mllm Length of output: 49 🏁 Script executed: # Check the full LPBQ quantization spec to understand semantics
cat mllm/compile/ir/linalg/Attribute.hpp | grep -A 30 "struct QuantizationSpecLPBQ"Repository: UbiquitousLearning/mllm Length of output: 1463 🏁 Script executed: # Let's check the relationship more carefully - look at how scale_level_0_int is used
rg -n "scale_level_0_int" --type cpp -B 2 -A 2 | head -100Repository: UbiquitousLearning/mllm Length of output: 3644 🏁 Script executed: # Check dimensions of weight tensors in Linear quantization - what are size(0) and size(1)?
rg -n "Linear.*weight\|register.*weight" mllm/backends/qnn/aot/passes/LLMQuantRecipePass.cpp -B 3 -A 3 | head -80Repository: UbiquitousLearning/mllm Length of output: 49 🏁 Script executed: # Look for any documentation or schema about blockwise expansion axis semantics
rg -n "Qnn_BlockwiseExpansion" --type cpp -B 5 -A 10 | head -150Repository: UbiquitousLearning/mllm Length of output: 5137 Mismatch between blockwise expansion axis and numBlocksPerAxis calculation.
Suggested fix blockwise_expansion.axis = cfg->ch_axis;
blockwise_expansion.scaleOffsets = nullptr; // Will be set by setBlockwiseQuantization
- blockwise_expansion.numBlocksPerAxis = v->tensor_.size(1) / cfg->block_size;
+ blockwise_expansion.numBlocksPerAxis = v->tensor_.size(cfg->ch_axis) / cfg->block_size;🤖 Prompt for AI Agents |
||
| blockwise_expansion.blockScaleStorageType = QNN_BLOCKWISE_EXPANSION_BITWIDTH_SCALE_STORAGE_8; | ||
| blockwise_expansion.blocksScale8 = cfg->scale_level_0_int.ptr<mllm_uint8_t>(); | ||
|
|
||
| tensor_wrapper_->setBlockwiseQuantization(blockwise_expansion, scale_offsets); | ||
| break; | ||
| } | ||
| default: break; | ||
| } | ||
| } | ||
|
|
||
| // QnnAOTNodeOperation implementations | ||
|
|
||
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Silent skip of out-of-bounds indices may leave output uninitialized.
When
idx < 0 || idx >= dim_sizeafter negative index adjustment, the code silently skips the copy, leaving that portion of the output tensor uninitialized. This could result in garbage values or undefined behavior downstream.Consider either:
🔎 Proposed fix: Log warning and zero-fill
📝 Committable suggestion
🤖 Prompt for AI Agents