Skip to content

Commit fb7b0e6

Browse files
Smit-createSmit-create
authored
Merge pull request #421 from Smit-create/complex_asr
Add `ComplexRe` and `ComplexIm` in ASR
2 parents 0e43498 + de63a21 commit fb7b0e6

File tree

6 files changed

+95
-13
lines changed

6 files changed

+95
-13
lines changed

‎src/libasr/ASR.asdl

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -238,6 +238,9 @@ expr
238238
| ArrayRef(symbol v, array_index* args, ttype type, expr? value)
239239
| DerivedRef(expr v, symbol m, ttype type, expr? value)
240240
| Cast(expr arg, cast_kind kind, ttype type, expr? value)
241+
| ComplexRe(expr arg, ttype type, expr? value)
242+
| ComplexIm(expr arg, ttype type, expr? value)
243+
241244

242245
-- `len` in Character:
243246
-- >=0 ... the length of the string, known at compile time

‎src/libasr/asr_utils.h

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -120,6 +120,8 @@ static inline ASR::ttype_t* expr_type(const ASR::expr_t *f)
120120
case ASR::exprType::ArrayRef: { return ((ASR::ArrayRef_t*)f)->m_type; }
121121
case ASR::exprType::DerivedRef: { return ((ASR::DerivedRef_t*)f)->m_type; }
122122
case ASR::exprType::Cast: { return ((ASR::Cast_t*)f)->m_type; }
123+
case ASR::exprType::ComplexRe: { return ((ASR::ComplexRe_t*)f)->m_type; }
124+
case ASR::exprType::ComplexIm: { return ((ASR::ComplexIm_t*)f)->m_type; }
123125
default : throw LFortranException("Not implemented");
124126
}
125127
}
@@ -268,6 +270,8 @@ static inline ASR::expr_t* expr_value(ASR::expr_t *f)
268270
case ASR::exprType::Var: { return EXPR2VAR(f)->m_value; }
269271
case ASR::exprType::StrOp: { return ASR::down_cast<ASR::StrOp_t>(f)->m_value; }
270272
case ASR::exprType::ImpliedDoLoop: { return ASR::down_cast<ASR::ImpliedDoLoop_t>(f)->m_value; }
273+
case ASR::exprType::ComplexRe: { return ASR::down_cast<ASR::ComplexRe_t>(f)->m_value; }
274+
case ASR::exprType::ComplexIm: { return ASR::down_cast<ASR::ComplexIm_t>(f)->m_value; }
271275
case ASR::exprType::ArrayConstant: // Drop through
272276
case ASR::exprType::IntegerConstant: // Drop through
273277
case ASR::exprType::RealConstant: // Drop through

‎src/libasr/codegen/asr_to_llvm.cpp

Lines changed: 67 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3358,6 +3358,73 @@ class ASRToLLVMVisitor : public ASR::BaseVisitor<ASRToLLVMVisitor>
33583358
arg_kind = ASRUtils::extract_kind_from_ttype_t(curr_type);
33593359
}
33603360

