Star649
435
代码
代码
Issues28
Pull Requests104
Actions
项目讨论
Wiki
分析
Star649
435
已合并
【feat】: atc support onnx int4 #4420
kantao1创建于 3 天前
【feat】: atc support onnx int4 #4420
已合并
kantao1创建于 3 天前
从kantao1/gesupport_onnx_int4合入到cann/gedevelop
共 7 个文件变更+216-1
Mgraph_metadef/proto/onnx/ge_onnx.proto+4-0
@@ -350,6 +350,10 @@ message TensorProto {
350 FLOAT8E4M3FN = 17; // float 8, mostly used for coefficients, supports nan, not inf350 FLOAT8E4M3FN = 17; // float 8, mostly used for coefficients, supports nan, not inf
351 FLOAT8E5M2 = 19; // follows IEEE 754, supports nan, inf, mostly used for gradients351 FLOAT8E5M2 = 19; // follows IEEE 754, supports nan, inf, mostly used for gradients
352 352 
353+ // 4-bit signed integer type (ONNX IR v9, opset 21+).
354+ // Packed storage: 2 int4 values per byte (lower index in lower nibble).
355+ INT4 = 22;
356+ 
353 // Future extensions go here.357 // Future extensions go here.
354 }358 }
355 359 
Mparser/parser/onnx/onnx_constant_parser.cc+15-0
@@ -68,6 +68,8 @@ Status OnnxConstantParser::ParseConvertData(const ge::onnx::TensorProto &tensor_
68 // for uint64 and uint32 values68 // for uint64 and uint32 values
69 {OnnxDataType::UINT64, tensor_proto.uint64_data_size()},69 {OnnxDataType::UINT64, tensor_proto.uint64_data_size()},
70 {OnnxDataType::UINT32, tensor_proto.uint64_data_size()},70 {OnnxDataType::UINT32, tensor_proto.uint64_data_size()},
71+ // for int4 values (packed: 8 int4 per int32)
72+ {OnnxDataType::INT4, tensor_proto.int32_data_size()},
71 };73 };
72 74 
73 int32_t datatype_val_size = 0;75 int32_t datatype_val_size = 0;
@@ -151,6 +153,19 @@ void OnnxConstantParser::ParseConvertDataElements(const ge::onnx::TensorProto &t
151 case OnnxDataType::UINT32:153 case OnnxDataType::UINT32:
152 (void)SetTensorData(tensor_proto.uint64_data_size(), tensor_proto.uint64_data(), count, tensor);154 (void)SetTensorData(tensor_proto.uint64_data_size(), tensor_proto.uint64_data(), count, tensor);
153 break;155 break;
156+ // for int4 values (packed: 8 int4 per int32, little-endian, lower index in lower nibble).
157+ // int32_data stores packed int4 values; direct byte copy since GE DT_INT4 is also packed.
158+ // Use GetSizeInBytes to get exact packed byte count (int32_data may have padding slots).
159+ case OnnxDataType::INT4: {
160+ int64_t byte_size = ge::GetSizeInBytes(static_cast<int64_t>(count), ge::DataType::DT_INT4);
161+ int64_t available = static_cast<int64_t>(tensor_proto.int32_data_size()) * static_cast<int64_t>(sizeof(int32_t));
162+ if (byte_size > 0 && available > 0) {
163+ int64_t copy_size = (byte_size < available) ? byte_size : available;
164+ tensor.SetData(PtrToPtr<const int32_t, const uint8_t>(tensor_proto.int32_data().data()),
165+ static_cast<size_t>(copy_size));
166+ }
167+ break;
168+ }
154 default:169 default:
155 break;170 break;
156 }171 }
Mparser/parser/onnx/onnx_file_constant_parser.cc+1-1
@@ -84,7 +84,7 @@ void OnnxFileConstantParser::ParseShape(const ge::onnx::TensorProto &tensor_prot
84Status OnnxFileConstantParser::ParseDataType(const ge::onnx::TensorProto &tensor_proto, ge::Operator &op_def) const {84Status OnnxFileConstantParser::ParseDataType(const ge::onnx::TensorProto &tensor_proto, ge::Operator &op_def) const {
85 int64_t data_type = tensor_proto.data_type();85 int64_t data_type = tensor_proto.data_type();
86 ge::DataType type = ge::OnnxUtil::ConvertOnnxDataType(data_type);86 ge::DataType type = ge::OnnxUtil::ConvertOnnxDataType(data_type);
87- if (type >= ge::DataType::DT_UNDEFINED) {87+ if (type == ge::DataType::DT_UNDEFINED) {
88 REPORT_INNER_ERR_MSG("E19999", "tensor_proto data type %" PRId64 " is undefined.", data_type);88 REPORT_INNER_ERR_MSG("E19999", "tensor_proto data type %" PRId64 " is undefined.", data_type);
89 GELOGE(domi::PARAM_INVALID, "[Check][Param] tensor_proto data type %" PRId64 " is undefined.", data_type);89 GELOGE(domi::PARAM_INVALID, "[Check][Param] tensor_proto data type %" PRId64 " is undefined.", data_type);
90 return FAILED;90 return FAILED;
Mparser/parser/onnx/onnx_util.cc+1-0
@@ -32,6 +32,7 @@ const std::map<uint32_t, ge::DataType> onnx_data_type_map = {
32 {OnnxDataType::BFLOAT16, ge::DataType::DT_BF16},32 {OnnxDataType::BFLOAT16, ge::DataType::DT_BF16},
33 {OnnxDataType::FLOAT8E5M2, ge::DataType::DT_FLOAT8_E5M2},33 {OnnxDataType::FLOAT8E5M2, ge::DataType::DT_FLOAT8_E5M2},
34 {OnnxDataType::FLOAT8E4M3FN, ge::DataType::DT_FLOAT8_E4M3FN},34 {OnnxDataType::FLOAT8E4M3FN, ge::DataType::DT_FLOAT8_E4M3FN},
35+ {OnnxDataType::INT4, ge::DataType::DT_INT4},
35};36};
36}37}
37 38 
Mparser/parser/onnx/onnx_util.h+1-0
@@ -38,6 +38,7 @@ enum OnnxDataType {
38 FLOAT8E4M3FNUZ = 18,38 FLOAT8E4M3FNUZ = 18,
39 FLOAT8E5M2 = 19,39 FLOAT8E5M2 = 19,
40 FLOAT8E5M2FNUZ = 20,40 FLOAT8E5M2FNUZ = 20,
41+ INT4 = 22,
41};42};
42}43}
43 44 
Mtests/parser/st/testcase/test_onnx_parser.cc+90-0
@@ -301,4 +301,94 @@ TEST_F(STestOnnxParser, onnx_test_SetExternalPath) {
301 auto ret = modelParser.SetExternalPath("/usr/local", model_proto);301 auto ret = modelParser.SetExternalPath("/usr/local", model_proto);
302 EXPECT_EQ(ret, SUCCESS);302 EXPECT_EQ(ret, SUCCESS);
303}303}
304+ 
305+static ge::onnx::ModelProto CreateInt4ModelProto(bool use_raw_data) {
306+ ge::onnx::ModelProto model_proto;
307+ auto *onnx_graph = model_proto.mutable_graph();
308+ 
309+ auto *input = onnx_graph->add_input();
310+ input->set_name("A");
311+ auto *in_type = input->mutable_type()->mutable_tensor_type();
312+ in_type->set_elem_type(OnnxDataType::FLOAT);
313+ in_type->mutable_shape()->add_dim()->set_dim_value(8);
314+ 
315+ auto *output = onnx_graph->add_output();
316+ output->set_name("Y");
317+ auto *out_type = output->mutable_type()->mutable_tensor_type();
318+ out_type->set_elem_type(OnnxDataType::FLOAT);
319+ out_type->mutable_shape()->add_dim()->set_dim_value(8);
320+ 
321+ auto *const_node = onnx_graph->add_node();
322+ const_node->set_op_type(kOpTypeConstant);
323+ const_node->add_output("const_int4_out");
324+ auto *attr = const_node->add_attribute();
325+ attr->set_name(ge::kAttrNameValue);
326+ auto *tensor_proto = attr->mutable_t();
327+ tensor_proto->set_data_type(OnnxDataType::INT4);
328+ tensor_proto->add_dims(8);
329+ if (use_raw_data) {
330+ tensor_proto->set_raw_data(std::string("\x10\x32\x54\x76", 4));
331+ } else {
332+ tensor_proto->add_int32_data(0x76543210);
333+ }
334+ 
335+ auto *identity_node = onnx_graph->add_node();
336+ identity_node->set_op_type("Identity");
337+ identity_node->add_input("A");
338+ identity_node->add_output("Y");
339+ 
340+ auto *op_st = model_proto.add_opset_import();
341+ op_st->set_domain("ai.onnx");
342+ op_st->set_version(11);
343+ return model_proto;
344+}
345+ 
346+static void VerifyInt4ConstantNode(const ge::Graph &graph) {
347+ auto compute_graph = ge::GraphUtilsEx::GetComputeGraph(graph);
348+ ASSERT_NE(compute_graph, nullptr);
349+ ge::NodePtr constant_node = nullptr;
350+ for (const auto &node : compute_graph->GetAllNodes()) {
351+ if (node->GetType() == "Const") {
352+ constant_node = node;
353+ break;
354+ }
355+ }
356+ ASSERT_NE(constant_node, nullptr);
357+ std::shared_ptr<const ge::GeTensor> tensor = nullptr;
358+ EXPECT_EQ(ge::AttrUtils::GetTensor(constant_node->GetOpDesc(), ge::kAttrNameValue, tensor), true);
359+ ASSERT_NE(tensor, nullptr);
360+ EXPECT_EQ(tensor->GetTensorDesc().GetDataType(), ge::DataType::DT_INT4);
361+ const ge::TensorData &tensor_data = tensor->GetData();
362+ EXPECT_EQ(tensor_data.GetSize(), 4U);
363+ const uint8_t *data = tensor_data.GetData();
364+ ASSERT_NE(data, nullptr);
365+ EXPECT_EQ(data[0], 0x10);
366+ EXPECT_EQ(data[1], 0x32);
367+ EXPECT_EQ(data[2], 0x54);
368+ EXPECT_EQ(data[3], 0x76);
369+}
370+ 
371+/**
372+ * 用例描述：测试ONNX INT4类型Constant节点raw_data路径的解析
373+ * 预期结果：解析成功，DataType为DT_INT4，数据为packed字节 0x10,0x32,0x54,0x76
374+ */
375+TEST_F(STestOnnxParser, onnx_parser_int4_const_raw_data) {
376+ OnnxModelParser modelParser;
377+ auto model_proto = CreateInt4ModelProto(true);
378+ ge::Graph graph;
379+ EXPECT_EQ(modelParser.ModelParseToGraph(model_proto, graph), SUCCESS);
380+ VerifyInt4ConstantNode(graph);
381+}
382+ 
383+/**
384+ * 用例描述：测试ONNX INT4类型Constant节点int32_data路径的解析
385+ * 预期结果：解析成功，DataType为DT_INT4，数据为packed字节（little-endian）
386+ */
387+TEST_F(STestOnnxParser, onnx_parser_int4_const_int32_data) {
388+ OnnxModelParser modelParser;
389+ auto model_proto = CreateInt4ModelProto(false);
390+ ge::Graph graph;
391+ EXPECT_EQ(modelParser.ModelParseToGraph(model_proto, graph), SUCCESS);
392+ VerifyInt4ConstantNode(graph);
393+}
304} // namespace ge394} // namespace ge
Mtests/parser/ut/parser/testcase/onnx_parser_testcase/onnx_parser_unittest.cc+104-0
@@ -1710,4 +1710,108 @@ TEST_F(UtestOnnxParser, onnx_model_parse_with_multiple_inputs) {
1710 auto ret = model_parser.ModelParseToGraph(model_proto, root_graph);1710 auto ret = model_parser.ModelParseToGraph(model_proto, root_graph);
1711 EXPECT_EQ(ret, SUCCESS);1711 EXPECT_EQ(ret, SUCCESS);
1712}1712}
1713+ 
1714+// ============ DT_INT4 support tests ============
1715+ 
1716+TEST_F(UtestOnnxParser, ParseConvertDataType_Int4) {
1717+ OnnxConstantParser constant_parser;
1718+ ge::onnx::TensorProto tensor_proto;
1719+ tensor_proto.set_data_type(OnnxDataType::INT4);
1720+ ge::Tensor tensor;
1721+ Status ret = constant_parser.ParseConvertDataType(tensor_proto, tensor);
1722+ EXPECT_EQ(ret, SUCCESS);
1723+ EXPECT_EQ(tensor.GetTensorDesc().GetDataType(), ge::DataType::DT_INT4);
1724+}
1725+ 
1726+TEST_F(UtestOnnxParser, ParseConvertData_Int4_RawData) {
1727+ OnnxConstantParser constant_parser;
1728+ ge::onnx::TensorProto tensor_proto;
1729+ tensor_proto.set_data_type(OnnxDataType::INT4);
1730+ // 2 int4 values packed in 1 byte: 0x21 => low nibble=1, high nibble=2
1731+ tensor_proto.set_raw_data(std::string(1, static_cast<char>(0x21)));
1732+ ge::Tensor tensor;
1733+ TensorDesc tensor_desc = tensor.GetTensorDesc();
1734+ tensor_desc.SetDataType(ge::DataType::DT_INT4);
1735+ tensor.SetTensorDesc(tensor_desc);
1736+ int count = 2;
1737+ Status ret = constant_parser.ParseConvertData(tensor_proto, tensor, count);
1738+ EXPECT_EQ(ret, SUCCESS);
1739+ EXPECT_EQ(tensor.GetSize(), 1U);
1740+ auto data = tensor.GetData();
1741+ ASSERT_NE(data, nullptr);
1742+ EXPECT_EQ(data[0], 0x21);
1743+}
1744+ 
1745+TEST_F(UtestOnnxParser, ParseConvertData_Int4_Int32Data) {
1746+ OnnxConstantParser constant_parser;
1747+ ge::onnx::TensorProto tensor_proto;
1748+ tensor_proto.set_data_type(OnnxDataType::INT4);
1749+ // 1 int32 = 8 packed int4 values: 0x76543210
1750+ tensor_proto.add_int32_data(0x76543210);
1751+ ge::Tensor tensor;
1752+ TensorDesc tensor_desc = tensor.GetTensorDesc();
1753+ tensor_desc.SetDataType(ge::DataType::DT_INT4);
1754+ tensor.SetTensorDesc(tensor_desc);
1755+ int count = 8;
1756+ Status ret = constant_parser.ParseConvertData(tensor_proto, tensor, count);
1757+ EXPECT_EQ(ret, SUCCESS);
1758+ EXPECT_EQ(tensor.GetSize(), sizeof(int32_t));
1759+ auto data = tensor.GetData();
1760+ ASSERT_NE(data, nullptr);
1761+ // little-endian: 0x10, 0x32, 0x54, 0x76
1762+ EXPECT_EQ(data[0], 0x10);
1763+ EXPECT_EQ(data[1], 0x32);
1764+ EXPECT_EQ(data[2], 0x54);
1765+ EXPECT_EQ(data[3], 0x76);
1766+}
1767+ 
1768+TEST_F(UtestOnnxParser, ParseConvertData_Int4_Int32Data_OddCount) {
1769+ OnnxConstantParser constant_parser;
1770+ ge::onnx::TensorProto tensor_proto;
1771+ tensor_proto.set_data_type(OnnxDataType::INT4);
1772+ // 2 int32s = 16 packed int4 slots, but count=9 (7 padding slots)
1773+ tensor_proto.add_int32_data(0x76543210);
1774+ tensor_proto.add_int32_data(0xFEDCBA98);
1775+ ge::Tensor tensor;
1776+ TensorDesc tensor_desc = tensor.GetTensorDesc();
1777+ tensor_desc.SetDataType(ge::DataType::DT_INT4);
1778+ tensor.SetTensorDesc(tensor_desc);
1779+ int count = 9;
1780+ Status ret = constant_parser.ParseConvertData(tensor_proto, tensor, count);
1781+ EXPECT_EQ(ret, SUCCESS);
1782+ // GE expects ceil(9*4/8) = 5 bytes, not 8 (2 int32s)
1783+ EXPECT_EQ(tensor.GetSize(), 5U);
1784+ auto data = tensor.GetData();
1785+ ASSERT_NE(data, nullptr);
1786+ // first 4 bytes from int32[0] (little-endian)
1787+ EXPECT_EQ(data[0], 0x10);
1788+ EXPECT_EQ(data[1], 0x32);
1789+ EXPECT_EQ(data[2], 0x54);
1790+ EXPECT_EQ(data[3], 0x76);
1791+ // 5th byte = low nibble of int32[1] = 0x98
1792+ EXPECT_EQ(data[4], 0x98);
1793+}
1794+ 
1795+TEST_F(UtestOnnxParser, ParseConvertData_Int4_Int32Data_InsufficientData) {
1796+ OnnxConstantParser constant_parser;
1797+ ge::onnx::TensorProto tensor_proto;
1798+ tensor_proto.set_data_type(OnnxDataType::INT4);
1799+ // 1 int32 = 4 bytes available, but count=16 expects 8 bytes
1800+ tensor_proto.add_int32_data(0x76543210);
1801+ ge::Tensor tensor;
1802+ TensorDesc tensor_desc = tensor.GetTensorDesc();
1803+ tensor_desc.SetDataType(ge::DataType::DT_INT4);
1804+ tensor.SetTensorDesc(tensor_desc);
1805+ int count = 16;
1806+ Status ret = constant_parser.ParseConvertData(tensor_proto, tensor, count);
1807+ EXPECT_EQ(ret, SUCCESS);
1808+ // copy_size = min(8, 4) = 4 bytes, no over-read
1809+ EXPECT_EQ(tensor.GetSize(), 4U);
1810+ auto data = tensor.GetData();
1811+ ASSERT_NE(data, nullptr);
1812+ EXPECT_EQ(data[0], 0x10);
1813+ EXPECT_EQ(data[1], 0x32);
1814+ EXPECT_EQ(data[2], 0x54);
1815+ EXPECT_EQ(data[3], 0x76);
1816+}
1713} // namespace ge1817} // namespace ge