This document is relevant for: Trn2, Trn3

nki.isa.memset#

nki.isa.memset(dst: NkiTensor, value, engine=engine.unknown, name=None)[source]#

Initialize dst by filling it with a compile-time constant value, using Vector or GpSimd Engine. The memset instruction supports all valid NKI dtypes (see Supported Data Types). dst can have at most four free dimensions (maximum rank 5, including the partition dimension).

Parameters:
  • dst – destination tile to initialize.

  • value – the constant value to initialize with

  • engine – specify which engine to use for memset: nki.isa.engine.vector or nki.isa.engine.gpsimd ; nki.isa.engine.unknown by default, lets compiler select the best engine for the given input tile shape

Note

For x4 packed types (float8_e4m3fn_x4, float8_e5m2_x4, float4_e2m1fn_x4), only value=0 is supported.

This document is relevant for: Trn2, Trn3