mindspore 2.3.0__cp310-cp310-win_amd64.whl → 2.4.0__cp310-cp310-win_amd64.whl

mindspore/_extends/pijit/pijit_func_white_list.py

@@ -27,13 +27,129 @@ from mindspore.ops.composite.base import GradOperation, _Grad
 from mindspore.ops._primitive_cache import _get_cache_prim
 from mindspore.common.api import jit
 from mindspore.common.tensor import Tensor
-from mindspore.common._register_for_tensor import Registry, tensor_operator_registry
-from mindspore._c_expression import MetaFuncGraph_, function_id, Primitive_, PrimitiveFunction_
+from mindspore.common._register_for_tensor import Registry
+from mindspore._c_expression import MetaFuncGraph_, function_id
 from mindspore._c_expression import Tensor as Tensor_
 from mindspore._extends.parse.resources import convert_object_map
 from mindspore import _checkparam as validator
 from mindspore import Parameter, ParameterTuple
 from mindspore.common.initializer import Zero
+from mindspore.ops.function import array_func
+from mindspore.ops import operations as P
+from mindspore.ops import functional as F
+from mindspore._c_expression.np_dtypes import np_version_valid
+from mindspore.common.dtype import type_size_in_bytes
+from mindspore.communication._comm_helper import _is_initialized, _get_rank_helper, _get_local_rank_helper, \
+    _get_size_helper, _get_local_size_helper, _get_world_rank_from_group_rank_helper, _get_group_ranks, \
+    _get_group_rank_from_world_rank_helper, _set_elegant_exit_handle
+from mindspore import SummaryCollector
+from mindspore.train import ModelCheckpoint, LossMonitor
+from mindspore.train.model import _FrameworkProfilerCallback
+from mindspore.train.data_sink import _init_sink_dataset
+from mindspore.train.summary import SummaryRecord
+from mindspore.train._utils import _exec_datagraph
+from mindspore.train.summary.writer import BaseWriter
+from mindspore.train.serialization import _exec_save, load, export_split_mindir, obfuscate_model, _parse_ckpt_proto, \
+    _generate_front_info_for_param_data_file, _get_data_file, _encrypt_data, _split_save, _save_mindir_together, \
+    _load_into_param_dict
+from mindspore.parallel import _cost_model_context
+from mindspore.parallel._offload_context import offload_context
+from mindspore.run_check._check_version import check_version_and_env_config
+from mindspore.dataset.callback.ds_callback import DSCallback, WaitedDSCallback
+from mindspore.dataset.transforms.c_transforms import TensorOperation as CTensorOperation, OneHot as COneHot, \
+    Fill as CFill, TypeCast as CTypeCast, Slice as CSlice, Mask as CMask, PadEnd as CPadEnd, \
+    Concatenate as CConcatenate, Duplicate as CDuplicate, Unique as CUnique, Compose as CCompose, \
+    RandomApply as CRandomApply, RandomChoice as CRandomChoice, Plugin as CPlugin
+from mindspore.dataset.transforms.transforms import TensorOperation, Compose, Concatenate, Duplicate, Fill, Mask, \
+    OneHot, PadEnd, Plugin, RandomApply, RandomChoice, Slice, TypeCast, Unique
+from mindspore.dataset.text.transforms import AddToken, JiebaTokenizer, Lookup, Ngram, SentencePieceTokenizer, \
+    SlidingWindow, ToNumber, ToVectors, Truncate, TruncateSequencePair, UnicodeCharTokenizer, WordpieceTokenizer, \
+    BasicTokenizer, BertTokenizer, CaseFold, FilterWikipediaXML, NormalizeUTF8, RegexReplace, RegexTokenizer, \
+    UnicodeScriptTokenizer, WhitespaceTokenizer
+from mindspore.dataset.core.datatypes import nptype_to_detype, mstype_to_detype, mstypelist_to_detypelist
+from mindspore.dataset.audio.utils import create_dct, linear_fbanks, melscale_fbanks
+from mindspore.dataset.audio.transforms import AllpassBiquad, AmplitudeToDB, Angle, BandBiquad, BandpassBiquad, \
+    BandrejectBiquad, BassBiquad, Biquad, ComplexNorm, ComputeDeltas, Contrast, DBToAmplitude, DCShift, DeemphBiquad, \
+    DetectPitchFrequency, Dither, EqualizerBiquad, Fade, Filtfilt, Flanger, FrequencyMasking, Gain, GriffinLim, \
+    HighpassBiquad, InverseMelScale, InverseSpectrogram, LFCC, LFilter, LowpassBiquad, Magphase, MaskAlongAxis, \
+    MaskAlongAxisIID, MelScale, MelSpectrogram, MFCC, MuLawDecoding, MuLawEncoding, Overdrive, Phaser, PhaseVocoder, \
+    PitchShift, Resample, RiaaBiquad, SlidingWindowCmn, SpectralCentroid, Spectrogram, TimeMasking, TimeStretch, \
+    TrebleBiquad, Vad, Vol
+from mindspore.dataset.engine.datasets_audio import CMUArcticDataset, GTZANDataset, LibriTTSDataset, LJSpeechDataset, \
+    SpeechCommandsDataset, TedliumDataset, YesNoDataset
+from mindspore.dataset.engine.cache_client import DatasetCache
+from mindspore.dataset.engine.iterators import Iterator
+from mindspore.dataset.engine.datasets_standard_format import CSVDataset, MindDataset, TFRecordDataset
+from mindspore.dataset.engine.datasets_text import AGNewsDataset, AmazonReviewDataset, CLUEDataset, CoNLL2000Dataset, \
+    DBpediaDataset, EnWik9Dataset, IMDBDataset, IWSLT2016Dataset, IWSLT2017Dataset, Multi30kDataset, WikiTextDataset, \
+    PennTreebankDataset, SogouNewsDataset, SQuADDataset, SST2Dataset, TextFileDataset, UDPOSDataset, \
+    YelpReviewDataset, YahooAnswersDataset
+from mindspore.dataset.engine.datasets_user_defined import GeneratorDataset
+from mindspore.dataset.engine.datasets_vision import Caltech256Dataset, CelebADataset, Cifar100Dataset, \
+    Cifar10Dataset, CityscapesDataset, CocoDataset, DIV2KDataset, EMnistDataset, FakeImageDataset, \
+    FashionMnistDataset, FlickrDataset, Food101Dataset, ImageFolderDataset, KITTIDataset, KMnistDataset, \
+    LFWDataset, LSUNDataset, ManifestDataset, MnistDataset, OmniglotDataset, PhotoTourDataset, \
+    Places365Dataset, QMnistDataset, RandomDataset, RenderedSST2Dataset, SBUDataset, SemeionDataset, \
+    STL10Dataset, SUN397Dataset, USPSDataset, VOCDataset, WIDERFaceDataset
+from mindspore.dataset.engine.queue import _SharedQueue
+from mindspore.dataset.engine.datasets import Dataset, BucketBatchByLengthDataset, BatchDataset, \
+    PaddedBatchDataset, SyncWaitDataset, ShuffleDataset, MapDataset, FilterDataset, RepeatDataset, SkipDataset, \
+    TakeDataset, ZipDataset, ConcatDataset, RenameDataset, ProjectDataset, _ToDevice, TransferDataset, Schema
+from mindspore.dataset.engine.samplers import Sampler, DistributedSampler, PKSampler, RandomSampler, SubsetSampler, \
+    SequentialSampler, SubsetRandomSampler, WeightedRandomSampler
+from mindspore.dataset.vision.c_transforms import ImageTensorOperation, AdjustGamma, AutoAugment, AutoContrast, \
+    BoundingBoxAugment, CenterCrop, ConvertColor, Crop, CutMixBatch, CutOut, Decode, Equalize, GaussianBlur, \
+    HorizontalFlip, HWC2CHW, Invert, MixUpBatch, Normalize, NormalizePad, Pad, RandomAdjustSharpness, RandomAffine, \
+    RandomAutoContrast, RandomColor, RandomColorAdjust, RandomCrop, RandomCropDecodeResize, RandomCropWithBBox, \
+    RandomEqualize, RandomHorizontalFlip, RandomHorizontalFlipWithBBox, RandomInvert, RandomLighting, RandomRotation, \
+    RandomPosterize, RandomResizedCrop, RandomResizedCropWithBBox, RandomResize, RandomResizeWithBBox, \
+    RandomSelectSubpolicy, RandomSharpness, RandomSolarize, RandomVerticalFlip, RandomVerticalFlipWithBBox, Rescale, \
+    Resize, ResizeWithBBox, RgbToBgr, Rotate, SlicePatches, UniformAugment, VerticalFlip
+from mindspore.dataset.vision.utils import encode_jpeg, encode_png, get_image_num_channels, get_image_size, \
+    read_file, read_image, read_video, read_video_timestamps, write_file, write_jpeg, write_png
+from mindspore.dataset.vision.transforms import AdjustBrightness, AdjustContrast, AdjustGamma as VAdjustGamma, \
+    AdjustHue, AdjustSaturation, AdjustSharpness, Affine, AutoAugment as VAutoAugment, AutoContrast as VAutoContrast, \
+    BoundingBoxAugment as VBoundingBoxAugment, CenterCrop as VCenterCrop, ConvertColor as VConvertColor, \
+    Crop as VCrop, CutMixBatch as VCutMixBatch, CutOut as VCutOut, Decode as VDecode, DecodeVideo, \
+    Equalize as VEqualize, Erase, GaussianBlur as VGaussianBlur, HorizontalFlip as VHorizontalFlip, \
+    HWC2CHW as VHWC2CHW, Invert as VInvert, MixUpBatch as VMixUpBatch, Normalize as VNormalize, \
+    NormalizePad as VNormalizePad, Pad as VPad, PadToSize, Perspective, Posterize, RandAugment, \
+    RandomAdjustSharpness as VRandomAdjustSharpness, RandomAffine as VRandomAffine, \
+    RandomAutoContrast as VRandomAutoContrast, RandomColor as VRandomColor, RandomColorAdjust as VRandomColorAdjust, \
+    RandomCrop as VRandomCrop, RandomCropDecodeResize as VRandomCropDecodeResize, \
+    RandomCropWithBBox as VRandomCropWithBBox, RandomEqualize as VRandomEqualize, RandomResize as VRandomResize, \
+    RandomHorizontalFlip as VRandomHorizontalFlip, RandomHorizontalFlipWithBBox as VRandomHorizontalFlipWithBBox, \
+    RandomInvert as VRandomInvert, RandomLighting as VRandomLighting, RandomPosterize as VRandomPosterize, \
+    RandomResizedCrop as VRandomResizedCrop, RandomResizedCropWithBBox as VRandomResizedCropWithBBox, \
+    RandomResizeWithBBox as VRandomResizeWithBBox, RandomRotation as VRandomRotation, \
+    RandomSelectSubpolicy as VRandomSelectSubpolicy, RandomSharpness as VRandomSharpness, \
+    RandomSolarize as VRandomSolarize, RandomVerticalFlip as VRandomVerticalFlip, \
+    RandomVerticalFlipWithBBox as VRandomVerticalFlipWithBBox, Rescale as VRescale, Resize as VResize, ResizedCrop, \
+    ResizeWithBBox as VResizeWithBBox, Rotate as VRotate, SlicePatches as VSlicePatches, Solarize, ToTensor,\
+    TrivialAugmentWide, UniformAugment as VUniformAugment, VerticalFlip as VVerticalFlip
+from mindspore.profiler.profiling import Profiler
+from mindspore.communication._hccl_management import get_rank_size, get_rank_id
+from mindspore.communication._comm_helper import _create_group_helper, _destroy_group_helper
+from mindspore.communication.management import _set_rank_from_mpi, init as cinit, release as crelease
+from mindspore.hal.stream import Stream, synchronize, set_cur_stream, current_stream, default_stream
+from mindspore.hal.event import Event
+from mindspore.hal.memory import memory_stats, memory_reserved, max_memory_allocated, reset_peak_memory_stats, \
+    memory_summary, memory_allocated, max_memory_reserved, reset_max_memory_allocated, reset_max_memory_reserved
+from mindspore.multiprocessing import Process
+from mindspore.mindrecord.shardsegment import ShardSegment
+from mindspore.mindrecord.shardreader import ShardReader
+from mindspore.mindrecord.shardindexgenerator import ShardIndexGenerator
+from mindspore.mindrecord.shardwriter import ShardWriter
+from mindspore.mindrecord.shardheader import ShardHeader
+from mindspore.mindrecord.config import encrypt, decrypt
+from mindspore.parallel.mpi._mpi_config import _MpiConfig
+from mindspore.parallel._ps_context import ps_context
+from mindspore.parallel.algo_parameter_config import _AlgoParameterConfig
+from mindspore.parallel._utils import _reset_op_id, _reset_op_id_with_offset
+from mindspore.parallel._recovery_context import recovery_context
+from mindspore.parallel._auto_parallel_context import _AutoParallelContext
+from mindspore.common.api import ms_memory_recycle
+from mindspore.context import _Context
 
 
 def _get_after_grad_code():
