From b1d38deafc3228acc3a06053e8e0359da617e659 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E7=9F=B3=E6=99=93=E4=BC=9F?=
 <39303645+Shixiaowei02@users.noreply.github.com>
Date: Thu, 3 Mar 2022 16:46:54 +0800
Subject: [PATCH] mlir attr types for infrt place, test=develop (#40087)

* mlir attr types for infrt place, test=develop

* fix a bug, test=develop
---
 paddle/infrt/dialect/infrt/common_type.cc     | 33 +++++++------
 paddle/infrt/dialect/infrt/common_type.h      | 11 +++--
 paddle/infrt/dialect/infrt/infrt_dialect.cc   |  1 +
 paddle/infrt/dialect/infrt/infrt_ops_base.td  | 46 ++++++++++++++++++-
 paddle/infrt/dialect/phi/ir/infrt_phi_base.td | 38 +++++++--------
 .../infrt/dialect/phi/ir/infrt_phi_tensor.td  | 10 ++--
 paddle/infrt/dialect/phi/ir/phi_base.cc       | 40 +---------------
 paddle/infrt/dialect/phi/ir/phi_base.h        |  7 +--
 paddle/infrt/host_context/kernel_frame.cc     | 37 ++++++---------
 .../phi/infershaped/phi_kernel_launcher.h     | 20 --------
 paddle/infrt/kernel/phi/registry.cc           | 11 -----
 .../tests/dialect/pten/dense_tensor.mlir      |  8 ++--
 tools/infrt/generate_phi_kernel_dialect.py    |  2 +-
 13 files changed, 121 insertions(+), 143 deletions(-)

diff --git a/paddle/infrt/dialect/infrt/common_type.cc b/paddle/infrt/dialect/infrt/common_type.cc
index 5cbd7b2cd6..00684c5052 100644
--- a/paddle/infrt/dialect/infrt/common_type.cc
+++ b/paddle/infrt/dialect/infrt/common_type.cc
@@ -43,46 +43,49 @@ llvm::Optional<PrecisionType> GetPrecisionType(llvm::StringRef key) {
     return llvm::None;
 }
 
-llvm::raw_ostream &operator<<(llvm::raw_ostream &os, TargetType type) {
+llvm::StringRef GetString(TargetType type) {
+  llvm::StringRef str;
   switch (type) {
     case (TargetType::CPU):
-      os << "CPU";
+      str = "CPU";
       break;
     case (TargetType::GPU):
-      os << "GPU";
+      str = "GPU";
       break;
     default:
-      os << "Unsupported";
+      str = "Unsupported";
   }
-  return os;
+  return str;
 }
 
-llvm::raw_ostream &operator<<(llvm::raw_ostream &os, LayoutType type) {
+llvm::StringRef GetString(LayoutType type) {
+  llvm::StringRef str;
   switch (type) {
     case (LayoutType::NCHW):
-      os << "NCHW";
+      str = "NCHW";
       break;
     case (LayoutType::NHWC):
-      os << "NHWC";
+      str = "NHWC";
       break;
     default:
-      os << "Unsupported";
+      str = "Unsupported";
   }
-  return os;
+  return str;
 }
 
-llvm::raw_ostream &operator<<(llvm::raw_ostream &os, PrecisionType type) {
+llvm::StringRef GetString(PrecisionType type) {
+  llvm::StringRef str;
   switch (type) {
     case (PrecisionType::FLOAT32):
-      os << "FP32";
+      str = "FP32";
       break;
     case (PrecisionType::FLOAT16):
-      os << "FP16";
+      str = "FP16";
       break;
     default:
-      os << "Unsupported";
+      str = "Unsupported";
   }
-  return os;
+  return str;
 }
 
 }  // namespace infrt
diff --git a/paddle/infrt/dialect/infrt/common_type.h b/paddle/infrt/dialect/infrt/common_type.h
index 436e7920ca..2ebe2b8ccd 100644
--- a/paddle/infrt/dialect/infrt/common_type.h
+++ b/paddle/infrt/dialect/infrt/common_type.h
@@ -54,8 +54,13 @@ llvm::Optional<TargetType> GetTargetType(llvm::StringRef key);
 llvm::Optional<LayoutType> GetLayoutType(llvm::StringRef key);
 llvm::Optional<PrecisionType> GetPrecisionType(llvm::StringRef key);
 
-llvm::raw_ostream &operator<<(llvm::raw_ostream &os, TargetType type);
-llvm::raw_ostream &operator<<(llvm::raw_ostream &os, LayoutType type);
-llvm::raw_ostream &operator<<(llvm::raw_ostream &os, PrecisionType type);
+llvm::StringRef GetString(TargetType type);
+llvm::StringRef GetString(LayoutType type);
+llvm::StringRef GetString(PrecisionType type);
 
+template <typename T>
+llvm::raw_ostream &operator<<(llvm::raw_ostream &os, T type) {
+  os << GetString(type);
+  return os;
+}
 }  // end namespace infrt
diff --git a/paddle/infrt/dialect/infrt/infrt_dialect.cc b/paddle/infrt/dialect/infrt/infrt_dialect.cc
index abb60016f9..400e4921c9 100644
--- a/paddle/infrt/dialect/infrt/infrt_dialect.cc
+++ b/paddle/infrt/dialect/infrt/infrt_dialect.cc
@@ -14,6 +14,7 @@
 
 #include "paddle/infrt/dialect/infrt/infrt_dialect.h"
 
+#include <llvm/ADT/TypeSwitch.h>
 #include <mlir/IR/Builders.h>
 #include <mlir/IR/BuiltinOps.h>
 #include <mlir/IR/DialectImplementation.h>
diff --git a/paddle/infrt/dialect/infrt/infrt_ops_base.td b/paddle/infrt/dialect/infrt/infrt_ops_base.td
index f19912dc0c..8a6eb76656 100644
--- a/paddle/infrt/dialect/infrt/infrt_ops_base.td
+++ b/paddle/infrt/dialect/infrt/infrt_ops_base.td
@@ -10,16 +10,59 @@ def Infrt_Dialect : Dialect {
 
   let name = "infrt";
   let cppNamespace = "::infrt";
+  let useDefaultAttributePrinterParser = 1;
 }
 
 // Type definitions
-
 // Base class for Infrt dialect types.
 class Infrt_Type<string name, list<Trait> traits = [],
                    string baseCppClass = "::mlir::Type">
     : TypeDef<Infrt_Dialect, name, traits, baseCppClass> {
 }
 
+class Infrt_EnumParam<string cppEnumType, string stringToSymbolFnName,
+  string symbolToStringFnName, string desc = ""> : TypeParameter<cppEnumType, desc> {
+  let parser = [{[&]() -> ::mlir::FailureOr<}] # cppEnumType # [{> {
+    ::llvm::StringRef enumKeyword;
+    if (::mlir::failed($_parser.parseKeyword(&enumKeyword)))
+      return ::mlir::failure();
+    auto maybeEnum = }] # stringToSymbolFnName # [{(enumKeyword);
+    if (maybeEnum)
+      return *maybeEnum;
+    llvm_unreachable("}] # cppEnumType # [{ can not be found.");
+    return {};
+  }()}];
+  let printer = "$_printer << " # symbolToStringFnName # "($_self)";
+}
+
+def TargetParam : Infrt_EnumParam<"::infrt::TargetType", "GetTargetType", "GetString">;
+def PrecisionParam : Infrt_EnumParam<"::infrt::PrecisionType", "GetPrecisionType", "GetString">;
+def LayoutParam : Infrt_EnumParam<"::infrt::LayoutType", "GetLayoutType", "GetString">;
+
+def TargetAttr : AttrDef<Infrt_Dialect, "Target"> {
+  let mnemonic = "target";
+  let parameters = (ins
+    TargetParam:$target
+  );
+  let assemblyFormat = "`<` $target `>`";
+}
+
+def PrecisionAttr : AttrDef<Infrt_Dialect, "Precision"> {
+  let mnemonic = "precision";
+  let parameters = (ins
+    PrecisionParam:$precision
+  );
+  let assemblyFormat = "`<` $precision `>`";
+}
+
+def LayoutAttr : AttrDef<Infrt_Dialect, "Layout"> {
+  let mnemonic = "layout";
+  let parameters = (ins
+    LayoutParam:$layout
+  );
+  let assemblyFormat = "`<` $layout `>`";
+}
+
 def LoDTensor : Infrt_Type<"LoDTensor"> {
   let summary = "infrt lod tensor";
   let description = [{lod_tensor<3x64x3x3xf32, 3>}];
@@ -37,7 +80,6 @@ def DenseTensor : Infrt_Type<"DenseTensor"> {
     "::infrt::TargetType":$target,
     "::infrt::PrecisionType":$precision,
     "::infrt::LayoutType":$layout
-
   );
 }
 
