wehub-resource-sync
154 lines
4.0 KiB
Python
154 lines
4.0 KiB
Python
# Licensed to the Apache Software Foundation (ASF) under one
|
|
# or more contributor license agreements. See the NOTICE file
|
|
# distributed with this work for additional information
|
|
# regarding copyright ownership. The ASF licenses this file
|
|
# to you under the Apache License, Version 2.0 (the
|
|
# "License"); you may not use this file except in compliance
|
|
# with the License. You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing,
|
|
# software distributed under the License is distributed on an
|
|
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
|
# KIND, either express or implied. See the License for the
|
|
# specific language governing permissions and limitations
|
|
# under the License.
|
|
"""Trainium-owned NKI intrinsic Python wrappers."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from tvm.tirx.op import call_intrin
|
|
|
|
|
|
def nki_load(res, data):
|
|
return call_intrin("", "tirx.nki.load", res, data)
|
|
|
|
|
|
def nki_store(res, data):
|
|
return call_intrin("", "tirx.nki.store", res, data)
|
|
|
|
|
|
def nki_tensor_copy(res, data):
|
|
return call_intrin("", "tirx.nki.tensor_copy", res, data)
|
|
|
|
|
|
def nki_matmul(res, lhs, rhs, accum=True):
|
|
return call_intrin("", "tirx.nki.matmul", res, lhs, rhs, accum)
|
|
|
|
|
|
def nki_activation(result, data, opcode, bias=0.0, scale=1.0):
|
|
return call_intrin("", "tirx.nki.activation", result, data, opcode, bias, scale)
|
|
|
|
|
|
def nki_reciprocal(result, data):
|
|
return call_intrin("", "tirx.nki.reciprocal", result, data)
|
|
|
|
|
|
def nki_tensorreduce(result, data, opcode, negate, *axes):
|
|
return call_intrin("", "tirx.nki.tensorreduce", result, data, opcode, negate, *axes)
|
|
|
|
|
|
def nki_tensortensor(result, operand0, operand1, opcode):
|
|
return call_intrin("", "tirx.nki.tensortensor", result, operand0, operand1, opcode)
|
|
|
|
|
|
def nki_tensorscalar(result, operand0, operand1, opcode, reverse=False):
|
|
return call_intrin("", "tirx.nki.tensorscalar", result, operand0, operand1, opcode, reverse)
|
|
|
|
|
|
def nki_memset(result, value):
|
|
return call_intrin("", "tirx.nki.memset", result, value)
|
|
|
|
|
|
def nki_activation_reduce(reduce_res, act_res, data, opcode, reduce_opcode, bias=0.0, scale=1.0):
|
|
return call_intrin(
|
|
"",
|
|
"tirx.nki.activation_reduce",
|
|
reduce_res,
|
|
act_res,
|
|
data,
|
|
opcode,
|
|
reduce_opcode,
|
|
bias,
|
|
scale,
|
|
)
|
|
|
|
|
|
def nki_tensorscalar_reduce(
|
|
reduce_res, tensorscalar_res, operand0, operand1, opcode, reduce_opcode, reverse=False
|
|
):
|
|
return call_intrin(
|
|
"",
|
|
"tirx.nki.tensorscalar_reduce",
|
|
reduce_res,
|
|
tensorscalar_res,
|
|
operand0,
|
|
operand1,
|
|
opcode,
|
|
reduce_opcode,
|
|
reverse,
|
|
)
|
|
|
|
|
|
def nki_identity(result, size):
|
|
return call_intrin("", "tirx.nki.identity", result, size)
|
|
|
|
|
|
def nki_scalar_tensor_tensor(
|
|
result, data, operand0, operand1, opcode0, opcode1, reverse0=False, reverse1=False
|
|
):
|
|
return call_intrin(
|
|
"",
|
|
"tirx.nki.scalar_tensor_tensor",
|
|
result,
|
|
data,
|
|
operand0,
|
|
operand1,
|
|
opcode0,
|
|
opcode1,
|
|
reverse0,
|
|
reverse1,
|
|
)
|
|
|
|
|
|
def nki_scalar_tensor_scalar(
|
|
result, data, operand0, operand1, opcode0, opcode1, reverse0=False, reverse1=False
|
|
):
|
|
return call_intrin(
|
|
"",
|
|
"tirx.nki.scalar_tensor_scalar",
|
|
result,
|
|
data,
|
|
operand0,
|
|
operand1,
|
|
opcode0,
|
|
opcode1,
|
|
reverse0,
|
|
reverse1,
|
|
)
|
|
|
|
|
|
def nki_affine_select(result, pred, true_value, false_value):
|
|
return call_intrin("", "tirx.nki.affine_select", result, pred, true_value, false_value)
|
|
|
|
|
|
__all__ = [
|
|
"nki_activation",
|
|
"nki_activation_reduce",
|
|
"nki_affine_select",
|
|
"nki_identity",
|
|
"nki_load",
|
|
"nki_matmul",
|
|
"nki_memset",
|
|
"nki_reciprocal",
|
|
"nki_scalar_tensor_scalar",
|
|
"nki_scalar_tensor_tensor",
|
|
"nki_store",
|
|
"nki_tensor_copy",
|
|
"nki_tensorreduce",
|
|
"nki_tensorscalar",
|
|
"nki_tensorscalar_reduce",
|
|
"nki_tensortensor",
|
|
]
|