Add tensor_scalar_cumulative API

Author: jiancheng-aws

KLR/Core/Basic.lean

@@ -180,6 +180,7 @@ partial def operatorBasicTensors : Operator → List TensorRef
   | .rng r | .rand2 r | .randGetState r  => [r.dst]
   | .setRngSeed r | .randSetState r => [r.src]
   | .extendedInst _ => []
+  | .tensorScalarCumulative t => [t.dst, t.src]
 
 partial def operatorAdditionalTensors : Operator → List TensorName
   | .ncActivate d => (tensors d.scale) ++ (tensors d.bias) ++ (tensors d.reduceRes)
@@ -200,6 +201,7 @@ partial def operatorAdditionalTensors : Operator → List TensorName
   | .send s => tensors s.srcs
   | .recv r => tensors r.dsts
   | .rand2 r => tensors r.min ++ tensors r.max
+  | .tensorScalarCumulative t => (tensors t.imm0) ++ (tensors t.imm1)
   | _ => []
 
 instance : Tensors Operator where

KLR/Core/LowerAP.lean

@@ -191,6 +191,12 @@ def Operator.lowerAccessPatterns (k : Operator) : KLR.Err Operator :=
   | .setRngSeed r => return .setRngSeed { r with src := (<- r.src.lowerAccessPatterns)}
   | .randSetState r => return .randSetState { r with src := (<- r.src.lowerAccessPatterns)}
   | .extendedInst i => return .extendedInst i
+  | .tensorScalarCumulative op => return .tensorScalarCumulative { op with
+      dst := <- op.dst.lowerAccessPatterns
+      src := <- op.src.lowerAccessPatterns
+      imm0 := <- Operand.lowerAccessPatterns op.imm0
+      imm1 := <- op.imm1.mapM Operand.lowerAccessPatterns
+    }
 
 def Stmt.lowerAccessPatterns : Stmt → KLR.Err Stmt
   | .oper op name pos => return .oper (<- op.lowerAccessPatterns) name pos

KLR/Core/Operators.lean

@@ -931,6 +931,19 @@ structure ExtendedInst where
 deriving BEq, FromCBOR, FromJson, FromSexp, Repr, ToCBOR, ToJson, ToSexp
 
 @[serde tag = 208]
+structure TensorScalarCumulative where
+  dst : TensorRef
+  src : TensorRef
+  op0 : AluOp
+  op1: AluOp
+  imm0: Operand
+  imm1: Option Operand
+  reduceCmd: AccumCmd
+  reverse : TensorScalarReverseOps
+  dtype : Option Dtype
+  deriving BEq, FromCBOR, FromJson, FromSexp, Repr, ToCBOR, ToJson, ToSexp
+
+@[serde tag = 209]
 inductive Operator where
   | activate (op : Activate)
   | ncActivate (op : NcActivate)
@@ -998,9 +1011,10 @@ inductive Operator where
   | setRngSeed (op : SetRngSeed)
   | randSetState (op : RandSetState)
   | extendedInst (op : ExtendedInst)
+  | tensorScalarCumulative (op: TensorScalarCumulative)
   deriving BEq, FromCBOR, FromJson, FromSexp, Repr, ToCBOR, ToJson, ToSexp
 
-@[serde tag = 209]
+@[serde tag = 210]
 inductive TGROperator where
   | activate (op : Activate)
   | affineSelect (op : AffineSelect)

KLR/Extract/Extract/Basic.lean

@@ -366,6 +366,7 @@ def klrAST: MetaM (List LeanType) := do
     `KLR.Core.SetRngSeed,
     `KLR.Core.RandSetState,
     `KLR.Core.ExtendedInst,
+    `KLR.Core.TensorScalarCumulative,
     `KLR.Core.Operator,
     -- Core.Basic
     `KLR.Core.Stmt,

KLR/Trace/ISA.lean

@@ -426,6 +426,36 @@ nki builtin.isa.tensor_scalar_reduce
     }) name
     return .none
 
+nki builtin.isa.tensor_scalar_cumulative
+  (dst: Access)
+  (src: Access)
+  (op0: AluOp)
+  (op1: AluOp)
+  (imm0: Sum Immediate Access)
+  (imm1: Option (Sum Immediate Access) := none)
+  (reduce_cmd: AccumCmd := AccumCmd.ZeroAccumulate)
+  (mask: Option Immediate := none)
+  (name : Option String := none) := do
+    if mask.isSome then
+      throw maskNotSupported
+    Trace.add_stmt $ .oper (.tensorScalarCumulative {
+      dst := .abstract dst
+      src := .abstract src
+      op0 := op0
+      op1 := op1
+      imm0 := match imm0 with
+        | .inl i => .imm i
+        | .inr t => .tile $ .abstract t
+      imm1 := match imm1 with
+        | some (.inl i) => some $ .imm i
+        | some (.inr t) => some $ .tile $ .abstract t
+        | none => .none
+      reduceCmd := reduce_cmd
+      reverse := TensorScalarReverseOps.none
+      dtype := dst.tensor.dtype
+    }) name
+    return .none
+
 nki builtin.isa.tensor_copy
  (dst: Access)
  (src : Access)

