Kokkos-aware Clad #783

kliegeois · 2024-02-22T20:12:20Z

This PR provides the implementation of a first set of features for the automatic generation of gradients of Kokkos-based code.

The content of this PR supports reverse mode of Kokkos parallel_for, deep_copy, and view accesses.

@vgvassilev @brian-kelley

…ide the forward pass

…zation problem with a simple algorithm

…tions

github-actions

clang-tidy made some suggestions

There were too many comments to post at once. Showing the first 10 out of 136. Check the log or trigger a new build to see more.

github-actions · 2024-02-22T21:09:28Z

include/clad/Differentiator/CladUtils.h

@@ -174,6 +179,22 @@ namespace clad {
    /// otherwise returns false.
    bool HasAnyReferenceOrPointerArgument(const clang::FunctionDecl* FD);

+    /// Returns true if `constructedTypeName` is a string describing Kokkos::TeamPolicy type.
+    bool IsKokkosTeamPolicy(const std::string constructedTypeName);


warning: parameter 'constructedTypeName' is const-qualified in the function declaration; const-qualification of parameters only has an effect in function definitions [readability-avoid-const-params-in-decls]

Suggested change

bool IsKokkosTeamPolicy(const std::string constructedTypeName);

bool IsKokkosTeamPolicy(std::string constructedTypeName);

github-actions · 2024-02-22T21:09:28Z

include/clad/Differentiator/CladUtils.h

+
+    /// Returns true if `constructedTypeName` is a string describing Kokkos::TeamThreadRange,
+    /// Kokkos::ThreadVectorRange, or Kokkos::TeamVectorRange type.
+    bool IsKokkosRange(const std::string constructedTypeName);


warning: parameter 'constructedTypeName' is const-qualified in the function declaration; const-qualification of parameters only has an effect in function definitions [readability-avoid-const-params-in-decls]

Suggested change

bool IsKokkosRange(const std::string constructedTypeName);

bool IsKokkosRange(std::string constructedTypeName);

github-actions · 2024-02-22T21:09:28Z

include/clad/Differentiator/CladUtils.h

+    bool IsKokkosRange(const std::string constructedTypeName);
+
+    /// Returns true if `constructedTypeName` is a string describing Kokkos::Member type.
+    bool IsKokkosMember(const std::string constructedTypeName);


warning: parameter 'constructedTypeName' is const-qualified in the function declaration; const-qualification of parameters only has an effect in function definitions [readability-avoid-const-params-in-decls]

Suggested change

bool IsKokkosMember(const std::string constructedTypeName);

bool IsKokkosMember(std::string constructedTypeName);

github-actions · 2024-02-22T21:09:29Z

include/clad/Differentiator/CladUtils.h

+    bool IsKokkosMember(const std::string constructedTypeName);
+
+    /// Returns true if `constructedTypeName` is a string describing Kokkos::View type.
+    bool IsKokkosView(const std::string constructedTypeName);


warning: parameter 'constructedTypeName' is const-qualified in the function declaration; const-qualification of parameters only has an effect in function definitions [readability-avoid-const-params-in-decls]

Suggested change

bool IsKokkosView(const std::string constructedTypeName);

bool IsKokkosView(std::string constructedTypeName);

github-actions · 2024-02-22T21:09:29Z

include/clad/Differentiator/KokkosViewAccessVisitor.h

+#ifndef CLAD_KOKKOS_VIEW_ACCESS_VISITOR_H
+#define CLAD_KOKKOS_VIEW_ACCESS_VISITOR_H


warning: header guard does not follow preferred style [llvm-header-guard]

Suggested change

#ifndef CLAD_KOKKOS_VIEW_ACCESS_VISITOR_H

#define CLAD_KOKKOS_VIEW_ACCESS_VISITOR_H

#ifndef CLAD_DIFFERENTIATOR_KOKKOSVIEWACCESSVISITOR_H

#define CLAD_DIFFERENTIATOR_KOKKOSVIEWACCESSVISITOR_H

include/clad/Differentiator/KokkosViewAccessVisitor.h:423:

- #endif // CLAD_KOKKOS_VIEW_ACCESS_VISITOR_H + #endif // CLAD_DIFFERENTIATOR_KOKKOSVIEWACCESSVISITOR_H

github-actions · 2024-02-22T21:09:29Z

include/clad/Differentiator/KokkosViewAccessVisitor.h

+  using namespace clang;
+
+  if (!Stmt1 || !Stmt2) {
+    return !Stmt1 && !Stmt2;


warning: implicit conversion 'const clang::Stmt *' -> bool [readability-implicit-bool-conversion]

Suggested change

return !Stmt1 && !Stmt2;

return (Stmt1 == nullptr) && !Stmt2;

github-actions · 2024-02-22T21:09:29Z

include/clad/Differentiator/KokkosViewAccessVisitor.h

+  using namespace clang;
+
+  if (!Stmt1 || !Stmt2) {
+    return !Stmt1 && !Stmt2;


warning: implicit conversion 'const clang::Stmt *' -> bool [readability-implicit-bool-conversion]

Suggested change

return !Stmt1 && !Stmt2;

return !Stmt1 && (Stmt2 == nullptr);

github-actions · 2024-02-22T21:09:29Z

include/clad/Differentiator/KokkosViewAccessVisitor.h

+  case Stmt::NullStmtClass:
+    return true;
+  case Stmt::CStyleCastExprClass: {
+    const CStyleCastExpr* CastExpr1 = cast<CStyleCastExpr>(Stmt1);


warning: use auto when initializing with a template cast to avoid duplicating the type name [modernize-use-auto]

Suggested change

const CStyleCastExpr* CastExpr1 = cast<CStyleCastExpr>(Stmt1);

const auto* CastExpr1 = cast<CStyleCastExpr>(Stmt1);

github-actions · 2024-02-22T21:09:29Z

include/clad/Differentiator/KokkosViewAccessVisitor.h

+    return true;
+  case Stmt::CStyleCastExprClass: {
+    const CStyleCastExpr* CastExpr1 = cast<CStyleCastExpr>(Stmt1);
+    const CStyleCastExpr* CastExpr2 = cast<CStyleCastExpr>(Stmt2);


warning: use auto when initializing with a template cast to avoid duplicating the type name [modernize-use-auto]

Suggested change

const CStyleCastExpr* CastExpr2 = cast<CStyleCastExpr>(Stmt2);

const auto* CastExpr2 = cast<CStyleCastExpr>(Stmt2);

github-actions · 2024-02-22T21:09:29Z

include/clad/Differentiator/KokkosViewAccessVisitor.h

+    return CastExpr1->getTypeAsWritten() == CastExpr2->getTypeAsWritten();
+  }
+  case Stmt::ReturnStmtClass: {
+    const ReturnStmt *ReturnStmt1 = cast<ReturnStmt>(Stmt1);


warning: use auto when initializing with a template cast to avoid duplicating the type name [modernize-use-auto]

Suggested change

const ReturnStmt *ReturnStmt1 = cast<ReturnStmt>(Stmt1);

const auto *ReturnStmt1 = cast<ReturnStmt>(Stmt1);

parth-07 · 2024-03-03T08:49:58Z

include/clad/Differentiator/CladUtils.h

@@ -153,6 +153,11 @@ namespace clad {
    clang::DeclContext* GetOutermostDC(clang::Sema& semaRef,
                                       clang::DeclContext* DC);

+    clang::Expr* GetUnresolvedLookup(clang::Sema& semaRef,


ASTContext can be obtained through Sema. We don't need to pass both of them to a function.

parth-07 · 2024-03-03T08:50:29Z

include/clad/Differentiator/CladUtils.h

@@ -174,6 +179,22 @@ namespace clad {
    /// otherwise returns false.
    bool HasAnyReferenceOrPointerArgument(const clang::FunctionDecl* FD);

+    /// Returns true if `constructedTypeName` is a string describing Kokkos::TeamPolicy type.
+    bool IsKokkosTeamPolicy(const std::string constructedTypeName);


Perhaps kokkos specific functions should be in a separate namespace altogether?

parth-07 · 2024-03-03T08:52:12Z

include/clad/Differentiator/CladUtils.h

+
+    /// Returns true if `constructedTypeName` is a string describing Kokkos::TeamThreadRange,
+    /// Kokkos::ThreadVectorRange, or Kokkos::TeamVectorRange type.
+    bool IsKokkosRange(const std::string constructedTypeName);


We should pass std::string as const std::string& instead of const std::string. const in parameters not being passed as reference is meaningless for the caller.

parth-07 · 2024-03-03T08:54:23Z

lib/Differentiator/BaseForwardModeVisitor.cpp

+    auto MCE = dyn_cast<clang::CXXMemberCallExpr>(CE);
+
+    if (utils::IsKokkosView(MCE->getObjectType().getAsString())) {
+      //Member function called from a Kokkos::View; nothing to do here


For supporting things such as this, we should use custom derivatives instead of modifying the core Clad code. Modifying core clad is not a scalable solution.

parth-07 · 2024-03-03T08:54:58Z

lib/Differentiator/BaseForwardModeVisitor.cpp

+      return StmtDiff(Clone(CE));
+    }
+  }
+  if (isa<CXXOperatorCallExpr>(CE)) {


This should be implemented as custom derivative as well.

parth-07 · 2024-03-03T08:56:24Z

lib/Differentiator/BaseForwardModeVisitor.cpp

+    }
+
+    // Returning the function call and zero derivative
+    if (isKokkosViewAccess) {


All these changes can be implemented with more consistency and much less efforts using the custom derivative functionality.

parth-07 · 2024-03-03T08:59:17Z

lib/Differentiator/ReverseModeVisitor.cpp

@@ -99,7 +100,9 @@ Expr* getArraySizeExpr(const ArrayType* AT, ASTContext& context,
  }

  ReverseModeVisitor::ReverseModeVisitor(DerivativeBuilder& builder)
-      : VisitorBase(builder), m_Result(nullptr) {}
+      : VisitorBase(builder), m_Result(nullptr) {
+      m_KVAV = new KokkosViewAccessVisitor(m_Sema, m_Context);


It seems that we are unconditionally creating and deleting m_KVAV and we exactly know when to create and delete it. Why does m_KVAV needs to be a pointer? Why cannot we have m_KVAV as non-pointer instance member?

This commit adds a test from the original Kokkos PR vgvassilev#783 and provides the implementation of the simplest Kokkos features in the form of custom pushforwards. This commit only addresses the forward mode.

This commit adds a test from the original Kokkos PR #783 and provides the implementation of the simplest Kokkos features in the form of custom pushforwards. This commit only addresses the forward mode.

vgvassilev · 2024-07-23T17:23:30Z

@kliegeois, @brian-kelley, is there a point to try to rebase this PR or @gojakuch should continue picking random things from it as it is now?

kliegeois · 2024-07-23T19:40:41Z

@vgvassilev I am not sure what is best.
If we rebase, would @gojakuch still create different PRs that refactor this PR?
If yes, the rebase is not necessarily useful although it will be easier to highlight the remaining pieces.

kliegeois added 30 commits February 22, 2024 12:36

Modify Clad to use Kokkos views in forward mode

082bb20

Add parallel_for

9b9612e

Add Kokkos subviews

9284719

Add Kokkos examples

65557a9

Add CMakeLists.txt for the Kokkos example and push the generated file

322ca21

Start to work on the reverse mode

02dd2ce

Propagate derivatives inside Kokkos view

7e9c540

Modify the test to use the generated file

2cd5231

Add a first implementation of a deep_copy

db086a5

Update the implementation of the deep_copy

ef1d117

Reset _d_a to zero after a reverse deep_copy

43fa4d6

Fix generated deep_copy by removing not needed arguments

092f425

Clean a bit of the copy paste

1eb525e

Move towards the usage of kokkos_builtin_derivative::parallel_sum ins…

91ec308

…ide the forward pass

Move towards the usage of subview

f87ccd2

Update the isKokkosView logic

ee88c2e

Use the fact that a Kokkos view is passed by reference

1118bbe

Add a gradient of a 2D->double function

3fe3122

Use the gradient of the 2D-> double function to solve a simple optimi…

d6c4633

…zation problem with a simple algorithm

Reenable a forward test

814e46d

Clean files

aef1422

Add a python script to postprocess the generated Derivative.cpp

230baba

update the gradient descent test

f5b2f4b

add timers

09956bc

Modify the test to be run on GPU and generalize the parallel_sum func…

b7f2d71

…tions

Use more template type member to deduce some of the temporary types

09d3dec

Add create_mirror_view

327c88a

Use f_view in double f(double x, double y)

ddd5763

Use a partial template instanciation of ViewSum and ViewAdd

d01770d

Modify ViewSum and ViewAdd to be compiled with CUDA

36aec0a

kliegeois added 2 commits February 22, 2024 12:37

Add comments

816cfad

link Kokkos::kokkos

3a585dd

github-actions bot reviewed Feb 22, 2024

View reviewed changes

parth-07 requested changes Mar 3, 2024

View reviewed changes

vgvassilev added this to the v1.6 milestone Mar 7, 2024

gojakuch mentioned this pull request Mar 18, 2024

Add Kokkos unittests #826

Merged

vgvassilev modified the milestones: v1.6, v1.7 Jul 6, 2024

gojakuch mentioned this pull request Jul 10, 2024

Add basic Kokkos support for the original tests of #783 #977

Merged

vgvassilev modified the milestones: v1.7, 1.8 Jul 18, 2024

vgvassilev modified the milestones: v1.8, v1.9 Dec 10, 2024

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Kokkos-aware Clad #783

Kokkos-aware Clad #783

kliegeois commented Feb 22, 2024

github-actions bot left a comment

github-actions bot Feb 22, 2024

github-actions bot Feb 22, 2024

github-actions bot Feb 22, 2024

github-actions bot Feb 22, 2024

github-actions bot Feb 22, 2024

github-actions bot Feb 22, 2024

github-actions bot Feb 22, 2024

github-actions bot Feb 22, 2024

github-actions bot Feb 22, 2024

github-actions bot Feb 22, 2024

parth-07 Mar 3, 2024

parth-07 Mar 3, 2024

parth-07 Mar 3, 2024

parth-07 Mar 3, 2024

parth-07 Mar 3, 2024

parth-07 Mar 3, 2024

parth-07 Mar 3, 2024

vgvassilev commented Jul 23, 2024

kliegeois commented Jul 23, 2024

	bool IsKokkosTeamPolicy(const std::string constructedTypeName);
	bool IsKokkosTeamPolicy(std::string constructedTypeName);

	bool IsKokkosRange(const std::string constructedTypeName);
	bool IsKokkosRange(std::string constructedTypeName);

	bool IsKokkosMember(const std::string constructedTypeName);
	bool IsKokkosMember(std::string constructedTypeName);

	bool IsKokkosView(const std::string constructedTypeName);
	bool IsKokkosView(std::string constructedTypeName);

		#ifndef CLAD_KOKKOS_VIEW_ACCESS_VISITOR_H
		#define CLAD_KOKKOS_VIEW_ACCESS_VISITOR_H

	return !Stmt1 && !Stmt2;
	return (Stmt1 == nullptr) && !Stmt2;

	return !Stmt1 && !Stmt2;
	return !Stmt1 && (Stmt2 == nullptr);

	const CStyleCastExpr* CastExpr1 = cast<CStyleCastExpr>(Stmt1);
	const auto* CastExpr1 = cast<CStyleCastExpr>(Stmt1);

	const CStyleCastExpr* CastExpr2 = cast<CStyleCastExpr>(Stmt2);
	const auto* CastExpr2 = cast<CStyleCastExpr>(Stmt2);

	const ReturnStmt *ReturnStmt1 = cast<ReturnStmt>(Stmt1);
	const auto *ReturnStmt1 = cast<ReturnStmt>(Stmt1);

Kokkos-aware Clad #783

Are you sure you want to change the base?

Kokkos-aware Clad #783

Conversation

kliegeois commented Feb 22, 2024

github-actions bot left a comment

Choose a reason for hiding this comment

github-actions bot Feb 22, 2024

Choose a reason for hiding this comment

github-actions bot Feb 22, 2024

Choose a reason for hiding this comment

github-actions bot Feb 22, 2024

Choose a reason for hiding this comment

github-actions bot Feb 22, 2024

Choose a reason for hiding this comment

github-actions bot Feb 22, 2024

Choose a reason for hiding this comment

github-actions bot Feb 22, 2024

Choose a reason for hiding this comment

github-actions bot Feb 22, 2024

Choose a reason for hiding this comment

github-actions bot Feb 22, 2024

Choose a reason for hiding this comment

github-actions bot Feb 22, 2024

Choose a reason for hiding this comment

github-actions bot Feb 22, 2024

Choose a reason for hiding this comment

parth-07 Mar 3, 2024

Choose a reason for hiding this comment

parth-07 Mar 3, 2024

Choose a reason for hiding this comment

parth-07 Mar 3, 2024

Choose a reason for hiding this comment

parth-07 Mar 3, 2024

Choose a reason for hiding this comment

parth-07 Mar 3, 2024

Choose a reason for hiding this comment

parth-07 Mar 3, 2024

Choose a reason for hiding this comment

parth-07 Mar 3, 2024

Choose a reason for hiding this comment

vgvassilev commented Jul 23, 2024

kliegeois commented Jul 23, 2024