From 1ba0c7aea80cf92a54add732ec8ea8d4baf2ad52 Mon Sep 17 00:00:00 2001
From: tj-sun <tianjiao.sun.2010@gmail.com>
Date: Tue, 6 Mar 2018 17:30:07 +0000
Subject: [PATCH] add pragma to iname

---
 loopy/kernel/__init__.py   | 8 ++++++++
 loopy/target/c/__init__.py | 9 ++++++++-
 2 files changed, 16 insertions(+), 1 deletion(-)

diff --git a/loopy/kernel/__init__.py b/loopy/kernel/__init__.py
index 550288642..6be38e45c 100644
--- a/loopy/kernel/__init__.py
+++ b/loopy/kernel/__init__.py
@@ -142,6 +142,11 @@ class LoopKernel(ImmutableRecordWithoutPickling):
         A :class:`dict` mapping inames (as strings)
         to instances of :class:`loopy.kernel.data.IndexTag`.
 
+    .. attribute:: iname_to_pragma
+
+        A :class:`dict` mapping inames (as strings)
+        to pragma (as #pragma 'value here').
+
     .. attribute:: function_manglers
     .. attribute:: symbol_manglers
 
@@ -195,6 +200,7 @@ class LoopKernel(ImmutableRecordWithoutPickling):
             local_sizes={},
             temporary_variables={},
             iname_to_tag={},
+            iname_to_pragma={},
             substitutions={},
             function_manglers=[
                 default_function_mangler,
@@ -280,6 +286,7 @@ class LoopKernel(ImmutableRecordWithoutPickling):
                 temporary_variables=temporary_variables,
                 local_sizes=local_sizes,
                 iname_to_tag=iname_to_tag,
+                iname_to_pragma=iname_to_pragma,
                 substitutions=substitutions,
                 cache_manager=cache_manager,
                 applied_iname_rewrites=applied_iname_rewrites,
@@ -1340,6 +1347,7 @@ class LoopKernel(ImmutableRecordWithoutPickling):
             "local_sizes",
             "temporary_variables",
             "iname_to_tag",
+            "iname_to_pragma",
             "substitutions",
             "iname_slab_increments",
             "loop_priority",
diff --git a/loopy/target/c/__init__.py b/loopy/target/c/__init__.py
index 8e69793e8..03cb45224 100644
--- a/loopy/target/c/__init__.py
+++ b/loopy/target/c/__init__.py
@@ -922,7 +922,7 @@ class CASTBuilder(ASTBuilderBase):
         from pymbolic.mapper.stringifier import PREC_NONE
         from cgen import For, InlineInitializer
 
-        return For(
+        loop = For(
                 InlineInitializer(
                     POD(self, iname_dtype, iname),
                     ecm(lbound, PREC_NONE, "i")),
@@ -935,6 +935,13 @@ class CASTBuilder(ASTBuilderBase):
                 "++%s" % iname,
                 inner)
 
+        if iname in codegen_state.kernel.iname_to_pragma:
+            from cgen import Block, Pragma
+            pragma = Pragma(codegen_state.kernel.iname_to_pragma[iname])
+            loop = Block([pragma, loop])
+
+        return loop
+
     def emit_initializer(self, codegen_state, dtype, name, val_str, is_const):
         decl = POD(self, dtype, name)
 
-- 
GitLab