interop/klr/NKI.asdl

@@ -337,6 +337,8 @@ SetRngSeed = (TensorRef src)
 RandSetState = (TensorRef src, Engine engine)
 
 ExtendedInst = (Nat opcode, Bool hasRead, Bool hasWrite, Nat ports, Nat* data0, Nat* data1)
+
+TensorScalarCumulative = (TensorRef dst, TensorRef src, AluOp op0, AluOp op1, Operand imm0, Operand? imm1, AccumCmd reduceCmd, TensorScalarReverseOps reverse, Dtype? dtype)
 Operator =
  | activate(Activate op)
  | ncActivate(NcActivate op)
@@ -404,6 +406,7 @@ Operator =
  | setRngSeed(SetRngSeed op)
  | randSetState(RandSetState op)
  | extendedInst(ExtendedInst op)
+ | tensorScalarCumulative(TensorScalarCumulative op)
 
 Stmt =
  | oper(Operator op, String? name, Pos pos)

interop/klr/klir_ast.hpp

@@ -986,6 +986,18 @@ struct ExtendedInst final {
   List<Nat> data1;
 };
 
+struct TensorScalarCumulative final {
+  Ptr<TensorRef> dst;
+  Ptr<TensorRef> src;
+  AluOp op0;
+  AluOp op1;
+  Ptr<Operand> imm0;
+  Option<Ptr<Operand>> imm1;
+  AccumCmd reduceCmd;
+  TensorScalarReverseOps reverse;
+  Option<Dtype> dtype;
+};
+
 struct Operator {
   enum class Tag {
     activate = 1,
@@ -1054,6 +1066,7 @@ struct Operator {
     setRngSeed,
     randSetState,
     extendedInst,
+    tensorScalarCumulative,
   };
   Tag tag;
   Operator(Tag tag) : tag(tag) {}
@@ -1390,6 +1403,12 @@ struct OperatorExtendedInstWrapper final : Operator {
   OperatorExtendedInstWrapper() : Operator(Tag::extendedInst) {}
 };
 
+struct OperatorTensorScalarCumulativeWrapper final : Operator {
+  Ptr<TensorScalarCumulative> op;
+  OperatorTensorScalarCumulativeWrapper()
+      : Operator(Tag::tensorScalarCumulative) {}
+};
+
 struct Stmt {
   enum class Tag {
     oper = 1,

interop/klr/klir_serde.cpp

@@ -3205,11 +3205,30 @@ Ptr<ExtendedInst> ExtendedInst_des(FILE *in) {
   return x;
 }
 
+Ptr<TensorScalarCumulative> TensorScalarCumulative_des(FILE *in) {
+  u8 t, c, l;
+  if (!deserialize_tag(in, &t, &c, &l))
+    throw std::runtime_error("Could not find tag");
+  if (t != 208 || c != 0 || l != 9)
+    throw std::runtime_error("Invalid Tag");
+  Ptr<TensorScalarCumulative> x = ptr<TensorScalarCumulative>();
+  x->dst = TensorRef_des(in);
+  x->src = TensorRef_des(in);
+  x->op0 = AluOp_des(in);
+  x->op1 = AluOp_des(in);
+  x->imm0 = Operand_des(in);
+  x->imm1 = Option_Operand_des(in);
+  x->reduceCmd = AccumCmd_des(in);
+  x->reverse = TensorScalarReverseOps_des(in);
+  x->dtype = Option_Dtype_des(in);
+  return x;
+}
+
 Ptr<Operator> Operator_des(FILE *in) {
   u8 t, c, l;
   if (!deserialize_tag(in, &t, &c, &l))
     throw std::runtime_error("Could not read tag");
-  if (t != 208)
+  if (t != 209)
     throw std::runtime_error("Unexpected type tag");
   switch (c) {
   case 0: {
@@ -3749,6 +3768,15 @@ Ptr<Operator> Operator_des(FILE *in) {
     return x;
     break;
   }
+  case 66: {
+    if (l != 1)
+      throw std::runtime_error("Wrong number of elements");
+    Ptr<OperatorTensorScalarCumulativeWrapper> x =
+        ptr<OperatorTensorScalarCumulativeWrapper>();
+    x->op = TensorScalarCumulative_des(in);
+    return x;
+    break;
+  }
   default:
     throw std::runtime_error("Invalid value tag");
   }

interop/klr/klir_serde.hpp

@@ -134,6 +134,7 @@ Ptr<RandGetState> RandGetState_des(FILE *in);
 Ptr<SetRngSeed> SetRngSeed_des(FILE *in);
 Ptr<RandSetState> RandSetState_des(FILE *in);
 Ptr<ExtendedInst> ExtendedInst_des(FILE *in);
+Ptr<TensorScalarCumulative> TensorScalarCumulative_des(FILE *in);
 Ptr<Operator> Operator_des(FILE *in);
 Ptr<Stmt> Stmt_des(FILE *in);
 Ptr<Block> Block_des(FILE *in);