@@ -51,6 +167,107 @@ def _get_after_grad_code():
     return codes
 
 
+def _get_dataset_forbidden_code():
+    """Get the forbidden function which should be broken in graph"""
+    codes = []
+    codes.extend([DSCallback.__init__, WaitedDSCallback.__init__])
+    codes.extend([CTensorOperation.__call__, COneHot.parse, CFill.__init__, CFill.parse, CTypeCast.parse, \
+                  CSlice.parse, CMask.__init__, CMask.parse, CPadEnd.__init__, CPadEnd.parse, CConcatenate.__init__, \
+                  CConcatenate.parse, CDuplicate.parse, CUnique.parse, CCompose.parse, CRandomApply.parse, \
+                  CRandomChoice.parse, CPlugin.parse])
+    codes.extend([TensorOperation.__call__, Compose.parse, Concatenate.__init__, Concatenate.parse, Duplicate.parse, \
+                  Fill.__init__, Fill.parse, Mask.__init__, Mask.parse, OneHot.parse, PadEnd.__init__, PadEnd.parse, \
+                  Plugin.parse, RandomApply.parse, RandomChoice.parse, Slice.parse, TypeCast.parse, Unique.parse])
+    codes.extend([AddToken.parse, JiebaTokenizer.parse, Lookup.parse, Ngram.parse, SentencePieceTokenizer.parse, \
+                  SlidingWindow.parse, ToNumber.parse, ToVectors.parse, Truncate.parse, TruncateSequencePair.parse, \
+                  UnicodeCharTokenizer.parse, WordpieceTokenizer.parse, BasicTokenizer.parse, BertTokenizer.parse, \
+                  CaseFold.parse, FilterWikipediaXML.parse, NormalizeUTF8.parse, RegexReplace.parse, \
+                  RegexTokenizer.parse, UnicodeScriptTokenizer.parse, WhitespaceTokenizer.parse])
+    codes.extend([create_dct, linear_fbanks, melscale_fbanks])
+    codes.extend([AllpassBiquad.parse, AmplitudeToDB.parse, Angle.parse, BandBiquad.parse, BandpassBiquad.parse, \
+                  BandrejectBiquad.parse, BassBiquad.parse, Biquad.parse, ComplexNorm.parse, ComputeDeltas.parse, \
+                  Contrast.parse, DBToAmplitude.parse, DCShift.parse, DeemphBiquad.parse, DetectPitchFrequency.parse, \
+                  Dither.parse, EqualizerBiquad.parse, Fade.parse, Filtfilt.parse, Flanger.parse, \
+                  FrequencyMasking.parse, Gain.parse, GriffinLim.parse, HighpassBiquad.parse, InverseMelScale.parse, \
+                  InverseSpectrogram.parse, LFCC.parse, LFilter.parse, LowpassBiquad.parse, Magphase.parse, \
+                  MaskAlongAxis.parse, MaskAlongAxisIID.parse, MelScale.parse, MelSpectrogram.parse, MFCC.parse, \
+                  MuLawDecoding.parse, MuLawEncoding.parse, Overdrive.parse, Phaser.parse, PhaseVocoder.parse, \
+                  PhaseVocoder.__init__, PitchShift.parse, Resample.parse, RiaaBiquad.parse, SlidingWindowCmn.parse, \
+                  SpectralCentroid.parse, Spectrogram.parse, TimeMasking.parse, TimeStretch.parse, \
+                  TrebleBiquad.parse, Vad.parse, Vol.parse])
+    codes.extend([CMUArcticDataset.parse, GTZANDataset.parse, LibriTTSDataset.parse, LJSpeechDataset.parse, \
+                  SpeechCommandsDataset.parse, TedliumDataset.parse, YesNoDataset.parse])
+    codes.extend([DatasetCache.__init__, Iterator.__init__])
+    codes.extend([CSVDataset.parse, MindDataset.parse, TFRecordDataset.parse])
+    codes.extend([AGNewsDataset.parse, AmazonReviewDataset.parse, CLUEDataset.parse, CoNLL2000Dataset.parse, \
+                  DBpediaDataset.parse, EnWik9Dataset.parse, IMDBDataset.parse, IWSLT2016Dataset.parse, \
+                  IWSLT2017Dataset.parse, Multi30kDataset.parse, PennTreebankDataset.parse, SogouNewsDataset.parse, \
+                  SQuADDataset.parse, SST2Dataset.parse, TextFileDataset.parse, UDPOSDataset.parse, \
+                  WikiTextDataset.parse, YahooAnswersDataset.parse, YelpReviewDataset.parse])
+    codes.extend([GeneratorDataset.parse, Caltech256Dataset.parse, CelebADataset.parse, Cifar10Dataset.parse, \
+                  Cifar100Dataset.parse, CityscapesDataset.parse, CocoDataset.parse, DIV2KDataset.parse, \
+                  EMnistDataset.parse, FakeImageDataset.parse, FashionMnistDataset.parse, FlickrDataset.parse, \
+                  Food101Dataset.parse, ImageFolderDataset.parse, KITTIDataset.parse, KMnistDataset.parse, \
+                  LFWDataset.parse, LSUNDataset.parse, ManifestDataset.parse, MnistDataset.parse, VOCDataset.parse, \
+                  OmniglotDataset.parse, PhotoTourDataset.parse, Places365Dataset.parse, QMnistDataset.parse, \
+                  RandomDataset.parse, RenderedSST2Dataset.parse, SBUDataset.parse, SemeionDataset.parse, \
+                  STL10Dataset.parse, SUN397Dataset.parse, USPSDataset.parse, WIDERFaceDataset.parse])
+    codes.extend([_SharedQueue.put, _SharedQueue.get])
+    codes.extend([BucketBatchByLengthDataset.parse, BatchDataset.parse, \
+                  PaddedBatchDataset.parse, SyncWaitDataset.parse, ShuffleDataset.parse, MapDataset.parse, \
+                  FilterDataset.parse, RepeatDataset.parse, SkipDataset.parse, TakeDataset.parse, ZipDataset.parse, \
+                  ConcatDataset.parse, RenameDataset.parse, ProjectDataset.parse, _ToDevice.__init__, \
+                  TransferDataset.parse, Schema.__init__, Schema.add_column, Dataset.save])
+    codes.extend([Sampler.parse, DistributedSampler.parse, DistributedSampler.parse_for_minddataset, PKSampler.parse, \
+                  PKSampler.parse_for_minddataset, RandomSampler.parse, RandomSampler.parse_for_minddataset, \
+                  SequentialSampler.parse, SequentialSampler.parse_for_minddataset, SubsetSampler.parse, \
+                  SubsetSampler.parse_for_minddataset, SubsetRandomSampler.parse, \
+                  SubsetRandomSampler.parse_for_minddataset, WeightedRandomSampler.parse])
+    codes.extend([ImageTensorOperation.__call__, AdjustGamma.parse, AutoAugment.parse, AutoContrast.parse, Pad.parse, \
+                  BoundingBoxAugment.parse, CenterCrop.parse, ConvertColor.parse, Crop.parse, CutMixBatch.parse, \
+                  CutOut.parse, Decode.parse, Equalize.parse, GaussianBlur.parse, HorizontalFlip.parse, \
+                  Invert.parse, MixUpBatch.parse, Normalize.parse, NormalizePad.parse, HWC2CHW.parse, \
+                  RandomAdjustSharpness.parse, RandomAffine.parse, RandomAutoContrast.parse, RandomColor.parse, \
+                  RandomColorAdjust.parse, RandomCrop.parse, RandomCropDecodeResize.parse, RandomCropWithBBox.parse, \
+                  RandomEqualize.parse, RandomHorizontalFlip.parse, RandomHorizontalFlipWithBBox.parse, \
+                  RandomInvert.parse, RandomLighting.parse, RandomPosterize.parse, RandomResizedCrop.parse, \
+                  RandomResizedCropWithBBox.parse, RandomResize.parse, RandomResizeWithBBox.parse, Resize.parse, \
+                  RandomRotation.parse, RandomSelectSubpolicy.parse, RandomSharpness.parse, RandomSolarize.parse, \
+                  RandomVerticalFlip.parse, RandomVerticalFlipWithBBox.parse, Rescale.parse, VerticalFlip.parse, \
+                  ResizeWithBBox.parse, RgbToBgr.parse, Rotate.parse, SlicePatches.parse, UniformAugment.parse])
+    codes.extend([encode_jpeg, encode_png, get_image_num_channels, get_image_size, read_file, read_image, read_video, \
+                  read_video_timestamps, write_file, write_jpeg, write_png])
+    codes.extend([AdjustBrightness.parse, AdjustContrast.parse, VAdjustGamma.parse, AdjustHue.parse, \
+                  AdjustSaturation.parse, AdjustSharpness.parse, Affine.parse, VAutoAugment.parse, \
+                  VAutoContrast.parse, VBoundingBoxAugment.parse, VCenterCrop.parse, VConvertColor.parse, \
+                  VCrop.parse, VCutMixBatch.parse, VCutOut.parse, VDecode.parse, DecodeVideo.parse, \
+                  VEqualize.parse, Erase.parse, VGaussianBlur.parse, VHorizontalFlip.parse, \
+                  VHWC2CHW.parse, VInvert.parse, VMixUpBatch.parse, VNormalize.parse, VNormalizePad.parse, \
+                  VPad.parse, PadToSize.parse, Perspective.parse, Posterize.parse, RandAugment.parse, \
+                  VRandomAdjustSharpness.parse, VRandomAffine.parse, \
+                  VRandomAutoContrast.parse, VRandomColor.parse, VRandomColorAdjust.parse, \
+                  VRandomCrop.parse, VRandomCropDecodeResize.parse, \
+                  VRandomCropWithBBox.parse, VRandomEqualize.parse, VRandomResize.parse, \
+                  VRandomHorizontalFlip.parse, VRandomHorizontalFlipWithBBox.parse, \
+                  VRandomInvert.parse, VRandomLighting.parse, VRandomPosterize.parse, \
+                  VRandomResizedCrop.parse, VRandomResizedCropWithBBox.parse, \
+                  VRandomResizeWithBBox.parse, VRandomRotation.parse, \
+                  VRandomSelectSubpolicy.parse, VRandomSharpness.parse, \
+                  VRandomSolarize.parse, VRandomVerticalFlip.parse, \
+                  VRandomVerticalFlipWithBBox.parse, VRescale.parse, VResize.parse, ResizedCrop.parse, \
+                  VResizeWithBBox.parse, VRotate.parse, VSlicePatches.parse, Solarize.parse, ToTensor.parse,\
+                  TrivialAugmentWide.parse, VUniformAugment.parse, VVerticalFlip])
+    return codes
+
+
+def _get_math_code():
+    """Get the math builtin function which should be guarded in graph"""
+    codes = []
+    for i in dir(math):
+        codes.append(getattr(math, i))
+    return codes
+
+
 def _get_psjit_code():
     """Get the code object of 'staging_specialize'"""
     @jit
