bound check kernel implementation

Reviewed By: fxzhang Differential Revision: D48986845
pytorch · Sep 27, 2023 · 376bc32 · 376bc32
1 parent 0a53b9f
commit 376bc32
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 0 deletions.
diff --git a/lib/Backends/Interpreter/InterpreterNodes.cpp b/lib/Backends/Interpreter/InterpreterNodes.cpp
@@ -2248,6 +2248,11 @@ void BoundInterpreterFunction::fwdPrunedArrayLookupInst(
   llvm_unreachable("PrunedArrayLookup instruction is not supported yet");
 }
 
+void BoundInterpreterFunction::fwdBoundsCheckIndicesInst(
+    const BoundsCheckIndicesInst *) {
+  llvm_unreachable("BoundsCheckIndices instruction is not supported yet");
+}
+
 template <typename ElemTy>
 void BoundInterpreterFunction::fwdTanhInstFloatImpl(const TanhInst *I) {
   staticAssertFloatingPointType(ElemTy);

diff --git a/tools/ClassGen/InstrGen.cpp b/tools/ClassGen/InstrGen.cpp
@@ -1163,6 +1163,19 @@ int main(int argc, char **argv) {
                   {"IndexRemapping", "ElemKind::Int32ITy"})
       .autoVerify(VerifyKind::SameElementType,
                   {"IndexRemappingOffsets", "ElemKind::Int64ITy"});
+
+  BB.newInstr("BoundsCheckIndices")
+      .addOperand("RowsPerTable", OperandKind::In)
+      .addOperand("Indices", OperandKind::In)
+      .addOperand("Offsets", OperandKind::In)
+      .addOperand("Warning", OperandKind::Out)
+      .addOperand("Weights", OperandKind::In)
+      .addMember(MemberType::Unsigned, "BoundsCheckMode")
+      .autoVerify(VerifyKind::SameElementType,
+                  {"Indices", "ElemKind::Int32ITy"})
+      .autoVerify(VerifyKind::SameElementType,
+                  {"Offsets", "ElemKind::Int32ITy"});
+
   BB.newInstr("IndexPut")
       .addOperand("Dest", OperandKind::Out)
       .addOperand("Input", OperandKind::In)