3361+
void visit_ComplexRe(const ASR::ComplexRe_t &x) {
3362+
this->visit_expr_wrapper(x.m_arg, true);
3363+
ASR::ttype_t* curr_type = extract_ttype_t_from_expr(x.m_arg);
3364+
int arg_kind = ASRUtils::extract_kind_from_ttype_t(curr_type);
3365+
int dest_kind = ASRUtils::extract_kind_from_ttype_t(x.m_type);
3366+
llvm::Value *re;
3367+
if (arg_kind == 4 && dest_kind == 4) {
3368+
// complex(4) -> real(4)
3369+
re = complex_re(tmp, complex_type_4);
3370+
tmp = re;
3371+
} else if (arg_kind == 4 && dest_kind == 8) {
3372+
// complex(4) -> real(8)
3373+
re = complex_re(tmp, complex_type_4);
3374+
tmp = builder->CreateFPExt(re, llvm::Type::getDoubleTy(context));
3375+
} else if (arg_kind == 8 && dest_kind == 4) {
3376+
// complex(8) -> real(4)
3377+
re = complex_re(tmp, complex_type_8);
3378+
tmp = builder->CreateFPTrunc(re, llvm::Type::getFloatTy(context));
3379+
} else if (arg_kind == 8 && dest_kind == 8) {
3380+
// complex(8) -> real(8)
3381+
re = complex_re(tmp, complex_type_8);
3382+
tmp = re;
3383+
} else {
3384+
std::string msg = "Conversion from " + std::to_string(arg_kind) +
3385+
" to " + std::to_string(dest_kind) + " not implemented yet.";
3386+
throw CodeGenError(msg);
3387+
}
3388+
}
3389+
3390+
void visit_ComplexIm(const ASR::ComplexIm_t &x) {
3391+
ASR::ttype_t* curr_type = extract_ttype_t_from_expr(x.m_arg);
3392+
int arg_kind = ASRUtils::extract_kind_from_ttype_t(curr_type);
3393+
llvm::Function *fn = nullptr;
3394+
llvm::Type *ret_type = nullptr, *complex_type = nullptr;
3395+
llvm::AllocaInst *arg = nullptr;
3396+
std::string runtime_func_name = "";
3397+
if (arg_kind == 4) {
3398+
runtime_func_name = "_lfortran_complex_aimag_32";
3399+
ret_type = llvm::Type::getFloatTy(context);
3400+
complex_type = complex_type_4;
3401+
arg = builder->CreateAlloca(complex_type_4,
3402+
nullptr);
3403+
} else {
3404+
runtime_func_name = "_lfortran_complex_aimag_64";
3405+
ret_type = llvm::Type::getDoubleTy(context);
3406+
complex_type = complex_type_8;
3407+
arg = builder->CreateAlloca(complex_type_8,
3408+
nullptr);
3409+
}
3410+
fn = module->getFunction(runtime_func_name);
3411+
if (!fn) {
3412+
llvm::FunctionType *function_type = llvm::FunctionType::get(
3413+
llvm::Type::getVoidTy(context), {
3414+
complex_type->getPointerTo(),
3415+
ret_type->getPointerTo(),
3416+
}, true);
3417+
fn = llvm::Function::Create(function_type,
3418+
llvm::Function::ExternalLinkage, runtime_func_name, *module);
3419+
}
3420+
this->visit_expr_wrapper(x.m_arg, true);
3421+
builder->CreateStore(tmp, arg);
3422+
llvm::AllocaInst *result = builder->CreateAlloca(ret_type, nullptr);
3423+
std::vector<llvm::Value*> args = {arg, result};
3424+
builder->CreateCall(fn, args);
3425+
tmp = CreateLoad(result);
3426+
}
3427+
33613428
void visit_Cast(const ASR::Cast_t &x) {
33623429
if (x.m_value) {
33633430
this->visit_expr_wrapper(x.m_value, true);

‎src/lpython/semantics/python_ast_to_asr.cpp

Lines changed: 6 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -2027,23 +2027,20 @@ class BodyVisitor : public CommonVisitor<BodyVisitor> {
20272027
std::string attr = x.m_attr;
20282028
if (attr == "imag") {
20292029
ASR::expr_t *val = ASR::down_cast<ASR::expr_t>(ASR::make_Var_t(al, x.base.base.loc, t));
2030-
ASR::symbol_t *fn_imag = resolve_intrinsic_function(x.base.base.loc, "_lpython_imag");
2031-
Vec<ASR::call_arg_t> args;
2032-
args.reserve(al, 1);
2033-
ASR::call_arg_t arg;
2034-
arg.loc = val->base.loc;
2035-
arg.m_value = val;
2036-
args.push_back(al, arg);
2037-
tmp = make_call_helper(al, fn_imag, current_scope, args, "_lpython_imag", x.base.base.loc);
2030+
int kind = ASRUtils::extract_kind_from_ttype_t(var->m_type);
2031+
ASR::ttype_t *dest_type = ASR::down_cast<ASR::ttype_t>(ASR::make_Real_t(al, x.base.base.loc,
2032+
kind, nullptr, 0));
2033+
tmp = ASR::make_ComplexIm_t(al, x.base.base.loc, val, dest_type, nullptr);
20382034
return;
20392035
} else if (attr == "real") {
20402036
ASR::expr_t *val = ASR::down_cast<ASR::expr_t>(ASR::make_Var_t(al, x.base.base.loc, t));
20412037
int kind = ASRUtils::extract_kind_from_ttype_t(var->m_type);
20422038
ASR::ttype_t *dest_type = ASR::down_cast<ASR::ttype_t>(ASR::make_Real_t(al, x.base.base.loc,
20432039
kind, nullptr, 0));
2044-
tmp = (ASR::asr_t*)ASR::down_cast<ASR::expr_t>(ASR::make_Cast_t(
2040+
ASR::expr_t *value = ASR::down_cast<ASR::expr_t>(ASR::make_Cast_t(
20452041
al, val->base.loc, val, ASR::cast_kindType::ComplexToReal, dest_type,
20462042
nullptr));
2043+
tmp = ASR::make_ComplexRe_t(al, x.base.base.loc, val, dest_type, ASRUtils::expr_value(value));
20472044
return;
20482045
} else {
20492046
throw SemanticError("'" + attr + "' is not implemented for Complex type",

‎src/runtime/impure/lfortran_intrinsics.c

Lines changed: 13 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -193,6 +193,16 @@ LFORTRAN_API double _lfortran_zaimag(double_complex_t x)
193193
return cimag(x);
194194
}
195195

196+
LFORTRAN_API void _lfortran_complex_aimag_32(struct _lfortran_complex_32 *x, float *res)
197+
{
198+
*res = x->im;
199+
}
200+
201+
LFORTRAN_API void _lfortran_complex_aimag_64(struct _lfortran_complex_64 *x, double *res)
202+
{
203+
*res = x->im;
204+
}
205+
196206
// exp -------------------------------------------------------------------------
197207

198208
LFORTRAN_API float _lfortran_sexp(float x)
@@ -857,7 +867,7 @@ LFORTRAN_API int64_t _lpython_open(char *path, char *flags)
857867
{
858868
FILE *fd;
859869
fd = fopen(path, flags);
860-
if (!fd)
870+
if (!fd)
861871
{
862872
printf("Error in opening the file!\n");
863873
perror(path);
@@ -869,7 +879,7 @@ LFORTRAN_API int64_t _lpython_open(char *path, char *flags)
869879
LFORTRAN_API char* _lpython_read(int64_t fd, int64_t n)
870880
{
871881
char *c = (char *) calloc(n, sizeof(char));
872-
if (fd < 0)
882+
if (fd < 0)
873883
{
874884
printf("Error in reading the file!\n");
875885
exit(1);
@@ -881,7 +891,7 @@ LFORTRAN_API char* _lpython_read(int64_t fd, int64_t n)
881891

882892
LFORTRAN_API void _lpython_close(int64_t fd)
883893
{
884-
if (fclose((FILE*)fd) != 0)
894+
if (fclose((FILE*)fd) != 0)
885895
{
886896
printf("Error in closing the file!\n");
887897
exit(1);

‎src/runtime/impure/lfortran_intrinsics.h

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -61,7 +61,8 @@ LFORTRAN_API void _lfortran_complex_div_64(struct _lfortran_complex_64* a,
6161
struct _lfortran_complex_64* b, struct _lfortran_complex_64 *result);
6262
LFORTRAN_API void _lfortran_complex_pow_64(struct _lfortran_complex_64* a,
6363
struct _lfortran_complex_64* b, struct _lfortran_complex_64 *result);
64-
64+
LFORTRAN_API void _lfortran_complex_aimag_32(struct _lfortran_complex_32 *x, float *res);
65+
LFORTRAN_API void _lfortran_complex_aimag_64(struct _lfortran_complex_64 *x, double *res);
6566
LFORTRAN_API float_complex_t _lfortran_csqrt(float_complex_t x);
6667
LFORTRAN_API double_complex_t _lfortran_zsqrt(double_complex_t x);
6768
LFORTRAN_API float _lfortran_caimag(float_complex_t x);

0 commit comments

Comments
 (0)