add nisa.nc_n_gather in place of nl.gather_flattened

Author: yongweiy

KLR/Core/Basic.lean

@@ -181,6 +181,7 @@ partial def operatorBasicTensors : Operator → List TensorRef
   | .setRngSeed r | .randSetState r => [r.src]
   | .extendedInst _ => []
   | .tensorScalarCumulative t => [t.dst, t.src]
+  | .ncNGather g => [g.dst, g.data, g.indices]
 
 partial def operatorAdditionalTensors : Operator → List TensorName
   | .ncActivate d => (tensors d.scale) ++ (tensors d.bias) ++ (tensors d.reduceRes)
@@ -202,6 +203,7 @@ partial def operatorAdditionalTensors : Operator → List TensorName
   | .recv r => tensors r.dsts
   | .rand2 r => tensors r.min ++ tensors r.max
   | .tensorScalarCumulative t => (tensors t.imm0) ++ (tensors t.imm1)
+  | .ncNGather _ => []
   | _ => []
 
 instance : Tensors Operator where

KLR/Core/LowerAP.lean

@@ -197,6 +197,11 @@ def Operator.lowerAccessPatterns (k : Operator) : KLR.Err Operator :=
       imm0 := <- Operand.lowerAccessPatterns op.imm0
       imm1 := <- op.imm1.mapM Operand.lowerAccessPatterns
     }
+  | .ncNGather op => return .ncNGather { op with
+      dst := <- op.dst.lowerAccessPatterns
+      data := <- op.data.lowerAccessPatterns
+      indices := <- op.indices.lowerAccessPatterns
+    }
 
 def Stmt.lowerAccessPatterns : Stmt → KLR.Err Stmt
   | .oper op name pos => return .oper (<- op.lowerAccessPatterns) name pos

KLR/Core/Operators.lean

@@ -944,6 +944,14 @@ structure TensorScalarCumulative where
   deriving BEq, FromCBOR, FromJson, FromSexp, Repr, ToCBOR, ToJson, ToSexp
 
 @[serde tag = 209]
+structure NcNGather where
+  dst : TensorRef
+  data : TensorRef
+  indices : TensorRef
+  dtype : Option Dtype
+  deriving BEq, FromCBOR, FromJson, FromSexp, Repr, ToCBOR, ToJson, ToSexp
+
+@[serde tag = 210]
 inductive Operator where
   | activate (op : Activate)
   | ncActivate (op : NcActivate)
@@ -1012,9 +1020,10 @@ inductive Operator where
   | randSetState (op : RandSetState)
   | extendedInst (op : ExtendedInst)
   | tensorScalarCumulative (op: TensorScalarCumulative)
+  | ncNGather (op: NcNGather)
   deriving BEq, FromCBOR, FromJson, FromSexp, Repr, ToCBOR, ToJson, ToSexp
 
-@[serde tag = 210]
+@[serde tag = 211]
 inductive TGROperator where
   | activate (op : Activate)
   | affineSelect (op : AffineSelect)

KLR/Extract/Extract/Basic.lean

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

KLR/Trace/ISA.lean

@@ -1180,3 +1180,19 @@ nki builtin.isa.extended_inst
     data1
   }) name
   return .none
+
+nki builtin.isa.nc_n_gather
+  (dst: Access)
+  (data: Access)
+  (indices: Access)
+  (mask: Option Immediate := none)
+  (dtype: Option Dtype := none)
+  (name : Option String := none) := do
+    if mask.isSome then throw maskNotSupported
+    Trace.add_stmt $ .oper (.ncNGather {
+      dst := .abstract dst
+      data := .abstract data
+      indices := .abstract indices
+      dtype := dtype.or dst.tensor.dtype
+    }) name
+    return .none

interop/klr/NKI.asdl

@@ -339,6 +339,8 @@ 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)
+
+NcNGather = (TensorRef dst, TensorRef data, TensorRef indices, Dtype? dtype)
 Operator =
  | activate(Activate op)
  | ncActivate(NcActivate op)
@@ -407,6 +409,7 @@ Operator =
  | randSetState(RandSetState op)
  | extendedInst(ExtendedInst op)
  | tensorScalarCumulative(TensorScalarCumulative op)