diff --git a/paddle/infrt/dialect/phi/ir/infrt_phi_base.td b/paddle/infrt/dialect/phi/ir/infrt_phi_base.td
index 907f912d9e..e9591e7f6d 100644
--- a/paddle/infrt/dialect/phi/ir/infrt_phi_base.td
+++ b/paddle/infrt/dialect/phi/ir/infrt_phi_base.td
@@ -2,6 +2,7 @@
 #define PHI_BASE
 
 include "mlir/IR/OpBase.td"
+include "paddle/infrt/dialect/infrt_base.td"
 
 def PHI_Dialect : Dialect {
   let name = "phi";
@@ -11,27 +12,28 @@ def PHI_Dialect : Dialect {
   }];
 
   let cppNamespace = "::infrt::phi";
-}
-
-class AllocatorTypeOf<string place, list<Trait> traits=[]>:
-    TypeDef<PHI_Dialect, place # "Allocator", traits> {
-    let summary = !strconcat("!phi.allocator_", place, " type");
-}
-
-class ContextTypeOf<string place, list<Trait> traits=[]>:
-    TypeDef<PHI_Dialect, place # "Context", traits> {
-    let summary = !strconcat("!phi.context_", place, " type");
+  let useDefaultTypePrinterParser = 1;
 }
 
 def PhiOpTrait : NativeOpTrait<"PhiOpTrait">;
 
