update code

dlinfer/graph/dicp/vendor/AtbGraph/codegen/atb_op.py

@@ -727,7 +727,7 @@ def ReduceSum(name, x, dim):
         param = infer_param.ReduceParam()
         param.name = name
         param.reduceType = infer_param.ReduceType.REDUCE_SUM
-        param.axis = dim if isinstance(dim, list) else [dim]
+        param.axis = dim
 
         op.set_input([x])
         op.set_param(param)
@@ -739,7 +739,7 @@ def ReduceMax(name, x, dim):
         param = infer_param.ReduceParam()
         param.name = name
         param.reduceType = infer_param.ReduceType.REDUCE_MAX
-        param.axis = dim if isinstance(dim, list) else [dim]
+        param.axis = dim
 
         op.set_input([x])
         op.set_param(param)
@@ -751,7 +751,7 @@ def ReduceMin(name, x, dim):
         param = infer_param.ReduceParam()
         param.name = name
         param.reduceType = infer_param.ReduceType.REDUCE_MIN
-        param.axis = dim if isinstance(dim, list) else [dim]
+        param.axis = dim
 
         op.set_input([x])
         op.set_param(param)

dlinfer/graph/dicp/vendor/AtbGraph/codegen/utils.py

@@ -163,5 +163,12 @@ def get_ascend_dtype(dtype: torch.dtype) -> str:
         raise RuntimeError(f"unknow torch data type ({dtype}) in get_ascend_dtype!")
 
 
+def get_reduce_dim(x, dim):
+    x_rank = len(x.node.meta["val"].shape)
+    dim = dim if isinstance(dim, list) else [dim]
+    dim = [(i + x_rank) % x_rank for i in dim]
+    return dim
+
+
 def remove_duplicates(lst):
     return list(OrderedDict.fromkeys(lst))

dlinfer/graph/dicp/vendor/AtbGraph/conversion.py

@@ -10,7 +10,10 @@
 from dlinfer.graph.dicp.dynamo_bridge.conversion import register_conversion_impl
 from dlinfer.graph.dicp.dynamo_bridge.op_transformer import SingleOpTransformer
 from dlinfer.graph.dicp.vendor.AtbGraph import ext_ops
-from dlinfer.graph.dicp.vendor.AtbGraph.codegen.utils import get_ascend_dtype
+from dlinfer.graph.dicp.vendor.AtbGraph.codegen.utils import (
+    get_ascend_dtype,
+    get_reduce_dim,
+)
 
 
 aten = torch.ops.aten
@@ -652,14 +655,17 @@ def aten_scalar_tensor(self, x, dtype, layout, device):
 
     @register_conversion(torch.ops.aten.sum.dim_IntList)
     def aten_reduce_sum(self, x, dim):
+        dim = get_reduce_dim(x, dim)
         return self.get_proxy(atb_op.ReduceSum, (x, dim))
 
     @register_conversion(torch.ops.aten.amax.default)
     def aten_reduce_sum(self, x, dim):
+        dim = get_reduce_dim(x, dim)
         return self.get_proxy(atb_op.ReduceMax, (x, dim))
 
     @register_conversion(torch.ops.aten.amin.default)
     def aten_reduce_sum(self, x, dim):
+        dim = get_reduce_dim(x, dim)
         return self.get_proxy(atb_op.ReduceMin, (x, dim))