+ | ncNGather(NcNGather op)
 
 Stmt =
  | oper(Operator op, String? name, Pos pos)

interop/klr/klir_ast.hpp

@@ -998,6 +998,13 @@ struct TensorScalarCumulative final {
   Option<Dtype> dtype;
 };
 
+struct NcNGather final {
+  Ptr<TensorRef> dst;
+  Ptr<TensorRef> data;
+  Ptr<TensorRef> indices;
+  Option<Dtype> dtype;
+};
+
 struct Operator {
   enum class Tag {
     activate = 1,
@@ -1067,6 +1074,7 @@ struct Operator {
     randSetState,
     extendedInst,
     tensorScalarCumulative,
+    ncNGather,
   };
   Tag tag;
   Operator(Tag tag) : tag(tag) {}
@@ -1409,6 +1417,11 @@ struct OperatorTensorScalarCumulativeWrapper final : Operator {
       : Operator(Tag::tensorScalarCumulative) {}
 };
 
+struct OperatorNcNGatherWrapper final : Operator {
+  Ptr<NcNGather> op;
+  OperatorNcNGatherWrapper() : Operator(Tag::ncNGather) {}
+};
+
 struct Stmt {
   enum class Tag {
     oper = 1,

interop/klr/klir_serde.cpp

@@ -3207,10 +3207,17 @@ Ptr<ExtendedInst> ExtendedInst_des(FILE *in) {
 
 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");
+  if (!deserialize_tag(in, &t, &c, &l)) {
+    std::ostringstream msg;
+    msg << "Could not find tag, expecting TensorScalarCumulative:208,0";
+    throw std::runtime_error(msg.str());
+  }
+  if (t != 208 || c != 0 || l != 9) {
+    std::ostringstream msg;
+    msg << "Expecting TensorScalarCumulative:(208,0,9)";
+    msg << " got:(" << (int)t << "," << (int)c << "," << (int)l << ")";
+    throw std::runtime_error(msg.str());
+  }
   Ptr<TensorScalarCumulative> x = ptr<TensorScalarCumulative>();
   x->dst = TensorRef_des(in);
   x->src = TensorRef_des(in);
@@ -3224,11 +3231,32 @@ Ptr<TensorScalarCumulative> TensorScalarCumulative_des(FILE *in) {
   return x;
 }
 
+Ptr<NcNGather> NcNGather_des(FILE *in) {
+  u8 t, c, l;
+  if (!deserialize_tag(in, &t, &c, &l)) {
+    std::ostringstream msg;
+    msg << "Could not find tag, expecting NcNGather:209,0";
+    throw std::runtime_error(msg.str());
+  }
+  if (t != 209 || c != 0 || l != 4) {
+    std::ostringstream msg;
+    msg << "Expecting NcNGather:(209,0,4)";
+    msg << " got:(" << (int)t << "," << (int)c << "," << (int)l << ")";
+    throw std::runtime_error(msg.str());
+  }
+  Ptr<NcNGather> x = ptr<NcNGather>();
+  x->dst = TensorRef_des(in);
+  x->data = TensorRef_des(in);
+  x->indices = TensorRef_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 != 209)
+  if (t != 210)
     throw std::runtime_error("Unexpected type tag");
   switch (c) {
   case 0: {
@@ -3777,6 +3805,14 @@ Ptr<Operator> Operator_des(FILE *in) {
     return x;
     break;
   }
+  case 67: {
+    if (l != 1)
+      throw std::runtime_error("Wrong number of elements");
+    Ptr<OperatorNcNGatherWrapper> x = ptr<OperatorNcNGatherWrapper>();
+    x->op = NcNGather_des(in);
+    return x;
+    break;
+  }
   default:
     throw std::runtime_error("Invalid value tag");
   }

interop/klr/klir_serde.hpp

@@ -135,6 +135,7 @@ 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<NcNGather> NcNGather_des(FILE *in);
 Ptr<Operator> Operator_des(FILE *in);
 Ptr<Stmt> Stmt_des(FILE *in);
 Ptr<Block> Block_des(FILE *in);