Implement real pruned_array_lookup op lowering

Reviewed By: jfix71, fxzhang Differential Revision: D46487261 fbshipit-source-id: 10532ea28ffc05f369cae8bbdfa61e4178367217
pytorch · Jun 26, 2023 · 29ffe6d · 29ffe6d
1 parent 35678b6
commit 29ffe6d
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 0 deletions.
diff --git a/lib/Backends/Interpreter/InterpreterNodes.cpp b/lib/Backends/Interpreter/InterpreterNodes.cpp
@@ -2243,6 +2243,11 @@ void BoundInterpreterFunction::fwdPadInst(const PadInst *) {
   llvm_unreachable("Pad instruction is not supported yet");
 }
 
+void BoundInterpreterFunction::fwdPrunedArrayLookupInst(
+    const PrunedArrayLookupInst *) {
+  llvm_unreachable("PrunedArrayLookup 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
@@ -1149,6 +1149,20 @@ int main(int argc, char **argv) {
       .addMember(MemberType::VectorUnsigned, "MeanDims")
       .autoVerify(VerifyKind::SameElementType, {"Dest", "Src"});
 
+  BB.newInstr("PrunedArrayLookup")
+      .addOperand("Dest", OperandKind::Out)
+      .addOperand("Indices", OperandKind::In)
+      .addOperand("Offsets", OperandKind::In)
+      .addOperand("IndexRemapping", OperandKind::In)
+      .addOperand("IndexRemappingOffsets", OperandKind::In)
+      .autoVerify(VerifyKind::SameElementType,
+                  {"Indices", "ElemKind::Int32ITy"})
+      .autoVerify(VerifyKind::SameElementType,
+                  {"Offsets", "ElemKind::Int32ITy"})
+      .autoVerify(VerifyKind::SameElementType,
+                  {"IndexRemapping", "ElemKind::Int32ITy"})
+      .autoVerify(VerifyKind::SameElementType,
+                  {"IndexRemappingOffsets", "ElemKind::Int64ITy"});
   //===--------------------------------------------------------------------===//
   //                Fillers
   //===--------------------------------------------------------------------===//