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 }