experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

[KERNEL_GEN] Fix func type for DeallocRawOp in TF Framework.

Browse Source 
PiperOrigin-RevId: 335374599
Change-Id: I2fd1782e0b4881834861ade76f3acbebafdd631a
This commit is contained in:
Alexander Belyaev 2020-10-05 02:35:52 -07:00 committed by TensorFlower Gardener
parent 5510653124
commit 45654084ab
2 changed files with 3 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
tensorflow/compiler/mlir/tools/kernel_gen/tests/tf_framework_legalize_to_llvm.mlir
View File

@ -55,7 +55,7 @@ func @alloc_raw(%ctx: !tf_framework.op_kernel_context,
// -----
// CHECK: llvm.func @_mlir_ciface_tf_dealloc_raw(!llvm.ptr<i8>)
// CHECK: llvm.func @_mlir_ciface_tf_dealloc_raw(!llvm.ptr<i8>, !llvm.ptr<i8>)
// CHECK-LABEL: llvm.func @dealloc_raw(
// CHECK-SAME: [[TF_CTX:%.*]]: !llvm.ptr<i8>,

3
tensorflow/compiler/mlir/tools/kernel_gen/transforms/tf_framework_legalize_to_llvm.cc
View File

@ -171,7 +171,8 @@ class DeallocRawOpConverter : public ConvertToLLVMCallOpPattern<DeallocRawOp> {
protected:
StringRef GetFuncName() const override { return kCInterfaceDealloc; }
LLVMType GetFuncType() const override {
return LLVM::LLVMType::getFunctionTy(getVoidType(), getVoidPtrType(),
return LLVM::LLVMType::getFunctionTy(getVoidType(),
{getVoidPtrType(), getVoidPtrType()},
/*isVarArg=*/false);
}
};