-def CPU_Allocator : AllocatorTypeOf<"CPU">;
-def GPU_Allocator : AllocatorTypeOf<"GPU">;
-
-def CPU_Context : ContextTypeOf<"CPU">;
-def GPU_Context : ContextTypeOf<"GPU">;
-
-def Allocator : AnyTypeOf<[CPU_Allocator, GPU_Allocator], "Allocator type">;
-def Context : AnyTypeOf<[CPU_Context, GPU_Context], "Context type">;
+class PHI_Type<string type, list<Trait> traits = []>
+   : TypeDef<PHI_Dialect, type, !listconcat(traits, [PhiOpTrait, IsolatedFromAbove])> {}
+
+def Allocator : PHI_Type<"Allocator"> {
+   let mnemonic = "allocator";
+   let parameters = (ins
+     TargetParam:$target
+   );
+   let assemblyFormat = "`<` $target `>`";
+ }
+
+ def Context : PHI_Type<"Context"> {
+   let mnemonic = "context";
+   let parameters = (ins
+     TargetParam:$target
+   );
+   let assemblyFormat = "`<` $target `>`";
+ }
 
 #endif
diff --git a/paddle/infrt/dialect/phi/ir/infrt_phi_tensor.td b/paddle/infrt/dialect/phi/ir/infrt_phi_tensor.td
index 39677871ff..3399c408d9 100644
--- a/paddle/infrt/dialect/phi/ir/infrt_phi_tensor.td
+++ b/paddle/infrt/dialect/phi/ir/infrt_phi_tensor.td
@@ -23,7 +23,7 @@ class PDT_Op<string mnemonic, list<OpTrait> traits = []> : Op<PHI_DenseTensorDia
 
 class CreateDenseTensorOp<string place, string dtype, string layout> 
       : PDT_Op<"create_dense_tensor." # place # "." # dtype # "." # layout, [NoSideEffect]> {
-  let arguments = (ins CPU_Allocator:$allocator, I64ArrayAttr:$dims, I64ArrayAttr:$lod);
+  let arguments = (ins Allocator:$allocator, I64ArrayAttr:$dims, I64ArrayAttr:$lod);
   let results = (outs DenseTensor:$output);
 }
 
@@ -47,13 +47,13 @@ class PrintDenseTensorOp:
 class CreateCPUAllocatorOp
       : PDT_Op<"create_allocator." # "cpu", [NoSideEffect]> {
   let arguments = (ins);
-  let results = (outs CPU_Allocator:$output);
+  let results = (outs Allocator:$output);
 }
 
 class CreateCPUContextOp
       : PDT_Op<"create_context." # "cpu", [NoSideEffect]> {
-  let arguments = (ins CPU_Allocator:$input);
-  let results = (outs CPU_Context:$output);
+  let arguments = (ins Allocator:$input);
+  let results = (outs Context:$output);
 }
 
 def PDT_CreateDenseTensorOp_cpu_f32_nchw : CreateDenseTensorOp<"cpu", "f32", "nchw">;
