feat: add dtype override to ap method

Author: govereau

KLR/Core/Tensor.lean

@@ -569,6 +569,7 @@ structure BirAccessPattern where
   scalarOffset : Option ScalarOffset
   vectorOffset : Option Access
   indirectDim : Int
+  dtypeOverride : Option Dtype := none
   deriving BEq, FromCBOR, FromJson, FromSexp, Repr, ToCBOR, ToJson, ToSexp
 
 

KLR/Trace/Term.lean

@@ -594,7 +594,8 @@ nki builtin.access.ap
     (offset : Nat := 0)
     (scalar_offset : Option (Sum Access Term) := none)
     (vector_offset : Option Access := none)
-    (indirect_dim : Int := 0) := do
+    (indirect_dim : Int := 0)
+    (dtype : Option Dtype := none) := do
   match self with
   | .simple t =>
       let pattern := pattern.map fun (s,c) => Core.APPair.mk s c 0
@@ -610,6 +611,7 @@ nki builtin.access.ap
         scalarOffset
         vectorOffset := vector_offset
         indirectDim := indirect_dim
+        dtypeOverride := dtype
       }
       return .access (.birPattern ap)
   -- TODO: need to figoure out how to combine cannonical form AP with user specified AP

interop/klr/NKI.asdl

@@ -52,7 +52,7 @@ ScalarOffset =
  | acc(Access a)
 
 
-BirAccessPattern = (TensorName tensor, Nat offset, APPair* pattern, ScalarOffset? scalarOffset, Access? vectorOffset, Int indirectDim)
+BirAccessPattern = (TensorName tensor, Nat offset, APPair* pattern, ScalarOffset? scalarOffset, Access? vectorOffset, Int indirectDim, Dtype? dtypeOverride)
 Access =
  | simple(TensorName tensor)
  | basic(AccessBasic access)

interop/klr/klir_ast.hpp

@@ -181,6 +181,7 @@ struct BirAccessPattern final {
   Option<Ptr<ScalarOffset>> scalarOffset;
   Option<Ptr<Access>> vectorOffset;
   Int indirectDim;
+  Option<Dtype> dtypeOverride;
 };
 
 struct Access {

interop/klr/klir_serde.cpp

@@ -749,9 +749,9 @@ Ptr<BirAccessPattern> BirAccessPattern_des(FILE *in) {
     msg << "Could not find tag, expecting BirAccessPattern:124,0";
     throw std::runtime_error(msg.str());
   }
-  if (t != 124 || c != 0 || l != 6) {
+  if (t != 124 || c != 0 || l != 7) {
     std::ostringstream msg;
-    msg << "Expecting BirAccessPattern:(124,0,6)";
+    msg << "Expecting BirAccessPattern:(124,0,7)";
     msg << " got:(" << (int)t << "," << (int)c << "," << (int)l << ")";
     throw std::runtime_error(msg.str());
   }
@@ -762,6 +762,7 @@ Ptr<BirAccessPattern> BirAccessPattern_des(FILE *in) {
   x->scalarOffset = Option_ScalarOffset_des(in);
   x->vectorOffset = Option_Access_des(in);
   x->indirectDim = Int_des(in);
+  x->dtypeOverride = Option_Dtype_des(in);
   return x;
 }
 
@@ -3207,10 +3208,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);