@@ -95,30 +312,13 @@ def _get_pijit_constexpr_code():
     return codes
 
 
-def _get_ms_api():
-    """Get ms api"""
-    target_types = Cell, types.FunctionType, Primitive_, PrimitiveFunction_
-    results = []
-    from mindspore.ops import operations as P
-    from mindspore.ops import functional as F
-    from mindspore.ops import composite as C
-    mods = P, F, C
-    for mod in mods:
-        for i in mod.__all__:
-            f = getattr(mod, i)
-            if isinstance(f, target_types):
-                results.append(f)
-    for f in tensor_operator_registry.values():
-        if isinstance(f, target_types):
-            results.append(f)
-    return results
-
-
 psjit_code = _get_psjit_code()
 constexpr_code = _get_constexpr_code()
 primexpr_code = _get_primexpr_code()
 
 primitive_key = id(Primitive.__call__)
+primitive_assign_key = id(P.Assign.__call__)
+
 constexpr_key = id(constexpr_code)
 primexpr_key = id(primexpr_code)
 meta_func_graph_key = id(MetaFuncGraph_)
@@ -153,6 +353,13 @@ FUNC_KEY_PSJIT_CONVERTMAP = 15 # "mindspore._extends.parse.resources.convert_obj
 FUNC_KEY_GRAPH_CELL = 16  # "mindspore.nn.cell.GraphCell"
 FUNC_KEY_MS_API = 17  # mindspore common api
 FUNC_KEY_MAPPING_GET = 18 # collections.abc.Mapping.get