@@ -63,7 +63,7 @@ def PDT_CreateContextOp_cpu : CreateCPUContextOp;
 def PDT_PrintDenseTensor_cpu : PrintDenseTensorOp;
 
 def FakeKernelOp : PDT_Op<"fake_phi_kernel"> {
-  let arguments = (ins CPU_Context:$dev_ctx, DenseTensor:$x, DenseTensor:$y, BoolAttr:$transpose_x, BoolAttr:$transpose_y);
+  let arguments = (ins Context:$dev_ctx, DenseTensor:$x, DenseTensor:$y, BoolAttr:$transpose_x, BoolAttr:$transpose_y);
   let results = (outs DenseTensor:$output);
 }
 
diff --git a/paddle/infrt/dialect/phi/ir/phi_base.cc b/paddle/infrt/dialect/phi/ir/phi_base.cc
index 7a6b3f3f0a..d8095d7f3f 100644
--- a/paddle/infrt/dialect/phi/ir/phi_base.cc
+++ b/paddle/infrt/dialect/phi/ir/phi_base.cc
@@ -14,6 +14,7 @@
 
 #include "paddle/infrt/dialect/phi/ir/phi_base.h"
 
+#include <llvm/include/llvm/ADT/TypeSwitch.h>
 #include <mlir/IR/Builders.h>
 #include <mlir/IR/Dialect.h>
 #include <mlir/IR/DialectImplementation.h>
@@ -27,27 +28,6 @@
 namespace infrt {
 namespace phi {
 
-void PHIDialect::printType(::mlir::Type type,
-                           mlir::DialectAsmPrinter& os) const {
-  if (type.isa<CPUAllocatorType>()) {
-    os << "CPU_Allocator";
-    return;
-  }
-  if (type.isa<GPUAllocatorType>()) {
-    os << "GPU_Allocator";
-    return;
-  }
-  if (type.isa<CPUContextType>()) {
-    os << "CPU_Context";
-    return;
-  }
-  if (type.isa<GPUContextType>()) {
-    os << "GPU_Context";
-    return;
-  }
-  llvm_unreachable("unexpected 'allocator/context' type kind");
-}
-
 void PHIDialect::initialize() {
   addOperations<
 #define GET_OP_LIST
@@ -59,24 +39,6 @@ void PHIDialect::initialize() {
       >();
 }
 
-mlir::Type PHIDialect::parseType(mlir::DialectAsmParser& parser) const {
-  llvm::StringRef keyword;
-  if (parser.parseKeyword(&keyword)) return mlir::Type();
-  if (keyword == "CPU_allocator") {
-    return CPUAllocatorType::get(parser.getContext());
-  } else if (keyword == "GPU_allocator") {
-    return GPUAllocatorType::get(parser.getContext());
-  } else if (keyword == "CPU_context") {
-    return CPUContextType::get(parser.getContext());
-  } else if (keyword == "GPU_context") {
-    return GPUContextType::get(parser.getContext());
-  } else {
-    llvm_unreachable("unexpected 'allocator/context' type kind");
-  }
-
-  return mlir::Type();
-}
-
 }  // namespace phi
 }  // namespace infrt
 
diff --git a/paddle/infrt/dialect/phi/ir/phi_base.h b/paddle/infrt/dialect/phi/ir/phi_base.h
index a08d8229fc..0ea1973a73 100644
--- a/paddle/infrt/dialect/phi/ir/phi_base.h
+++ b/paddle/infrt/dialect/phi/ir/phi_base.h
@@ -18,12 +18,10 @@
 #include <mlir/Interfaces/SideEffectInterfaces.h>
 
 #include <string>
+#include "paddle/infrt/dialect/infrt/common_type.h"
 
 #include "paddle/infrt/dialect/phi/ir/infrt_phi_baseDialect.h.inc"
 
-#define GET_TYPEDEF_CLASSES
-#include "paddle/infrt/dialect/phi/ir/infrt_phi_baseTypes.h.inc"
-
 #define GET_OP_CLASSES
 #include "paddle/infrt/dialect/phi/ir/infrt_phi_base.h.inc"
 
@@ -41,6 +39,9 @@ class PhiOpTrait : public OpTrait::TraitBase<ConcreteType, PhiOpTrait> {
 }  // namespace OpTrait
 }  // namespace mlir
 
