1 /*
2 * Copyright (c) 2025, Oracle and/or its affiliates. All rights reserved.
3 * DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.
4 *
5 * This code is free software; you can redistribute it and/or modify it
6 * under the terms of the GNU General Public License version 2 only, as
7 * published by the Free Software Foundation. Oracle designates this
8 * particular file as subject to the "Classpath" exception as provided
9 * by Oracle in the LICENSE file that accompanied this code.
10 *
11 * This code is distributed in the hope that it will be useful, but WITHOUT
12 * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
13 * FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License
14 * version 2 for more details (a copy is included in the LICENSE file that
15 * accompanied this code).
16 *
17 * You should have received a copy of the GNU General Public License version
18 * 2 along with this work; if not, write to the Free Software Foundation,
19 * Inc., 51 Franklin St, Fifth Floor, Boston, MA 02110-1301 USA.
20 *
21 * Please contact Oracle, 500 Oracle Parkway, Redwood Shores, CA 94065 USA
22 * or visit www.oracle.com if you need additional information or have any
23 * questions.
24 */
25 package hat.dialect;
26
27 import jdk.incubator.code.CodeContext;
28 import jdk.incubator.code.CodeTransformer;
29 import jdk.incubator.code.Op;
30 import jdk.incubator.code.TypeElement;
31 import jdk.incubator.code.Value;
32 import optkl.util.ops.Precedence;
33
34 import java.util.List;
35 import java.util.Map;
36
37 public abstract sealed class HATThreadOp extends HATOp implements Precedence.LoadOrConv
38 permits HATThreadOp.HATBlockThreadIdOp, HATThreadOp.HATGlobalSizeOp, HATThreadOp.HATGlobalThreadIdOp, HATThreadOp.HATLocalSizeOp, HATThreadOp.HATLocalThreadIdOp {
39 final private String name;
40 final private TypeElement resultType;
41 final private int dimension;
42
43 public HATThreadOp(String name, TypeElement resultType,int dimension, List<Value> operands) {
44 super(operands);
45 this.name = name;
46 this.resultType = resultType;
47 this.dimension = dimension;
48 }
49
50 protected HATThreadOp(HATThreadOp that, CodeContext cc) {
51 super(that, cc);
52 this.name =that.name;
53 this.resultType = that.resultType;
54 this.dimension = that.dimension;
55 }
56
57 public int getDimension() {
58 return dimension;
59 }
60
61
62 @Override
63 final public TypeElement resultType() {
64 return resultType;
65 }
66
67 @Override
68 final public Map<String, Object> externalize() {
69 return Map.of("hat.dialect." + name, this.getDimension());
70 }
71
72 public static final class HATLocalThreadIdOp extends HATThreadOp {
73
74 public HATLocalThreadIdOp(int dimension, TypeElement resultType) {
75 super("LocalThreadId",resultType,dimension, List.of());
76 }
77
78 public HATLocalThreadIdOp(HATLocalThreadIdOp op, CodeContext copyContext) {
79 super(op, copyContext);
80 }
81
82 @Override
83 public Op transform(CodeContext copyContext, CodeTransformer opTransformer) {
84 return new HATLocalThreadIdOp(this, copyContext);
85 }
86
87 public static HATLocalThreadIdOp of(int dimension, TypeElement resultType){
88 return new HATLocalThreadIdOp(dimension,resultType);
89 }
90 }
91
92 public static final class HATBlockThreadIdOp extends HATThreadOp {
93 public HATBlockThreadIdOp(int dimension, TypeElement resultType) {
94 super("BlockThreadId", resultType,dimension, List.of());
95 }
96
97 public HATBlockThreadIdOp(HATBlockThreadIdOp op, CodeContext copyContext) {
98 super(op, copyContext);
99 }
100
101 @Override
102 public Op transform(CodeContext copyContext, CodeTransformer opTransformer) {
103 return new HATBlockThreadIdOp(this, copyContext);
104 }
105
106
107 public static HATBlockThreadIdOp of(int dimension, TypeElement resultType){
108 return new HATBlockThreadIdOp(dimension,resultType);
109 }
110 }
111
112 public static final class HATLocalSizeOp extends HATThreadOp {
113
114 public HATLocalSizeOp(int dimension, TypeElement resultType) {
115 super("GlobalThreadSize",resultType,dimension, List.of());
116 }
117
118 public HATLocalSizeOp(HATLocalSizeOp op, CodeContext copyContext) {
119 super(op, copyContext);
120 }
121
122 @Override
123 public Op transform(CodeContext copyContext, CodeTransformer opTransformer) {
124 return new HATLocalSizeOp(this, copyContext);
125 }
126
127 public static HATThreadOp of(int dimension, TypeElement resultType){
128 return new HATLocalSizeOp(dimension, resultType);
129 }
130 }
131
132 public static final class HATGlobalThreadIdOp extends HATThreadOp {
133
134 public HATGlobalThreadIdOp(int dimension, TypeElement resultType) {
135 super("GlobalThreadId",resultType,dimension, List.of());
136 }
137
138 public HATGlobalThreadIdOp(HATGlobalThreadIdOp op, CodeContext copyContext) {
139 super(op, copyContext);
140 }
141
142 @Override
143 public Op transform(CodeContext copyContext, CodeTransformer opTransformer) {
144 return new HATGlobalThreadIdOp(this, copyContext);
145 }
146
147 public static HATGlobalThreadIdOp of(int dimension, TypeElement resultType){
148 return new HATGlobalThreadIdOp(dimension, resultType);
149 }
150 }
151
152 public static final class HATGlobalSizeOp extends HATThreadOp {
153 public HATGlobalSizeOp(int dimension, TypeElement resultType) {
154 super("GlobalThreadSize",resultType,dimension, List.of());
155 }
156
157 public HATGlobalSizeOp(HATGlobalSizeOp op, CodeContext copyContext) {
158 super(op, copyContext);
159 }
160
161 @Override
162 public Op transform(CodeContext copyContext, CodeTransformer opTransformer) {
163 return new HATGlobalSizeOp(this, copyContext);
164 }
165
166
167 static public HATGlobalSizeOp of(int dimension, TypeElement resultType){
168 return new HATGlobalSizeOp(dimension,resultType);
169 }
170 }
171 }