+FUNC_KEY_LIST_POP = 19  # list.pop
+FUNC_KEY_LIST_REMOVE = 20  # list.remove
+FUNC_KEY_LIST_REVERSE = 21  # list.reverse
+FUNC_KEY_DICT_ITEMS = 22  # dict.items
+FUNC_KEY_PRIMITIVE_ASSIGN = 23  # mindspore.ops.assign, Primitive("Assign")
+FUNC_KEY_TENSOR_SETITEM = 24  # Tensor.__setitem__
+FUNC_KEY_TENSOR_ASSIGN_VALUE = 25  # Tensor.assign_value
 
 # Initialized only once. This map will initialize by c++ when start pijit.
 # key is customer if fuzzy match. (Primitive, constexpr, primexpr, MetaFuncGraph)
@@ -162,6 +369,8 @@ FUNC_KEY_MAPPING_GET = 18 # collections.abc.Mapping.get
 _func_map = {
     # special function
     pijit_constexpr_key: FUNC_KEY_PIJIT_CONSTEXPR,
+    id(getattr(array_func, "_get_max_type")): FUNC_KEY_PIJIT_CONSTEXPR,
+    id(Cell.__getattr__): FUNC_KEY_PIJIT_CONSTEXPR,
     pijit_forbidden_key: FUNC_KEY_PIJIT_FORBIDDEN,
     primitive_key: FUNC_KEY_PRIMITIVE,
     constexpr_key: FUNC_KEY_CONSTEXPR,
@@ -172,6 +381,12 @@ _func_map = {
     id(_get_cache_prim): FUNC_KEY_GET_CACHE_PRIM,
     id(Registry.get): FUNC_KEY_REGISTRY_GET,
 
+    # tensor side-effect
+    primitive_assign_key: FUNC_KEY_PRIMITIVE_ASSIGN,
+    id(F.assign): FUNC_KEY_PRIMITIVE_ASSIGN,
+    id(Tensor.assign_value): FUNC_KEY_TENSOR_ASSIGN_VALUE,
+    id(Tensor.__setitem__): FUNC_KEY_TENSOR_SETITEM,
+
     # Tensor method
     id(Tensor.astype): FUNC_KEY_TENSOR_ASTYPE,
 
@@ -181,6 +396,7 @@ _func_map = {
     function_id(len): FUNC_KEY_BUILTIN_FUNC,
     function_id(abs): FUNC_KEY_BUILTIN_FUNC,
     function_id(max): FUNC_KEY_BUILTIN_FUNC,
+    function_id(min): FUNC_KEY_BUILTIN_FUNC,
     function_id(all): FUNC_KEY_BUILTIN_FUNC,
     function_id(any): FUNC_KEY_BUILTIN_FUNC,
     function_id(hash): FUNC_KEY_BUILTIN_FUNC,
@@ -189,6 +405,10 @@ _func_map = {
     function_id(callable): FUNC_KEY_BUILTIN_FUNC,
     function_id(getattr): FUNC_KEY_BUILTIN_FUNC,
     function_id(hasattr): FUNC_KEY_BUILTIN_FUNC,
+    function_id(chr): FUNC_KEY_BUILTIN_FUNC,
+    function_id(divmod): FUNC_KEY_BUILTIN_FUNC,
+    function_id(repr): FUNC_KEY_BUILTIN_FUNC,
+    function_id(type): FUNC_KEY_BUILTIN_FUNC,
 
     # types.MethodDescriptorType, types.WrapperDescriptorType
     function_id(tuple.__getitem__): FUNC_KEY_BUILTIN_FUNC,
@@ -232,7 +452,11 @@ _func_map = {
     function_id(str.format_map): FUNC_KEY_BUILTIN_FUNC,
     function_id(str.__format__): FUNC_KEY_BUILTIN_FUNC,
     function_id(list.append): FUNC_KEY_LIST_APPEND,
+    function_id(list.pop): FUNC_KEY_LIST_POP,
+    function_id(list.remove): FUNC_KEY_LIST_REMOVE,
+    function_id(list.reverse): FUNC_KEY_LIST_REVERSE,
     function_id(dict.pop): FUNC_KEY_DICT_POP,
+    function_id(dict.items): FUNC_KEY_DICT_ITEMS,
 
     # instancemethod
     function_id(Tensor_._flatten_tensors): FUNC_KEY_BUILTIN_FUNC,  # pylint: disable=protected-access
@@ -254,23 +478,125 @@ _func_map = {
 
     # other builtin function
     function_id(collections.abc.Mapping.get): FUNC_KEY_MAPPING_GET,
-    function_id(math.log): FUNC_KEY_BUILTIN_FUNC,
-
     function_id(numpy.isinf): FUNC_KEY_BUILTIN_FUNC,
     function_id(numpy.isnan): FUNC_KEY_BUILTIN_FUNC,
     function_id(numpy.abs): FUNC_KEY_BUILTIN_FUNC,
     function_id(numpy.log): FUNC_KEY_BUILTIN_FUNC,
+    function_id(os.getenv): FUNC_KEY_BUILTIN_FUNC,
 
     # const function
-    function_id(os.getenv): FUNC_KEY_PIJIT_CONSTEXPR,
     function_id(validator.check_number_range): FUNC_KEY_PIJIT_CONSTEXPR,
     function_id(validator.check_is_int): FUNC_KEY_PIJIT_CONSTEXPR,
     function_id(validator.check_is_number): FUNC_KEY_PIJIT_CONSTEXPR,
+    function_id(np_version_valid): FUNC_KEY_PIJIT_CONSTEXPR,
+    function_id(_is_initialized): FUNC_KEY_PIJIT_CONSTEXPR,
+    function_id(_set_elegant_exit_handle): FUNC_KEY_PIJIT_CONSTEXPR,
+    function_id(_cost_model_context.get_cost_model_context): FUNC_KEY_PIJIT_CONSTEXPR,
+    function_id(Stream.__repr__): FUNC_KEY_PIJIT_CONSTEXPR,
+    function_id(get_rank_size): FUNC_KEY_PIJIT_CONSTEXPR,
+    function_id(get_rank_id): FUNC_KEY_PIJIT_CONSTEXPR,
+    function_id(offload_context): FUNC_KEY_PIJIT_CONSTEXPR,
+    function_id(check_version_and_env_config): FUNC_KEY_PIJIT_CONSTEXPR,
+
+    # inner function
+    function_id(type_size_in_bytes): FUNC_KEY_BUILTIN_FUNC,
+    function_id(_get_rank_helper): FUNC_KEY_BUILTIN_FUNC,
+    function_id(_get_local_rank_helper): FUNC_KEY_BUILTIN_FUNC,
+    function_id(_get_size_helper): FUNC_KEY_BUILTIN_FUNC,
+    function_id(_get_local_size_helper): FUNC_KEY_BUILTIN_FUNC,
+    function_id(_get_world_rank_from_group_rank_helper): FUNC_KEY_BUILTIN_FUNC,
+    function_id(_get_group_ranks): FUNC_KEY_BUILTIN_FUNC,
+    function_id(_get_group_rank_from_world_rank_helper): FUNC_KEY_BUILTIN_FUNC,
+    function_id(nptype_to_detype): FUNC_KEY_BUILTIN_FUNC,
+    function_id(mstype_to_detype): FUNC_KEY_BUILTIN_FUNC,
+    function_id(mstypelist_to_detypelist): FUNC_KEY_BUILTIN_FUNC,
+
+    # no need to capture function in black list
+    function_id(SummaryCollector.__init__): FUNC_KEY_PIJIT_FORBIDDEN,
+    function_id(SummaryCollector.begin): FUNC_KEY_PIJIT_FORBIDDEN,
+    function_id(SummaryCollector.step_end): FUNC_KEY_PIJIT_FORBIDDEN,
+    function_id(SummaryCollector.epoch_end): FUNC_KEY_PIJIT_FORBIDDEN,
+    function_id(SummaryCollector.end): FUNC_KEY_PIJIT_FORBIDDEN,
+    function_id(ModelCheckpoint.step_end): FUNC_KEY_PIJIT_FORBIDDEN,
+    function_id(ModelCheckpoint.end): FUNC_KEY_PIJIT_FORBIDDEN,
+    function_id(LossMonitor.step_end): FUNC_KEY_PIJIT_FORBIDDEN,
+    function_id(LossMonitor.on_train_epoch_end): FUNC_KEY_PIJIT_FORBIDDEN,
+    function_id(SummaryRecord.__init__): FUNC_KEY_PIJIT_FORBIDDEN,
+    function_id(_exec_datagraph): FUNC_KEY_PIJIT_FORBIDDEN,
+    function_id(BaseWriter.writer): FUNC_KEY_PIJIT_FORBIDDEN,
+    function_id(_FrameworkProfilerCallback.step_begin): FUNC_KEY_PIJIT_FORBIDDEN,
+    function_id(_FrameworkProfilerCallback.step_end): FUNC_KEY_PIJIT_FORBIDDEN,
+    function_id(_init_sink_dataset): FUNC_KEY_PIJIT_FORBIDDEN,
+    function_id(_exec_save): FUNC_KEY_PIJIT_FORBIDDEN,
+    function_id(load): FUNC_KEY_PIJIT_FORBIDDEN,
+    function_id(export_split_mindir): FUNC_KEY_PIJIT_FORBIDDEN,
+    function_id(obfuscate_model): FUNC_KEY_PIJIT_FORBIDDEN,
+    function_id(_parse_ckpt_proto): FUNC_KEY_PIJIT_FORBIDDEN,
+    function_id(_generate_front_info_for_param_data_file): FUNC_KEY_PIJIT_FORBIDDEN,
+    function_id(_get_data_file): FUNC_KEY_PIJIT_FORBIDDEN,
+    function_id(_encrypt_data): FUNC_KEY_PIJIT_FORBIDDEN,
+    function_id(_split_save): FUNC_KEY_PIJIT_FORBIDDEN,
+    function_id(_save_mindir_together): FUNC_KEY_PIJIT_FORBIDDEN,
+    function_id(_load_into_param_dict): FUNC_KEY_PIJIT_FORBIDDEN,
+    function_id(Profiler.start): FUNC_KEY_PIJIT_FORBIDDEN,
+    function_id(_create_group_helper): FUNC_KEY_PIJIT_FORBIDDEN,
+    function_id(_destroy_group_helper): FUNC_KEY_PIJIT_FORBIDDEN,
+    function_id(_set_rank_from_mpi): FUNC_KEY_PIJIT_FORBIDDEN,
+    function_id(cinit): FUNC_KEY_PIJIT_FORBIDDEN,
+    function_id(crelease): FUNC_KEY_PIJIT_FORBIDDEN,
+    function_id(Stream.__init__): FUNC_KEY_PIJIT_FORBIDDEN,
+    function_id(Stream.synchronize): FUNC_KEY_PIJIT_FORBIDDEN,
+    function_id(Stream.query): FUNC_KEY_PIJIT_FORBIDDEN,
+    function_id(Stream.__eq__): FUNC_KEY_PIJIT_FORBIDDEN,
+    function_id(synchronize): FUNC_KEY_PIJIT_FORBIDDEN,
+    function_id(set_cur_stream): FUNC_KEY_PIJIT_FORBIDDEN,
+    function_id(current_stream): FUNC_KEY_PIJIT_FORBIDDEN,
+    function_id(default_stream): FUNC_KEY_PIJIT_FORBIDDEN,
+    function_id(Event.__init__): FUNC_KEY_PIJIT_FORBIDDEN,
+    function_id(Event.record): FUNC_KEY_PIJIT_FORBIDDEN,
+    function_id(Event.wait): FUNC_KEY_PIJIT_FORBIDDEN,
+    function_id(Event.synchronize): FUNC_KEY_PIJIT_FORBIDDEN,
+    function_id(Event.query): FUNC_KEY_PIJIT_FORBIDDEN,
+    function_id(Event.elapsed_time): FUNC_KEY_PIJIT_FORBIDDEN,
+    function_id(memory_stats): FUNC_KEY_PIJIT_FORBIDDEN,
+    function_id(memory_reserved): FUNC_KEY_PIJIT_FORBIDDEN,
+    function_id(max_memory_reserved): FUNC_KEY_PIJIT_FORBIDDEN,
+    function_id(reset_peak_memory_stats): FUNC_KEY_PIJIT_FORBIDDEN,
+    function_id(reset_peak_memory_stats): FUNC_KEY_PIJIT_FORBIDDEN,
+    function_id(memory_summary): FUNC_KEY_PIJIT_FORBIDDEN,
+    function_id(memory_allocated): FUNC_KEY_PIJIT_FORBIDDEN,
+    function_id(max_memory_allocated): FUNC_KEY_PIJIT_FORBIDDEN,
+    function_id(reset_max_memory_reserved): FUNC_KEY_PIJIT_FORBIDDEN,
+    function_id(reset_max_memory_allocated): FUNC_KEY_PIJIT_FORBIDDEN,
+    function_id(Process.run): FUNC_KEY_PIJIT_FORBIDDEN,
+    function_id(Process.start): FUNC_KEY_PIJIT_FORBIDDEN,
+    function_id(ShardSegment.__init__): FUNC_KEY_PIJIT_FORBIDDEN,
+    function_id(ShardReader.__init__): FUNC_KEY_PIJIT_FORBIDDEN,
+    function_id(ShardIndexGenerator.__init__): FUNC_KEY_PIJIT_FORBIDDEN,
+    function_id(ShardWriter.__init__): FUNC_KEY_PIJIT_FORBIDDEN,
+    function_id(ShardHeader.__init__): FUNC_KEY_PIJIT_FORBIDDEN,
+    function_id(encrypt): FUNC_KEY_PIJIT_FORBIDDEN,
+    function_id(decrypt): FUNC_KEY_PIJIT_FORBIDDEN,
+    function_id(_MpiConfig.__init__): FUNC_KEY_PIJIT_FORBIDDEN,
+    function_id(ps_context): FUNC_KEY_PIJIT_FORBIDDEN,
+    function_id(_AlgoParameterConfig.__init__): FUNC_KEY_PIJIT_FORBIDDEN,
+    function_id(_reset_op_id): FUNC_KEY_PIJIT_FORBIDDEN,
+    function_id(_reset_op_id_with_offset): FUNC_KEY_PIJIT_FORBIDDEN,
+    function_id(recovery_context): FUNC_KEY_PIJIT_FORBIDDEN,
+    function_id(_AutoParallelContext.__init__): FUNC_KEY_PIJIT_FORBIDDEN,
+    function_id(ms_memory_recycle): FUNC_KEY_PIJIT_FORBIDDEN,
+    function_id(_Context.__init__): FUNC_KEY_PIJIT_FORBIDDEN,
 }
 
 for after_grad in _get_after_grad_code():
     _func_map[id(after_grad)] = FUNC_KEY_GRAD_OPERATIONS_CODE
 
+for func in _get_dataset_forbidden_code():
+    _func_map[function_id(func)] = FUNC_KEY_PIJIT_FORBIDDEN
+
+for math_func in _get_math_code():
+    _func_map[function_id(math_func)] = FUNC_KEY_BUILTIN_FUNC
+
 for k, v in convert_object_map.items():
     key = id(k)
     if key not in _func_map and isinstance(v, Primitive):

mindspore/amp.py

@@ -21,8 +21,10 @@ from mindspore.common import mutable
 from mindspore.ops._primitive_cache import _get_cache_prim
 from mindspore.ops.operations.math_ops import NPUGetFloatStatusV2, NPUClearFloatStatusV2
 from mindspore.ops.operations.nn_ops import AllFinite
+from mindspore.run_check._check_version import AscendEnvChecker
 from mindspore import _checkparam as validator
 from mindspore._c_expression import MSContext
+from mindspore import log as logger
 from .common import dtype as mstype
 from . import context
 from . import ops
@@ -31,10 +33,9 @@ from .common.api import jit_class, jit
 from .common.parameter import Parameter
 from .common.tensor import Tensor
 from .train.loss_scale_manager import DynamicLossScaleManager, LossScaleManager, FixedLossScaleManager
-from .train.amp import build_train_network, auto_mixed_precision, custom_mixed_precision,\
+from .train.amp import build_train_network, auto_mixed_precision, custom_mixed_precision, \
     get_white_list, get_black_list
 
-
 _hypermap = ops.HyperMap()
 _partial = ops.Partial()
 
@@ -50,8 +51,8 @@ def _ascend_910a_target():
 
 
 @constexpr
-def _ascend_910bc_target():
-    return MSContext.get_instance().get_ascend_soc_version() in ["ascend910b", "ascend910c"]
+def _ascend_910b_target():
+    return MSContext.get_instance().get_ascend_soc_version() in ["ascend910b", "ascend910_93"]
 
 
 @constexpr
@@ -62,15 +63,24 @@ def _gpu_target():
 @constexpr
 def _enable_all_finite():
     """check whether enable all finite"""
+    logger.debug("Start enable all finite.")
+    if _ascend_target():
+        checker = AscendEnvChecker(None)
+        if not checker.check_custom_version():
+            logger.debug("Disable AllFinite due to version check failure.")
+            return False
     runtime_conf = os.environ.get('MS_DEV_RUNTIME_CONF')
     global_jit_config = context.get_jit_config()
     if runtime_conf is not None and ("all_finite:True" in runtime_conf or "all_finite:true" in runtime_conf):
+        logger.debug("Enable AllFinite through the environment variable MS_DEV_RUNTIME_CONF.")
         return True
 
     if runtime_conf is not None and ("all_finite:False" in runtime_conf or "all_finite:false" in runtime_conf):
+        logger.debug("Disable AllFinite through the environment variable MS_DEV_RUNTIME_CONF.")
         return False
 
     if global_jit_config:
+        logger.debug("Current global jit config is: {}".format(global_jit_config["jit_level"]))
         return global_jit_config["jit_level"] == "O0" or global_jit_config["jit_level"] == "O1"
     return False
 
@@ -105,7 +115,7 @@ def _all_finite(inputs, check_overflow_mode, enable_allfinite):
     """all finite check"""
     if _ascend_target():
         if (_ascend_910a_target()) or \
-           (_ascend_910bc_target() and check_overflow_mode == "SATURATION_MODE"):
+                (_ascend_910b_target() and check_overflow_mode == "SATURATION_MODE"):
             status = Tensor([0] * 8, mstype.int32)
             status = ops.depend(status, inputs)
             get_status = _get_cache_prim(NPUGetFloatStatusV2)()(status)
@@ -153,7 +163,7 @@ def all_finite(inputs):
 
     Tutorial Examples:
         - `Automatic Mix Precision - Loss Scaling
-          <https://mindspore.cn/tutorials/en/master/advanced/mixed_precision.html#loss-scaling>`_
+          <https://mindspore.cn/tutorials/en/master/beginner/mixed_precision.html#loss-scaling>`_
     """
     inputs = mutable(inputs)
     _check_overflow_mode = os.environ.get('MS_ASCEND_CHECK_OVERFLOW_MODE')
@@ -169,7 +179,7 @@ class LossScaler(ABC):
     to scale and unscale the loss value and gradients to avoid overflow, `adjust` is used to update the
     loss scale value.
 
-    For more information, refer to the `tutorials  <https://mindspore.cn/tutorials/en/master/advanced/
+    For more information, refer to the `tutorials  <https://mindspore.cn/tutorials/en/master/beginner/
     mixed_precision.html#loss-scaling>`_.
 
     .. warning::
@@ -200,6 +210,7 @@ class LossScaler(ABC):
         >>>
         >>> loss_scaler = MyLossScaler(1024)
     """
+
     @abstractmethod
     def scale(self, inputs):
         """
@@ -262,6 +273,7 @@ class StaticLossScaler(LossScaler):
         (Tensor(shape=[2], dtype=Float16, value= [ 1.4648e-03,  9.7656e-04]),
         Tensor(shape=[1], dtype=Float16, value= [ 1.1721e-03]))
     """
+
     def __init__(self, scale_value):
         scale_value = validator.check_value_type("scale_value", scale_value, [float, int])
         if scale_value < 1.0:
@@ -340,6 +352,7 @@ class DynamicLossScaler(LossScaler):
         >>> print(loss_scaler.scale_value.asnumpy())
         512.0
     """
+
     def __init__(self, scale_value, scale_factor, scale_window):
         scale_value = validator.check_value_type("scale_value", scale_value, [float, int])
         if scale_value < 1.0:
@@ -361,7 +374,7 @@ class DynamicLossScaler(LossScaler):
 
         Tutorial Examples:
             - `Automatic Mix Precision - Loss Scaling
-              <https://mindspore.cn/tutorials/en/master/advanced/mixed_precision.html#loss-scaling>`_
+              <https://mindspore.cn/tutorials/en/master/beginner/mixed_precision.html#loss-scaling>`_
         """
         inputs = mutable(inputs)
         return _grad_scale_map(self.scale_value, inputs)
@@ -378,7 +391,7 @@ class DynamicLossScaler(LossScaler):
 
         Tutorial Examples:
             - `Automatic Mix Precision - Loss Scaling
-              <https://mindspore.cn/tutorials/en/master/advanced/mixed_precision.html#loss-scaling>`_
+              <https://mindspore.cn/tutorials/en/master/beginner/mixed_precision.html#loss-scaling>`_
         """
         inputs = mutable(inputs)
         return _grad_unscale_map(self.scale_value, inputs)
@@ -392,7 +405,7 @@ class DynamicLossScaler(LossScaler):
 
         Tutorial Examples:
             - `Automatic Mix Precision - Loss Scaling
-              <https://mindspore.cn/tutorials/en/master/advanced/mixed_precision.html#loss-scaling>`_
+              <https://mindspore.cn/tutorials/en/master/beginner/mixed_precision.html#loss-scaling>`_
         """
         one = ops.ones((), self.scale_value.dtype)
         scale_mul_factor = self.scale_value * self.scale_factor
@@ -411,6 +424,7 @@ class DynamicLossScaler(LossScaler):
         ops.assign(self.counter, counter)
         return True
 
+
 __all__ = [
     "DynamicLossScaleManager", "LossScaleManager", "FixedLossScaleManager",
     "build_train_network", "DynamicLossScaler", "StaticLossScaler", "LossScaler",