+#define GET_TYPEDEF_CLASSES
+#include "paddle/infrt/dialect/phi/ir/infrt_phi_baseTypes.h.inc"
+
 namespace infrt {
 namespace phi {}  // namespace phi
 }  // namespace infrt
diff --git a/paddle/infrt/host_context/kernel_frame.cc b/paddle/infrt/host_context/kernel_frame.cc
index 14e88be4b9..266c145f47 100644
--- a/paddle/infrt/host_context/kernel_frame.cc
+++ b/paddle/infrt/host_context/kernel_frame.cc
@@ -30,28 +30,21 @@ std::ostream& operator<<(std::ostream& os, const KernelFrame& frame) {
 std::string KernelFrame::DumpArgTypes() const {
   std::stringstream ss;
   for (auto* value : GetValues(0, GetNumElements())) {
-    if (value->is_type<bool>()) {
-      ss << "bool (" << &value->get<bool>() << "), ";
-    } else if (value->is_type<tensor::DenseHostTensor>()) {
-      ss << "DenseHostTensor(" << &value->get<tensor::DenseHostTensor>()
-         << "), ";
-    } else if (value->is_type<float>()) {
-      ss << "float(" << &value->get<float>() << "), ";
-    } else if (value->is_type<int>()) {
-      ss << "int(" << &value->get<int>() << "), ";
-    } else if (value->is_type<phi::DenseTensor>()) {
-      ss << "phi::DenseTensor(" << &value->get<phi::DenseTensor>() << "), ";
-    } else if (value->is_type<phi::MetaTensor>()) {
-      ss << "phi::MetaTensor(" << &value->get<phi::MetaTensor>() << "), ";
-    } else if (value->is_type<::phi::CPUContext>()) {
-      ss << "phi::CPUContext(" << &value->get<::phi::CPUContext>() << "), ";
-    } else if (value->is_type<host_context::None>()) {
-      ss << "none(" << &value->get<host_context::None>() << "), ";
-    } else if (value->is_type<backends::CpuPhiContext>()) {
-      ss << "CpuPhiContext(" << &value->get<backends::CpuPhiContext>() << "), ";
-    } else {
-      ss << "typeid: " << value->index() << ", ";
-    }
+#define DUMP(type_name)                                    \
+  if (value->is_type<type_name>()) {                       \
+    ss << #type_name << &value->get<type_name>() << "), "; \
+  }
+    DUMP(bool);
+    DUMP(tensor::DenseHostTensor);
+    DUMP(float);
+    DUMP(int);
+    DUMP(::phi::DenseTensor);
+    DUMP(::phi::MetaTensor);
+    DUMP(::phi::CPUContext);
+    DUMP(host_context::None);
+    DUMP(backends::CpuPhiContext);
+#undef DUMP
+    ss << "typeid: " << value->index() << ", ";
   }
   return ss.str();
 }
diff --git a/paddle/infrt/kernel/phi/infershaped/phi_kernel_launcher.h b/paddle/infrt/kernel/phi/infershaped/phi_kernel_launcher.h
index a0a5b391ea..75c9e55477 100644
--- a/paddle/infrt/kernel/phi/infershaped/phi_kernel_launcher.h
+++ b/paddle/infrt/kernel/phi/infershaped/phi_kernel_launcher.h
@@ -24,26 +24,6 @@
 namespace infrt {
 namespace kernel {
 
-static void FakePhiInferShape(const ::phi::MetaTensor& a,
-                              const ::phi::MetaTensor& b,
-                              bool arg_0,
-                              bool arg_1,
-                              ::phi::MetaTensor* c) {
-  LOG(INFO) << "the ptr of c: " << c;
-  LOG(INFO) << "c->numel(): " << c->numel();
-}
-
-static void FakePhiKernel(const ::phi::CPUContext& /*Context*/,
-                          const ::phi::DenseTensor& a,
-                          const ::phi::DenseTensor& b,
-                          bool arg_0,
-                          bool arg_1,
-                          ::phi::DenseTensor* c) {
-  std::cout << "@FakePhiKernel@" << std::endl;
-  LOG(INFO) << "the ptr of c: " << c;
-  LOG(INFO) << "c->numel(): " << c->numel();
-}
-
 template <typename KernelFunc,
           KernelFunc kernel,
           typename InferShapedFunc,
diff --git a/paddle/infrt/kernel/phi/registry.cc b/paddle/infrt/kernel/phi/registry.cc
index 15e2d21005..cb09275c17 100644
--- a/paddle/infrt/kernel/phi/registry.cc
+++ b/paddle/infrt/kernel/phi/registry.cc
@@ -44,17 +44,6 @@ void RegisterPhiKernels(host_context::KernelRegistry* registry) {
                       INFRT_KERNEL(infrt::kernel::phi::FillDenseTensorF32));
   registry->AddKernel("phi_dt.print_tensor",
                       INFRT_KERNEL(infrt::kernel::phi::PrintDenseTensor));
-  registry->AddKernel(
-      "phi_dt.fake_phi_kernel",
-      std::bind(&KernelLauncherFunc<decltype(&FakePhiKernel),
-                                    &FakePhiKernel,
-                                    decltype(&FakePhiInferShape),
-                                    &FakePhiInferShape>,
-                KernelLauncher<decltype(&FakePhiKernel),
-                               &FakePhiKernel,
-                               decltype(&FakePhiInferShape),
-                               &FakePhiInferShape>(),
-                std::placeholders::_1));
 }
 
 }  // namespace kernel
diff --git a/paddle/infrt/tests/dialect/pten/dense_tensor.mlir b/paddle/infrt/tests/dialect/pten/dense_tensor.mlir
index 695143c93b..586af7a9c5 100644
--- a/paddle/infrt/tests/dialect/pten/dense_tensor.mlir
+++ b/paddle/infrt/tests/dialect/pten/dense_tensor.mlir
@@ -2,11 +2,11 @@
 
 // CHECK-LABEL: @sign_any_float32_execute
 func @sign_any_float32_execute() {
-  %allocator = "phi_dt.create_allocator.cpu" (): () -> !phi.CPU_allocator
-  %ctx = "phi_dt.create_context.cpu" (%allocator): (!phi.CPU_allocator) -> !phi.CPU_context
-  %t = "phi_dt.create_dense_tensor.cpu.f32.nchw" (%allocator) {dims=[1:i64], lod=[1:i64]}: (!phi.CPU_allocator) -> (!infrt.dense_tensor<CPU, FP32, NCHW>)
+  %allocator = "phi_dt.create_allocator.cpu" (): () -> !phi.allocator<CPU>
+  %ctx = "phi_dt.create_context.cpu" (%allocator): (!phi.allocator<CPU>) -> !phi.context<CPU>
+  %t = "phi_dt.create_dense_tensor.cpu.f32.nchw" (%allocator) {dims=[1:i64], lod=[1:i64]}: (!phi.allocator<CPU>) -> (!infrt.dense_tensor<CPU, FP32, NCHW>)
   "phi_dt.fill_dense_tensor.f32"(%t) {value=[3.8:f32]} : (!infrt.dense_tensor<CPU, FP32, NCHW>) -> ()
-  %e = "phi_cpu.sign.any.float32"(%ctx, %t) : (!phi.CPU_context, !infrt.dense_tensor<CPU, FP32, NCHW>) -> (!infrt.dense_tensor<CPU, FP32, NCHW>)
+  %e = "phi_cpu.sign.any.float32"(%ctx, %t) : (!phi.context<CPU>, !infrt.dense_tensor<CPU, FP32, NCHW>) -> (!infrt.dense_tensor<CPU, FP32, NCHW>)
 
   // CHECK: dense_tensor: shape=shape[1], values=[1]
   "phi_dt.print_tensor" (%e) : (!infrt.dense_tensor<CPU, FP32, NCHW>) -> ()
diff --git a/tools/infrt/generate_phi_kernel_dialect.py b/tools/infrt/generate_phi_kernel_dialect.py
index 80cf3958b1..8efa03306f 100644
--- a/tools/infrt/generate_phi_kernel_dialect.py
+++ b/tools/infrt/generate_phi_kernel_dialect.py
@@ -95,7 +95,7 @@ def generate_inputs_info(input_info):
 def generate_arguments_info(op_name, input_info, attr_info):
     input_args = generate_inputs_info(input_info)
     attr_args = generate_attrs_info(op_name, attr_info)
-    context_args = "CPU_Context:$dev_ctx"
+    context_args = "Context:$dev_ctx"
     argument_ = "{},{},{}".format(context_args, input_args, attr_args)
     return (("let arguments = (ins {});".format(argument_.strip(","))))
 
-- 
GitLab